issue 469: TQUIC tools支持报文丢失

Chi-Kai · Chi-Kai · commit 23b7f724f109 · 2025-07-30T11:25:27.000+08:00
diff --git a/src/lib.rs b/src/lib.rs
@@ -1225,6 +1225,7 @@ pub use crate::endpoint::Endpoint;
 pub use crate::error::Error;
 pub use crate::multipath_scheduler::MultipathAlgorithm;
 pub use crate::packet::PacketHeader;
+pub use crate::packet::PacketType;
 pub use crate::tls::TlsConfig;
 pub use crate::tls::TlsConfigSelector;
 
diff --git a/src/packet.rs b/src/packet.rs
@@ -76,7 +76,7 @@ const RETRY_INTEGRITY_NONCE_V1: [u8; aead::NONCE_LEN] = [
 ];
 
 /// QUIC packet type.
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub enum PacketType {
     /// The Version Negotiation packet is a response to a client packet that
     /// contains a version that is not supported by the server.
diff --git a/tools/src/bin/tquic_client.rs b/tools/src/bin/tquic_client.rs
@@ -66,6 +66,8 @@ use tquic::PacketInfo;
 use tquic::TlsConfig;
 use tquic::TransportHandler;
 use tquic_tools::ApplicationProto;
+use tquic_tools::LossPacketType;
+use tquic_tools::PacketLossConfig;
 use tquic_tools::QuicSocket;
 use tquic_tools::Result;
 
@@ -334,6 +336,40 @@ pub struct ClientOpt {
     /// The range of the request, like "0-1023".
     #[clap(long, value_name = "RANGE", help_heading = "Protocol")]
     pub range: Option<String>,
+    /// Packet loss rate (0.0 to 1.0) for testing.
+    #[clap(
+        long,
+        default_value = "0.0",
+        value_name = "RATE",
+        help_heading = "Packet Loss"
+    )]
+    pub packet_loss_rate: f64,
+
+    /// Specific packet numbers to drop (comma-separated).
+    #[clap(
+        long,
+        value_delimiter = ',',
+        value_name = "NUMBERS",
+        help_heading = "Packet Loss"
+    )]
+    pub packet_loss_numbers: Vec<u64>,
+
+    /// Packet types to drop.
+    #[clap(
+        long,
+        value_delimiter = ',',
+        value_name = "TYPES",
+        help_heading = "Packet Loss"
+    )]
+    pub packet_loss_types: Vec<LossPacketType>,
+
+    /// Apply packet loss to incoming packets.
+    #[clap(long, help_heading = "Packet Loss")]
+    pub packet_loss_incoming: bool,
+
+    /// Apply packet loss to outgoing packets.
+    #[clap(long, help_heading = "Packet Loss")]
+    pub packet_loss_outgoing: bool,
 }
 
 const MAX_BUF_SIZE: usize = 65536;
@@ -572,7 +608,32 @@ impl Worker {
                 false => SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0),
             }
         };
-        let mut sock = QuicSocket::new(&local, registry)?;
+
+        // Create packet loss configuration if needed
+        let packet_loss_config = if option.packet_loss_rate > 0.0
+            || !option.packet_loss_numbers.is_empty()
+            || !option.packet_loss_types.is_empty()
+        {
+            let mut config = PacketLossConfig::new()
+                .with_loss_rate(option.packet_loss_rate)
+                .with_drop_packet_numbers(option.packet_loss_numbers.clone())
+                .with_drop_incoming(option.packet_loss_incoming)
+                .with_drop_outgoing(option.packet_loss_outgoing);
+
+            let packet_types: Vec<_> = option
+                .packet_loss_types
+                .iter()
+                .map(|t| t.to_packet_type())
+                .collect();
+            config = config.with_drop_packet_types(packet_types);
+
+            Some(config)
+        } else {
+            None
+        };
+
+        let mut sock =
+            QuicSocket::with_packet_loss(&local, registry, packet_loss_config, option.cid_len)?;
 
         let mut assigned_addrs = Vec::new();
         assigned_addrs.push(sock.local_addr());
diff --git a/tools/src/bin/tquic_server.rs b/tools/src/bin/tquic_server.rs
@@ -45,6 +45,8 @@ use tquic::PacketInfo;
 use tquic::TlsConfig;
 use tquic::TransportHandler;
 use tquic_tools::ApplicationProto;
+use tquic_tools::LossPacketType;
+use tquic_tools::PacketLossConfig;
 use tquic_tools::QuicSocket;
 use tquic_tools::Result;
 
@@ -250,6 +252,41 @@ pub struct ServerOpt {
     /// Disable encryption on 1-RTT packets.
     #[clap(long, help_heading = "Misc")]
     pub disable_encryption: bool,
+
+    /// Packet loss rate (0.0 to 1.0) for testing.
+    #[clap(
+        long,
+        default_value = "0.0",
+        value_name = "RATE",
+        help_heading = "Packet Loss"
+    )]
+    pub packet_loss_rate: f64,
+
+    /// Specific packet numbers to drop (comma-separated).
+    #[clap(
+        long,
+        value_delimiter = ',',
+        value_name = "NUMBERS",
+        help_heading = "Packet Loss"
+    )]
+    pub packet_loss_numbers: Vec<u64>,
+
+    /// Packet types to drop.
+    #[clap(
+        long,
+        value_delimiter = ',',
+        value_name = "TYPES",
+        help_heading = "Packet Loss"
+    )]
+    pub packet_loss_types: Vec<LossPacketType>,
+
+    /// Apply packet loss to incoming packets.
+    #[clap(long, help_heading = "Packet Loss")]
+    pub packet_loss_incoming: bool,
+
+    /// Apply packet loss to outgoing packets.
+    #[clap(long, help_heading = "Packet Loss")]
+    pub packet_loss_outgoing: bool,
 }
 
 const MAX_BUF_SIZE: usize = 65536;
