tcp: Add source address option

This option allows to specify the source address and port for outgoing TCP connections. Signed-off-by: Konrad Gräfe <[email protected]>
lemunozm · May 3, 2023 · 6cd8a6d · 6cd8a6d
1 parent 73b71a1
commit 6cd8a6d
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 3 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -41,7 +41,7 @@ version = "0.26.2"
 default-features = false
 features = ["socket", "uio", "net"]
 
-[target.'cfg(target_os = "linux")'.dependencies.libc]
+[target.'cfg(unix)'.dependencies.libc]
 version = "0.2.137"
 default-features = false
 

diff --git a/src/adapters/tcp.rs b/src/adapters/tcp.rs
@@ -9,7 +9,7 @@ use crate::network::{RemoteAddr, Readiness, TransportConnect, TransportListen};
 use mio::net::{TcpListener, TcpStream};
 use mio::event::{Source};
 
-use socket2::{Socket};
+use socket2::{Socket, Domain, Type, Protocol};
 
 use std::net::{SocketAddr};
 use std::io::{self, ErrorKind, Read, Write};
@@ -27,6 +27,7 @@ pub const INPUT_BUFFER_SIZE: usize = u16::MAX as usize; // 2^16 - 1
 
 #[derive(Clone, Debug, Default)]
 pub struct TcpConnectConfig {
+    source_address: Option<SocketAddr>,
     keepalive: Option<TcpKeepalive>,
 }
 
@@ -36,6 +37,12 @@ impl TcpConnectConfig {
         self.keepalive = Some(keepalive);
         self
     }
+
+    /// Specify the source address and port.
+    pub fn with_source_address(mut self, source_address: SocketAddr) -> Self {
+        self.source_address = Some(source_address);
+        self
+    }
 }
 
 #[derive(Clone, Debug, Default)]
@@ -78,7 +85,30 @@ impl Remote for RemoteResource {
             _ => panic!("Internal error: Got wrong config"),
         };
         let peer_addr = *remote_addr.socket_addr();
-        let stream = TcpStream::connect(peer_addr)?;
+
+        let socket = Socket::new(
+            match peer_addr {
+                SocketAddr::V4 { .. } => Domain::IPV4,
+                SocketAddr::V6 { .. } => Domain::IPV6,
+            },
+            Type::STREAM,
+            Some(Protocol::TCP),
+        )?;
+        socket.set_nonblocking(true)?;
+
+        if let Some(source_address) = config.source_address {
+            socket.bind(&source_address.into())?;
+        }
+
+        match socket.connect(&peer_addr.into()) {
+            #[cfg(unix)]
+            Err(e) if e.raw_os_error() != Some(libc::EINPROGRESS) => return Err(e),
+            #[cfg(windows)]
+            Err(e) if e.kind() != io::ErrorKind::WouldBlock => return Err(e),
+            _ => {}
+        }
+
+        let stream = TcpStream::from_std(socket.into());
         let local_addr = stream.local_addr()?;
         Ok(ConnectionInfo {
             remote: Self { stream, keepalive: config.keepalive },