feat: add TcpStream bind_and_connect (#409)

Author: Sherlock-Holo

.gitignore

@@ -2,6 +2,6 @@
 /Cargo.lock
 /.vscode
 /.cargo
-/.idea
+.idea
 /.direnv
-.envrc
+.envrc

compio-net/src/tcp.rs

@@ -176,6 +176,22 @@ impl TcpStream {
         .await
     }
 
+    /// Bind to `bind_addr` then opens a TCP connection to a remote host.
+    pub async fn bind_and_connect(
+        bind_addr: SocketAddr,
+        addr: impl ToSocketAddrsAsync,
+    ) -> io::Result<Self> {
+        super::each_addr(addr, |addr| async move {
+            let addr = SockAddr::from(addr);
+            let bind_addr = SockAddr::from(bind_addr);
+
+            let socket = Socket::bind(&bind_addr, Type::STREAM, Some(Protocol::TCP)).await?;
+            socket.connect_async(&addr).await?;
+            Ok(Self { inner: socket })
+        })
+        .await
+    }
+
     /// Creates new TcpStream from a [`std::net::TcpStream`].
     pub fn from_std(stream: std::net::TcpStream) -> io::Result<Self> {
         Ok(Self {

compio-net/tests/tcp_connect.rs

@@ -1,5 +1,5 @@
 use std::{
-    net::{IpAddr, SocketAddr},
+    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
     panic::resume_unwind,
 };
 
@@ -42,6 +42,44 @@ test_connect_ip! {
     (connect_v6, "[::1]:0", SocketAddr::is_ipv6),
 }
 
+async fn test_bind_and_connect_ip_impl(
+    bind_addr: SocketAddr,
+    target: impl ToSocketAddrsAsync,
+    assert_fn: impl FnOnce(&SocketAddr) -> bool,
+) {
+    let listener = TcpListener::bind(target).await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    assert!(assert_fn(&addr));
+
+    let task = compio_runtime::spawn(async move {
+        let (socket, addr) = listener.accept().await.unwrap();
+        assert_eq!(addr, socket.peer_addr().unwrap());
+        socket
+    });
+
+    let mine = TcpStream::bind_and_connect(bind_addr, &addr).await.unwrap();
+    let theirs = task.await.unwrap_or_else(|e| resume_unwind(e));
+
+    assert_eq!(mine.local_addr().unwrap(), theirs.peer_addr().unwrap());
+    assert_eq!(theirs.local_addr().unwrap(), mine.peer_addr().unwrap());
+}
+
+macro_rules! test_bind_and_connect_ip {
+    ($(($ident:ident, $bind_addr:expr, $target:expr, $addr_f:path),)*) => {
+        $(
+            #[compio_macros::test]
+            async fn $ident() {
+                test_bind_and_connect_ip_impl($bind_addr, $target, $addr_f).await;
+            }
+        )*
+    }
+}
+
+test_bind_and_connect_ip! {
+    (bind_and_connect_v4, SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 0), "127.0.0.1:0", SocketAddr::is_ipv4),
+    (bind_and_connect_v6, SocketAddr::new(Ipv6Addr::LOCALHOST.into(), 0), "[::1]:0", SocketAddr::is_ipv6),
+}
+
 async fn test_connect_impl<A: ToSocketAddrsAsync>(mapping: impl FnOnce(&TcpListener) -> A) {
     let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
     let addr = mapping(&listener);