How do I build correct TCP tunnel (forwarder) with Tokio 0.3?

[vi]

Naive attempt:

#[tokio::main(flavor="current_thread")]
async fn main() {
    use tokio::net::{TcpListener, TcpStream};
    let ss = TcpListener::bind("127.0.0.1:5747".parse::<std::net::SocketAddr>().unwrap()).await.unwrap();
    while let Ok((cs,_)) = ss.accept().await {
        let g = TcpStream::connect("127.0.0.1:5746".parse::<std::net::SocketAddr>().unwrap()).await.unwrap();
        let (mut gr, mut gw) = g.into_split();
        let (mut csr, mut csw) = cs.into_split();
        use tokio::io::AsyncWriteExt;
        tokio::spawn(async move {
            let _ = tokio::io::copy(&mut gr, &mut csw).await;
            let _ = csw.shutdown();
        });
        tokio::spawn( async move {
            let _ = tokio::io::copy(&mut csr, &mut gw).await;
            let _ = gw.shutdown();
        });
    }
}

works reasonably, but fails to propagate socket close when the tunnel is clogged (unavailable for writing) in both directions.

I made a special checker for TCP tunnels: https://github.com/vi/tcptunnelchecker. It shows that socat tcp-l:5747,fork,reuseaddr tcp:127.0.0.1:5746 is a properly implemented tunnel, but the one above is not:

[ OK ] Trivial test 1
[ OK ] Trivial test 2
The other direction buffer: 8858075
One direction buffer: 11072137
[FAIL] Clogged close test 1 failed!
[FAIL] Clogged close test 2 failed!
[ OK ] Clogged close test 3 passed
[ OK ] Clogged close test 4 passed
[ OK ] Clogged close test 5 passed
[ OK ] Clogged close test 6 passed

[Darksonn]

More generally, how does socat do this? It is unclear to me.

[vi]

1. socat is waiting for both TCP sockets to become writable:

13587 23:05:05.235138 select(7, [], [5<TCP:[127.0.0.1:56300->127.0.0.1:5746]> 6<TCP:[127.0.0.1:5747->127.0.0.1:50314]>], [], NULL) =

When tunnel checker closes the socket, select(2) wakes up and signals writeability: = 1 (out [6])

2. socat reads a bit of incoming data from the other socket:

13587 23:05:05.391310 read(5<TCP:[127.0.0.1:56300->127.0.0.1:5746]>, "\0\0\0\0\0\0"..., 8192) = 8192

3. , then tries to write it to the outgoing socket:

13587 23:05:05.391759 write(6<TCP:[127.0.0.1:5747->127.0.0.1:50314]>, "\0\0\0"..., 8192) = -1 ECONNRESET (Connection reset by peer)

4. Upon getting the writing error it just bails out the child process:

13587 23:05:05.392724 write(2</dev/pts/46<char 136:46>>, "2020/11/28 23:05:05 socat[13587]"..., 92) = 92
13587 23:05:05.392834 shutdown(6<TCP:[127.0.0.1:5747->127.0.0.1:50314]>, SHUT_RDWR) = -1 ENOTCONN (Transport endpoint is not connected)
13587 23:05:05.392913 shutdown(5<TCP:[127.0.0.1:56300->127.0.0.1:5746]>, SHUT_RDWR) = 0
13587 23:05:05.393372 exit_group(1)     = ?

5. Even without shutdown(2)s OS would close those file descriptors anyway, causing ENOTCONN to propagate to the checker.

[Darksonn]

Ah, ok, so the difference is that socat doesn't read from one socket until the other socket is write ready. You can wait for a socket to become ready for writing with TcpStream::writable. You can't use tokio::io::copy with this method, as it is generic over IO resource, and the generic interface doesn't expose writable.

[vi]

It is also important to keep tabs on socket where to intent to write while being in the reading phase.
You are interested in the event that it is no longer possible to write to the socket, hence it's not worth reading from the other socket anymore either.

[vi]

You can wait for a socket to become ready for writing with TcpStream::writable

I need to wait for a writable TcpStream to stop being writable due to disconnection.

[vi]

Here is TCP forwarder that passes the tests like socat (unrefactored):

use std::net::{TcpListener,TcpStream,SocketAddr,Shutdown};

type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync + 'static>>;