@@ -315,7 +352,36 @@ impl Server {
         let registry = poll.registry();
 
         let handlers = ServerHandler::new(option)?;
-        let sock = Rc::new(QuicSocket::new(&option.listen, registry)?);
+
+        // Create packet loss configuration if needed
+        let packet_loss_config = if option.packet_loss_rate > 0.0
+            || !option.packet_loss_numbers.is_empty()
+            || !option.packet_loss_types.is_empty()
+        {
+            let mut config = PacketLossConfig::new()
+                .with_loss_rate(option.packet_loss_rate)
+                .with_drop_packet_numbers(option.packet_loss_numbers.clone())
+                .with_drop_incoming(option.packet_loss_incoming)
+                .with_drop_outgoing(option.packet_loss_outgoing);
+
+            let packet_types: Vec<_> = option
+                .packet_loss_types
+                .iter()
+                .map(|t| t.to_packet_type())
+                .collect();
+            config = config.with_drop_packet_types(packet_types);
+
+            Some(config)
+        } else {
+            None
+        };
+
+        let sock = Rc::new(QuicSocket::with_packet_loss(
+            &option.listen,
+            registry,
+            packet_loss_config,
+            option.cid_len,
+        )?);
 
         Ok(Server {
             endpoint: Endpoint::new(Box::new(config), true, Box::new(handlers), sock.clone()),
diff --git a/tools/src/common.rs b/tools/src/common.rs
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::cell::RefCell;
 use std::io::ErrorKind;
 use std::net::SocketAddr;
 
@@ -29,6 +30,9 @@ use slab::Slab;
 use tquic::PacketInfo;
 use tquic::PacketSendHandler;
 
+pub mod packet_loss;
+pub use packet_loss::{LossPacketType, PacketLossConfig, PacketLossSimulator};
+
 pub type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
 
 /// Supported application protocols.
@@ -98,10 +102,25 @@ pub struct QuicSocket {
 
     /// Local address of the initial socket.
     local_addr: SocketAddr,
+
+    /// Packet loss simulator for testing
+    packet_loss: RefCell<Option<PacketLossSimulator>>,
+
+    /// Connection ID length for packet parsing
+    dcid_len: usize,
 }
 
 impl QuicSocket {
     pub fn new(local: &SocketAddr, registry: &Registry) -> Result<Self> {
+        Self::with_packet_loss(local, registry, None, 8)
+    }
+
+    pub fn with_packet_loss(
+        local: &SocketAddr,
+        registry: &Registry,
+        packet_loss_config: Option<PacketLossConfig>,
+        dcid_len: usize,
+    ) -> Result<Self> {
         let mut socks = Slab::new();
         let mut addrs = FxHashMap::default();
 
@@ -113,10 +132,14 @@ impl QuicSocket {
         let socket = socks.get_mut(sid).unwrap();
         registry.register(socket, Token(sid), Interest::READABLE)?;
 
+        let packet_loss = RefCell::new(packet_loss_config.map(PacketLossSimulator::new));
+
         Ok(Self {
             socks,
             addrs,
             local_addr,
+            packet_loss,
+            dcid_len,
         })
     }
 
@@ -166,14 +189,38 @@ impl QuicSocket {
         };
 
         match socket.recv_from(buf) {
-            Ok((len, remote)) => Ok((len, socket.local_addr()?, remote)),
+            Ok((len, remote)) => {
+                // Check if packet should be dropped due to loss simulation
+                if let Ok(mut packet_loss) = self.packet_loss.try_borrow_mut() {
+                    if let Some(ref mut simulator) = *packet_loss {
+                        if simulator.should_drop_incoming(&buf[..len], self.dcid_len) {
+                            debug!("Simulating incoming packet loss - dropping packet");
+                            return Err(std::io::Error::new(
+                                ErrorKind::WouldBlock,
+                                "simulated packet loss",
+                            ));
+                        }
+                    }
+                }
+                Ok((len, socket.local_addr()?, remote))
+            }
             Err(e) => Err(e),
         }
     }
 
     /// Send data on the socket to the given address.
     /// Note: packets with unknown src address are dropped.
     pub fn send_to(&self, buf: &[u8], src: SocketAddr, dst: SocketAddr) -> std::io::Result<usize> {
+        // Check if packet should be dropped due to loss simulation
+        if let Ok(mut packet_loss) = self.packet_loss.try_borrow_mut() {
+            if let Some(ref mut simulator) = *packet_loss {
+                if simulator.should_drop_outgoing(buf, self.dcid_len) {
+                    debug!("Simulating outgoing packet loss - dropping packet");
+                    return Ok(buf.len()); // Pretend we sent it
+                }
+            }
+        }
+
         let sid = match self.addrs.get(&src) {
             Some(sid) => sid,
             None => {
diff --git a/tools/src/packet_loss.rs b/tools/src/packet_loss.rs
