tokio: update connect to async/await (#1375)

2026-09-08 00:00:13 +02:00 · 2019-08-02 10:03:06 -07:00
parent 81d789b88f
commit 144d980e5c
2 changed files with 208 additions and 252 deletions
@@ -0,0 +1,208 @@
 //! An example of hooking up stdin/stdout to either a TCP or UDP stream.
 //!
 //! This example will connect to a socket address specified in the argument list
 //! and then forward all data read on stdin to the server, printing out all data
 //! received on stdout. An optional `--udp` argument can be passed to specify
 //! that the connection should be made over UDP instead of TCP, translating each
 //! line entered on stdin to a UDP packet to be sent to the remote address.
 //!
 //! Note that this is not currently optimized for performance, especially
 //! around buffer management. Rather it's intended to show an example of
 //! working with a client.
 //!
 //! This example can be quite useful when interacting with the other examples in
 //! this repository! Many of them recommend running this as a simple "hook up
 //! stdin/stdout to a server" to get up and running.
 #![deny(warnings, rust_2018_idioms)]
 #![feature(async_await)]
 use futures::{SinkExt, Stream};
 use std::{env, error::Error, net::SocketAddr};
 use tokio::{
    codec::{FramedRead, FramedWrite},
    io,
    sync::{mpsc, oneshot},
 };
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn Error>> {
    let (tx, rx) = oneshot::channel();
    tokio::spawn(async move {
        run().await.unwrap();
        tx.send(()).unwrap();
    });
    rx.await.map_err(Into::into)
 }
 // Currently, we need to spawn the initial future due to https://github.com/tokio-rs/tokio/issues/1356
 async fn run() -> Result<(), Box<dyn Error>> {
    // Determine if we're going to run in TCP or UDP mode
    let mut args = env::args().skip(1).collect::<Vec<_>>();
    let tcp = match args.iter().position(|a| a == "--udp") {
        Some(i) => {
            args.remove(i);
            false
        }
        None => true,
    };
    // Parse what address we're going to connect to
    let addr = match args.first() {
        Some(addr) => addr,
        None => Err("this program requires at least one argument")?,
    };
    let addr = addr.parse::<SocketAddr>()?;
    let stdin = stdin();
    let stdout = FramedWrite::new(io::stdout(), codec::Bytes);
    if tcp {
        tcp::connect(&addr, stdin, stdout).await?;
    } else {
        udp::connect(&addr, stdin, stdout).await?;
    }
    Ok(())
 }
 // Temporary work around for stdin blocking the stream
 fn stdin() -> impl Stream<Item = Result<Vec<u8>, io::Error>> + Unpin {
    let mut stdin = FramedRead::new(io::stdin(), codec::Bytes);
    let (mut tx, rx) = mpsc::unbounded_channel();
    tokio::spawn(async move {
        tx.send_all(&mut stdin).await.unwrap();
    });
    rx
 }
 mod tcp {
    use super::codec;
    use futures::{future, Sink, SinkExt, Stream, StreamExt};
    use std::{error::Error, io, net::SocketAddr};
    use tokio::{
        codec::{FramedRead, FramedWrite},
        net::TcpStream,
    };
    pub async fn connect(
        addr: &SocketAddr,
        stdin: impl Stream<Item = Result<Vec<u8>, io::Error>> + Unpin,
        mut stdout: impl Sink<Vec<u8>, Error = io::Error> + Unpin,
    ) -> Result<(), Box<dyn Error>> {
        let (r, w) = TcpStream::connect(addr).await?.split();
        let sink = FramedWrite::new(w, codec::Bytes);
        let mut stream = FramedRead::new(r, codec::Bytes).filter_map(|i| match i {
            Ok(i) => future::ready(Some(i)),
            Err(e) => {
                println!("failed to read from socket; error={}", e);
                future::ready(None)
            }
        });
        match future::join(stdin.forward(sink), stdout.send_all(&mut stream)).await {
            (Err(e), _) | (_, Err(e)) => Err(e.into()),
            _ => Ok(()),
        }
    }
 }
 mod udp {
    use futures::{future, Sink, SinkExt, Stream, StreamExt};
    use std::{error::Error, io, net::SocketAddr};
    use tokio::net::udp::{
        split::{UdpSocketRecvHalf, UdpSocketSendHalf},
        UdpSocket,
    };
    pub async fn connect(
        addr: &SocketAddr,
        stdin: impl Stream<Item = Result<Vec<u8>, io::Error>> + Unpin,
        stdout: impl Sink<Vec<u8>, Error = io::Error> + Unpin,
    ) -> Result<(), Box<dyn Error>> {
        // We'll bind our UDP socket to a local IP/port, but for now we
        // basically let the OS pick both of those.
        let bind_addr = if addr.ip().is_ipv4() {
            "0.0.0.0:0".parse()?
        } else {
            "[::]:0".parse()?
        };
        let socket = UdpSocket::bind(&bind_addr)?;
        socket.connect(addr)?;
        let (mut r, mut w) = socket.split();
        future::try_join(send(stdin, &mut w), recv(stdout, &mut r)).await?;
        Ok(())
    }
    async fn send(
        mut stdin: impl Stream<Item = Result<Vec<u8>, io::Error>> + Unpin,
        writer: &mut UdpSocketSendHalf,
    ) -> Result<(), io::Error> {
        while let Some(item) = stdin.next().await {
            let buf = item?;
            writer.send(&buf[..]).await?;
        }
        Ok(())
    }
    async fn recv(
        mut stdout: impl Sink<Vec<u8>, Error = io::Error> + Unpin,
        reader: &mut UdpSocketRecvHalf,
    ) -> Result<(), io::Error> {
        loop {
            let mut buf = vec![0; 1024];
            let n = reader.recv(&mut buf[..]).await?;
            if n > 0 {
                stdout.send(buf).await?;
            }
        }
    }
 }
 mod codec {
    use bytes::{BufMut, BytesMut};
    use std::io;
    use tokio::codec::{Decoder, Encoder};
    /// A simple `Codec` implementation that just ships bytes around.
    ///
    /// This type is used for "framing" a TCP/UDP stream of bytes but it's really
    /// just a convenient method for us to work with streams/sinks for now.
    /// This'll just take any data read and interpret it as a "frame" and
    /// conversely just shove data into the output location without looking at
    /// it.
    pub struct Bytes;
    impl Decoder for Bytes {
        type Item = Vec<u8>;
        type Error = io::Error;
        fn decode(&mut self, buf: &mut BytesMut) -> io::Result<Option<Vec<u8>>> {
            if buf.len() > 0 {
                let len = buf.len();
                Ok(Some(buf.split_to(len).into_iter().collect()))
            } else {
                Ok(None)
            }
        }
    }
    impl Encoder for Bytes {
        type Item = Vec<u8>;
        type Error = io::Error;
        fn encode(&mut self, data: Vec<u8>, buf: &mut BytesMut) -> io::Result<()> {
            buf.put(&data[..]);
            Ok(())
        }
    }
 }
@@ -1,252 +0,0 @@
 //! An example of hooking up stdin/stdout to either a TCP or UDP stream.
 //!
 //! This example will connect to a socket address specified in the argument list
 //! and then forward all data read on stdin to the server, printing out all data
 //! received on stdout. An optional `--udp` argument can be passed to specify
 //! that the connection should be made over UDP instead of TCP, translating each
 //! line entered on stdin to a UDP packet to be sent to the remote address.
 //!
 //! Note that this is not currently optimized for performance, especially
 //! around buffer management. Rather it's intended to show an example of
 //! working with a client.
 //!
 //! This example can be quite useful when interacting with the other examples in
 //! this repository! Many of them recommend running this as a simple "hook up
 //! stdin/stdout to a server" to get up and running.
 #![deny(warnings, rust_2018_idioms)]
 use futures::sync::mpsc;
 use std::env;
 use std::io::{self, Read, Write};
 use std::net::SocketAddr;
 use std::thread;
 use tokio;
 use tokio::prelude::*;
 fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Determine if we're going to run in TCP or UDP mode
    let mut args = env::args().skip(1).collect::<Vec<_>>();
    let tcp = match args.iter().position(|a| a == "--udp") {
        Some(i) => {
            args.remove(i);
            false
        }
        None => true,
    };
    // Parse what address we're going to connect to
    let addr = match args.first() {
        Some(addr) => addr,
        None => Err("this program requires at least one argument")?,
    };
    let addr = addr.parse::<SocketAddr>()?;
    // Right now Tokio doesn't support a handle to stdin running on the event
    // loop, so we farm out that work to a separate thread. This thread will
    // read data (with blocking I/O) from stdin and then send it to the event
    // loop over a standard futures channel.
    let (stdin_tx, stdin_rx) = mpsc::channel(0);
    thread::spawn(|| read_stdin(stdin_tx));
    let stdin_rx = stdin_rx.map_err(|_| panic!("errors not possible on rx"));
    // Now that we've got our stdin read we either set up our TCP connection or
    // our UDP connection to get a stream of bytes we're going to emit to
    // stdout.
    let stdout = if tcp {
        tcp::connect(&addr, Box::new(stdin_rx))?
    } else {
        udp::connect(&addr, Box::new(stdin_rx))?
    };
    // And now with our stream of bytes to write to stdout, we execute that in
    // the event loop! Note that this is doing blocking I/O to emit data to
    // stdout, and in general it's a no-no to do that sort of work on the event
    // loop. In this case, though, we know it's ok as the event loop isn't
    // otherwise running anything useful.
    let mut out = io::stdout();
    tokio::run({
        stdout
            .for_each(move |chunk| out.write_all(&chunk))
            .map_err(|e| println!("error reading stdout; error = {:?}", e))
    });
    Ok(())
 }
 mod codec {
    use bytes::{BufMut, BytesMut};
    use std::io;
    use tokio::codec::{Decoder, Encoder};
    /// A simple `Codec` implementation that just ships bytes around.
    ///
    /// This type is used for "framing" a TCP/UDP stream of bytes but it's really
    /// just a convenient method for us to work with streams/sinks for now.
    /// This'll just take any data read and interpret it as a "frame" and
    /// conversely just shove data into the output location without looking at
    /// it.
    pub struct Bytes;
    impl Decoder for Bytes {
        type Item = BytesMut;
        type Error = io::Error;
        fn decode(&mut self, buf: &mut BytesMut) -> io::Result<Option<BytesMut>> {
            if buf.len() > 0 {
                let len = buf.len();
                Ok(Some(buf.split_to(len)))
            } else {
                Ok(None)
            }
        }
    }
    impl Encoder for Bytes {
        type Item = Vec<u8>;
        type Error = io::Error;
        fn encode(&mut self, data: Vec<u8>, buf: &mut BytesMut) -> io::Result<()> {
            buf.put(&data[..]);
            Ok(())
        }
    }
 }
 mod tcp {
    use tokio;
    use tokio::codec::Decoder;
    use tokio::net::TcpStream;
    use tokio::prelude::*;
    use crate::codec::Bytes;
    use bytes::BytesMut;
    use std::error::Error;
    use std::io;
    use std::net::SocketAddr;
    pub fn connect(
        addr: &SocketAddr,
        stdin: Box<dyn Stream<Item = Vec<u8>, Error = io::Error> + Send>,
    ) -> Result<Box<dyn Stream<Item = BytesMut, Error = io::Error> + Send>, Box<dyn Error>> {
        let tcp = TcpStream::connect(addr);
        // After the TCP connection has been established, we set up our client
        // to start forwarding data.
        //
        // First we use the `Io::framed` method with a simple implementation of
        // a `Codec` (listed below) that just ships bytes around. We then split
        // that in two to work with the stream and sink separately.
        //
        // Half of the work we're going to do is to take all data we receive on
        // `stdin` and send that along the TCP stream (`sink`). The second half
        // is to take all the data we receive (`stream`) and then write that to
        // stdout. We'll be passing this handle back out from this method.
        //
        // You'll also note that we *spawn* the work to read stdin and write it
        // to the TCP stream. This is done to ensure that happens concurrently
        // with us reading data from the stream.
        let stream = Box::new(
            tcp.map(move |stream| {
                let (sink, stream) = Bytes.framed(stream).split();
                tokio::spawn(stdin.forward(sink).then(|result| {
                    if let Err(e) = result {
                        println!("failed to write to socket: {}", e)
                    }
                    Ok(())
                }));
                stream
            })
            .flatten_stream(),
        );
        Ok(stream)
    }
 }
 mod udp {
    use std::error::Error;
    use std::io;
    use std::net::SocketAddr;
    use bytes::BytesMut;
    use tokio;
    use tokio::net::{UdpFramed, UdpSocket};
    use tokio::prelude::*;
    use crate::codec::Bytes;
    pub fn connect(
        &addr: &SocketAddr,
        stdin: Box<dyn Stream<Item = Vec<u8>, Error = io::Error> + Send>,
    ) -> Result<Box<dyn Stream<Item = BytesMut, Error = io::Error> + Send>, Box<dyn Error>> {
        // We'll bind our UDP socket to a local IP/port, but for now we
        // basically let the OS pick both of those.
        let addr_to_bind = if addr.ip().is_ipv4() {
            "0.0.0.0:0".parse()?
        } else {
            "[::]:0".parse()?
        };
        let udp = match UdpSocket::bind(&addr_to_bind) {
            Ok(udp) => udp,
            Err(_) => Err("failed to bind socket")?,
        };
        // Like above with TCP we use an instance of `Bytes` codec to transform
        // this UDP socket into a framed sink/stream which operates over
        // discrete values. In this case we're working with *pairs* of socket
        // addresses and byte buffers.
        let (sink, stream) = UdpFramed::new(udp, Bytes).split();
        // All bytes from `stdin` will go to the `addr` specified in our
        // argument list. Like with TCP this is spawned concurrently
        let forward_stdin = stdin
            .map(move |chunk| (chunk, addr))
            .forward(sink)
            .then(|result| {
                if let Err(e) = result {
                    println!("failed to write to socket: {}", e)
                }
                Ok(())
            });
        // With UDP we could receive data from any source, so filter out
        // anything coming from a different address
        let receive = stream.filter_map(move |(chunk, src)| {
            if src == addr {
                Some(chunk.into())
            } else {
                None
            }
        });
        let stream = Box::new(
            future::lazy(|| {
                tokio::spawn(forward_stdin);
                future::ok(receive)
            })
            .flatten_stream(),
        );
        Ok(stream)
    }
 }
 // Our helper method which will read data from stdin and send it along the
 // sender provided.
 fn read_stdin(mut tx: mpsc::Sender<Vec<u8>>) {
    let mut stdin = io::stdin();
    loop {
        let mut buf = vec![0; 1024];
        let n = match stdin.read(&mut buf) {
            Err(_) | Ok(0) => break,
            Ok(n) => n,
        };
        buf.truncate(n);
        tx = match tx.send(buf).wait() {
            Ok(tx) => tx,
            Err(_) => break,
        };
    }
 }