fn forward(mut s1: TcpStream, mut s2: TcpStream) -> Result<()> {
    use std::os::unix::io::AsRawFd;
    // nix = "0.19.0"
    use nix::fcntl::{fcntl,OFlag,FcntlArg};
    use nix::poll::{poll, PollFd, PollFlags};
    use std::io::{Read,Write,ErrorKind};
    type UnwrittedDataInBuffer = std::ops::Range<usize>;

    fcntl ( s1.as_raw_fd(), FcntlArg::F_SETFL(OFlag::O_NONBLOCK) )?;
    fcntl ( s2.as_raw_fd(), FcntlArg::F_SETFL(OFlag::O_NONBLOCK) )?;
    let mut buf_1to2 = [0u8; 1024];
    let mut buf_2to1 = [0u8; 1024];
    let mut datarange_1to2 : Option<UnwrittedDataInBuffer> = None;
    let mut datarange_2to1 : Option<UnwrittedDataInBuffer> = None;
    let mut finished_1to2 = false;
    let mut finished_2to1 = false;

    let mut polls : [PollFd; 2] = [PollFd::new(-1, PollFlags::empty()); 2]; 
    let mut pollflags : [PollFlags; 2] = [PollFlags::empty(); 2];
    //let mut active_polls : &mut [PollFd] ;

    loop {
        if finished_1to2 && finished_2to1 {
            drop(s1);
            drop(s2);
            return Ok(());
        }

        pollflags[0] = PollFlags::POLLHUP;
        pollflags[1] = PollFlags::POLLHUP;
        if datarange_1to2.is_none() {
            pollflags[0] |= PollFlags::POLLIN;
        } else {
            pollflags[1] |= PollFlags::POLLOUT;
        }

        if datarange_2to1.is_none() {
            pollflags[1] |= PollFlags::POLLIN;
        } else {
            pollflags[0] |= PollFlags::POLLOUT;
        }

        polls[0] = PollFd::new(s1.as_raw_fd(), pollflags[0]);
        polls[1] = PollFd::new(s2.as_raw_fd(), pollflags[1]);

        poll(&mut polls[..], -1)?;

        macro_rules! can_definitely_be_done_using_usual_function {
            ($pollidx_r:literal, $pollidx_w:literal, $datarange:ident, $buf:ident, $finished:ident, $socket_r:ident, $socket_w:ident,) => {
                if let Some(r) = $datarange.take() {
                    if polls[$pollidx_w].revents().map(|x|x.contains(PollFlags::POLLOUT)) == Some(true) {
                        match $socket_w.write(& $buf[r.clone()]) {
                            Ok(x) => {
                                if x == r.len() {
                                    $datarange = None;
                                } else {
                                    $datarange = Some((r.start + x)..(r.end));
                                }
                            }
                            Err(e) if e.kind() == ErrorKind::WouldBlock => {
                            }
                            Err(e) => Err(e)?,
                        }
                    } else {
                        $datarange = Some(r);
                    }
                } else {
                    if polls[$pollidx_w].revents().map(|x|x.contains(PollFlags::POLLHUP)) == Some(true) {
                        // Abort reading from socket when we know that
                        // writing to the other socket is destined to fail
                        $finished = true;
                        let _ = $socket_w.shutdown(Shutdown::Write);
                        // But for some strange reason absence of this section does not affect test results.
                    } else
                    if polls[$pollidx_r].revents().map(|x|x.contains(PollFlags::POLLIN)) == Some(true) {
                        match $socket_r.read(&mut $buf[..]) {
                            Ok(0) => {
                                $finished = true;
                                $socket_w.shutdown(Shutdown::Write)?;
                            }
                            Ok(x) => {
                                $datarange = Some(0..x);
                            }
                            Err(e) if e.kind() == ErrorKind::WouldBlock => {
                            }
                            Err(e) => Err(e)?,
                        };
                    }
                }
            }
        }

        can_definitely_be_done_using_usual_function!(
            0, 
            1,
            datarange_1to2,
            buf_1to2,
            finished_1to2,
            s1,
            s2,
        );

        can_definitely_be_done_using_usual_function!(
            1,
            0,
            datarange_2to1,
            buf_2to1,
            finished_2to1,
            s2,
            s1,
        );
    }
}

fn main() -> Result<()> {
    let ss = TcpListener::bind("127.0.0.1:5747".parse::<SocketAddr>().unwrap())?;
    while let Ok((cs,_)) = ss.accept() {
        std::thread::spawn(move || {
            let g = TcpStream::connect("127.0.0.1:5746".parse::<SocketAddr>().unwrap()).unwrap();
            if let Err(e) = forward(cs, g) {
                eprintln!("{}", e);
            }
        });
    }
    Ok(())
}

Also tokio 0.1-based websocat 1.6, when invoked as websocat -b tcp-l:127.0.0.1:5747 tcp:127.0.0.1:5746 also passed all the tests.

The naive forwarder above fails only the clogged-in-both-directions test.

[vi]

@Darksonn Why is it marked as the answer to the discussion? This sample is not Tokio-based. It is just to demonstrate that it is possible to implement a compliant forwarder on UNIX-like in Rust (without using Tokio).

this comment is closer to the real answer (where actual real answer should be a tokio-0.3-based, but non-io::copy-based forwarder that takes writable into account).

[Darksonn]

Maybe I was just too quick 😅

Feel free to mark whatever you wish as the answer.