tunneler 0.4.0

Tunnel TCP or UDP traffic over TCP, (mutual) TLS or DNS (authoritative server or direct connection)
Documentation
use std::error::Error;
use std::net::IpAddr;

use async_trait::async_trait;
use tokio::net::TcpStream;

use common::io::{copy, Stream};

#[async_trait]
pub(crate) trait Tunneler: Send {
    async fn tunnel(&mut self, mut client: Stream) -> Result<(), Box<dyn Error>>;
}

pub(crate) struct TcpTunneler {
    tunneled: Stream,
}

impl TcpTunneler {
    pub(crate) async fn new(address: IpAddr, port: u16) -> Result<Self, Box<dyn Error>> {
        let to_address = format!("{}:{}", address, port);
        log::debug!("connecting to {}", to_address);
        let stream = TcpStream::connect(to_address).await?;
        let (reader, writer) = stream.into_split();
        TcpTunneler::from_stream(Stream::new(reader, writer))
    }

    pub(crate) fn from_stream(stream: Stream) -> Result<Self, Box<dyn Error>> {
        Ok(Self { tunneled: stream })
    }
}

#[async_trait]
impl Tunneler for TcpTunneler {
    async fn tunnel(&mut self, mut client: Stream) -> Result<(), Box<dyn Error>> {
        let to_tunnel_future = copy(
            &mut client.reader,
            &mut self.tunneled.writer,
            "sending to tunnel",
        );
        let from_tunnel_future = copy(
            &mut self.tunneled.reader,
            &mut client.writer,
            "received from tunnel",
        );
        match tokio::try_join!(to_tunnel_future, from_tunnel_future) {
            Ok(_) => Ok(()),
            Err(e) => Err(e.into()),
        }
    }
}

#[cfg(test)]
mod tests {
    use tokio_test::io::Builder;

    use super::*;

    #[tokio::test]
    async fn tcp_single_read() -> Result<(), Box<dyn Error>> {
        let tunneled_client = Stream::new(
            Builder::new().build(),
            Builder::new().write(b"hello").build(),
        );
        let mut tunneler = TcpTunneler {
            tunneled: tunneled_client,
        };

        let client = Stream::new(
            Builder::new().read(b"hello").build(),
            Builder::new().build(),
        );
        tunneler.tunnel(client).await
    }

    #[tokio::test]
    async fn tcp_single_write() -> Result<(), Box<dyn Error>> {
        let tunneled_client = Stream::new(
            Builder::new().read(b"hello").build(),
            Builder::new().build(),
        );
        let mut tunneler = TcpTunneler {
            tunneled: tunneled_client,
        };

        let client = Stream::new(
            Builder::new().build(),
            Builder::new().write(b"hello").build(),
        );
        tunneler.tunnel(client).await
    }

    #[tokio::test]
    async fn tcp_single_read_write() -> Result<(), Box<dyn Error>> {
        let tunneled_client = Stream::new(
            Builder::new().read(b"hello").build(),
            Builder::new().write(b"world").build(),
        );
        let mut tunneler = TcpTunneler {
            tunneled: tunneled_client,
        };

        let client = Stream::new(
            Builder::new().read(b"world").build(),
            Builder::new().write(b"hello").build(),
        );
        tunneler.tunnel(client).await
    }

    #[tokio::test]
    async fn tcp_multiple_reads() -> Result<(), Box<dyn Error>> {
        let tunneled_client = Stream::new(
            Builder::new().build(),
            Builder::new().write(b"1").write(b"2").write(b"3").build(),
        );
        let mut tunneler = TcpTunneler {
            tunneled: tunneled_client,
        };

        let client = Stream::new(
            Builder::new().read(b"1").read(b"2").read(b"3").build(),
            Builder::new().build(),
        );
        tunneler.tunnel(client).await
    }

    #[tokio::test]
    async fn tcp_multiple_writes() -> Result<(), Box<dyn Error>> {
        let tunneled_client = Stream::new(
            Builder::new().build(),
            Builder::new().write(b"1").write(b"2").write(b"3").build(),
        );
        let mut tunneler = TcpTunneler {
            tunneled: tunneled_client,
        };

        let client = Stream::new(
            Builder::new().read(b"1").read(b"2").read(b"3").build(),
            Builder::new().build(),
        );
        tunneler.tunnel(client).await
    }
}