Function test_tcp_forward_roundtrip

port-forward/src/tcp.rs:164–206

()

    #[tokio::test]
    async fn test_tcp_forward_roundtrip() {
        // Echo server
        let echo = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let echo_addr = echo.local_addr().unwrap();
        tokio::spawn(async move {
            loop {
                let (mut conn, _) = echo.accept().await.unwrap();
                tokio::spawn(async move {
                    let mut buf = vec![0u8; 4096];
                    loop {
                        let n = match conn.read(&mut buf).await {
                            Ok(0) | Err(_) => break,
                            Ok(n) => n,
                        };
                        if conn.write_all(&buf[..n]).await.is_err() {
                            break;
                        }
                    }
                });
            }
        });

        // Start forwarder on a free port
        let tmp = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let fwd_addr = tmp.local_addr().unwrap();
        drop(tmp);

        let cancel = CancellationToken::new();
        tokio::spawn(run_tcp_forwarder(fwd_addr, echo_addr, cancel.child_token()));
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        let mut conn = TcpStream::connect(fwd_addr).await.unwrap();
        conn.write_all(b"hello splice").await.unwrap();

        let mut buf = vec![0u8; 64];
        let n = tokio::time::timeout(std::time::Duration::from_secs(2), conn.read(&mut buf))
            .await
            .unwrap()
            .unwrap();
        assert_eq!(&buf[..n], b"hello splice");

        cancel.cancel();
    }
}