fix(net): only create the needed icmp or udp sending socket (#647)

fujiapple852 · Aug 19, 2023 · b842e72 · b842e72
1 parent b00f0ed
commit b842e72
Showing 1 changed file with 17 additions and 16 deletions.
diff --git a/src/tracing/net/channel.rs b/src/tracing/net/channel.rs
@@ -29,8 +29,7 @@ pub struct TracerChannel {
     tos: TypeOfService,
     read_timeout: Duration,
     tcp_connect_timeout: Duration,
-    icmp_send_socket: Socket,
-    udp_send_socket: Socket,
+    send_socket: Option<Socket>,
     recv_socket: Socket,
     tcp_probes: ArrayVec<TcpProbe, MAX_TCP_PROBES>,
 }
@@ -50,8 +49,11 @@ impl TracerChannel {
         platform::startup()?;
         let ipv4_length_order =
             platform::PlatformIpv4FieldByteOrder::for_address(config.source_addr)?;
-        let icmp_send_socket = make_icmp_send_socket(config.source_addr)?;
-        let udp_send_socket = make_udp_send_socket(config.source_addr)?;
+        let send_socket = match config.protocol {
+            TracerProtocol::Icmp => Some(make_icmp_send_socket(config.source_addr)?),
+            TracerProtocol::Udp => Some(make_udp_send_socket(config.source_addr)?),
+            TracerProtocol::Tcp => None,
+        };
         let recv_socket = make_recv_socket(config.source_addr)?;
         Ok(Self {
             protocol: config.protocol,
@@ -64,8 +66,7 @@ impl TracerChannel {
             tos: config.tos,
             read_timeout: config.read_timeout,
             tcp_connect_timeout: config.tcp_connect_timeout,
-            icmp_send_socket,
-            udp_send_socket,
+            send_socket,
             recv_socket,
             tcp_probes: ArrayVec::new(),
         })
@@ -101,18 +102,18 @@ impl TracerChannel {
     /// Dispatch a ICMP probe.
     #[instrument(skip_all)]
     fn dispatch_icmp_probe(&mut self, probe: Probe) -> TraceResult<()> {
-        match (self.src_addr, self.dest_addr) {
-            (IpAddr::V4(src_addr), IpAddr::V4(dest_addr)) => ipv4::dispatch_icmp_probe(
-                &mut self.icmp_send_socket,
+        match (self.src_addr, self.dest_addr, self.send_socket.as_mut()) {
+            (IpAddr::V4(src_addr), IpAddr::V4(dest_addr), Some(s)) => ipv4::dispatch_icmp_probe(
+                s,
                 probe,
                 src_addr,
                 dest_addr,
                 self.packet_size,
                 self.payload_pattern,
                 self.ipv4_length_order,
             ),
-            (IpAddr::V6(src_addr), IpAddr::V6(dest_addr)) => ipv6::dispatch_icmp_probe(
-                &mut self.icmp_send_socket,
+            (IpAddr::V6(src_addr), IpAddr::V6(dest_addr), Some(s)) => ipv6::dispatch_icmp_probe(
+                s,
                 probe,
                 src_addr,
                 dest_addr,
@@ -126,9 +127,9 @@ impl TracerChannel {
     /// Dispatch a UDP probe.
     #[instrument(skip_all)]
     fn dispatch_udp_probe(&mut self, probe: Probe) -> TraceResult<()> {
-        match (self.src_addr, self.dest_addr) {
-            (IpAddr::V4(src_addr), IpAddr::V4(dest_addr)) => ipv4::dispatch_udp_probe(
-                &mut self.udp_send_socket,
+        match (self.src_addr, self.dest_addr, self.send_socket.as_mut()) {
+            (IpAddr::V4(src_addr), IpAddr::V4(dest_addr), Some(s)) => ipv4::dispatch_udp_probe(
+                s,
                 probe,
                 src_addr,
                 dest_addr,
@@ -137,8 +138,8 @@ impl TracerChannel {
                 self.multipath_strategy,
                 self.ipv4_length_order,
             ),
-            (IpAddr::V6(src_addr), IpAddr::V6(dest_addr)) => ipv6::dispatch_udp_probe(
-                &mut self.udp_send_socket,
+            (IpAddr::V6(src_addr), IpAddr::V6(dest_addr), Some(s)) => ipv6::dispatch_udp_probe(
+                s,
                 probe,
                 src_addr,
                 dest_addr,