feat: add error mapping for ipv4/ipv6 tcp bind()

fujiapple852 · Mar 13, 2023 · 8707c96 · 8707c96
1 parent 203a3d7
commit 8707c96
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 36 deletions.
diff --git a/src/tracing/net/ipv4.rs b/src/tracing/net/ipv4.rs
@@ -150,34 +150,39 @@ pub fn dispatch_tcp_probe(
     port_direction: PortDirection,
     tos: TypeOfService,
 ) -> TraceResult<Socket> {
+    fn process_result(addr: SocketAddr, res: std::io::Result<()>) -> TraceResult<()> {
+        match res {
+            Ok(_) => Ok(()),
+            Err(err) => {
+                if let Some(code) = err.raw_os_error() {
+                    if platform::is_not_in_progress_error(code) {
+                        match err.kind() {
+                            ErrorKind::AddrInUse | ErrorKind::AddrNotAvailable => {
+                                Err(AddressNotAvailable(addr))
+                            }
+                            _ => Err(TracerError::IoError(err)),
+                        }
+                    } else {
+                        Ok(())
+                    }
+                } else {
+                    Err(TracerError::IoError(err))
+                }
+            }
+        }
+    }
     let (src_port, dest_port) = match port_direction {
         PortDirection::FixedSrc(src_port) => (src_port.0, probe.sequence.0),
         PortDirection::FixedDest(dest_port) => (probe.sequence.0, dest_port.0),
         PortDirection::FixedBoth(_, _) | PortDirection::None => unimplemented!(),
     };
     let mut socket = Socket::new_stream_socket_ipv4()?;
     let local_addr = SocketAddr::new(IpAddr::V4(src_addr), src_port);
-    socket.bind(local_addr)?;
+    process_result(local_addr, socket.bind(local_addr))?;
     socket.set_ttl(u32::from(probe.ttl.0))?;
     socket.set_tos(u32::from(tos.0))?;
     let remote_addr = SocketAddr::new(IpAddr::V4(dest_addr), dest_port);
-    match socket.connect(remote_addr) {
-        Ok(_) => {}
-        Err(err) => {
-            if let Some(code) = err.raw_os_error() {
-                if platform::is_not_in_progress_error(code) {
-                    return match err.kind() {
-                        ErrorKind::AddrInUse | ErrorKind::AddrNotAvailable => {
-                            Err(AddressNotAvailable(local_addr))
-                        }
-                        _ => Err(TracerError::IoError(err)),
-                    };
-                }
-            } else {
-                return Err(TracerError::IoError(err));
-            }
-        }
-    }
+    process_result(remote_addr, socket.connect(remote_addr))?;
     Ok(socket)
 }
 

diff --git a/src/tracing/net/ipv6.rs b/src/tracing/net/ipv6.rs
@@ -105,33 +105,38 @@ pub fn dispatch_tcp_probe(
     dest_addr: Ipv6Addr,
     port_direction: PortDirection,
 ) -> TraceResult<Socket> {
+    fn process_result(addr: SocketAddr, res: std::io::Result<()>) -> TraceResult<()> {
+        match res {
+            Ok(_) => Ok(()),
+            Err(err) => {
+                if let Some(code) = err.raw_os_error() {
+                    if platform::is_not_in_progress_error(code) {
+                        match err.kind() {
+                            ErrorKind::AddrInUse | ErrorKind::AddrNotAvailable => {
+                                Err(AddressNotAvailable(addr))
+                            }
+                            _ => Err(TracerError::IoError(err)),
+                        }
+                    } else {
+                        Ok(())
+                    }
+                } else {
+                    Err(TracerError::IoError(err))
+                }
+            }
+        }
+    }
     let (src_port, dest_port) = match port_direction {
         PortDirection::FixedSrc(src_port) => (src_port.0, probe.sequence.0),
         PortDirection::FixedDest(dest_port) => (probe.sequence.0, dest_port.0),
         PortDirection::FixedBoth(_, _) | PortDirection::None => unimplemented!(),
     };
     let mut socket = Socket::new_stream_socket_ipv6()?;
     let local_addr = SocketAddr::new(IpAddr::V6(src_addr), src_port);
-    socket.bind(local_addr)?;
+    process_result(local_addr, socket.bind(local_addr))?;
     socket.set_unicast_hops_v6(probe.ttl.0)?;
     let remote_addr = SocketAddr::new(IpAddr::V6(dest_addr), dest_port);
-    match socket.connect(remote_addr) {
-        Ok(_) => {}
-        Err(err) => {
-            if let Some(code) = err.raw_os_error() {
-                if platform::is_not_in_progress_error(code) {
-                    return match err.kind() {
-                        ErrorKind::AddrInUse | ErrorKind::AddrNotAvailable => {
-                            Err(AddressNotAvailable(local_addr))
-                        }
-                        _ => Err(TracerError::IoError(err)),
-                    };
-                }
-            } else {
-                return Err(TracerError::IoError(err));
-            }
-        }
-    }
+    process_result(remote_addr, socket.bind(remote_addr))?;
     Ok(socket)
 }