diff --git a/examples/client.rs b/examples/client.rs new file mode 100644 index 0000000..16324e4 --- /dev/null +++ b/examples/client.rs @@ -0,0 +1,116 @@ +//! A simple example of hooking up stdin/stdout to a WebSocket stream. +//! +//! This example will connect to a server specified in the argument list and +//! then forward all data read on stdin to the server, printing out all data +//! received on stdout. +//! +//! 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. +//! +//! You can use this example together with the `server` example. + +extern crate futures; +extern crate tokio_core; +extern crate tokio_tungstenite; +extern crate tungstenite; +extern crate url; + +use std::env; +use std::io::{self, Read, Write}; +use std::net::ToSocketAddrs; +use std::thread; + +use futures::sync::mpsc; +use futures::{Future, Sink, Stream}; +use tokio_core::net::TcpStream; +use tokio_core::reactor::Core; +use tokio_tungstenite::ClientHandshakeExt; +use tungstenite::handshake::client::{ClientHandshake, Request}; +use tungstenite::protocol::Message; + +fn main() { + // Specify the server address to which the client will be connecting. + let connect_addr = env::args().nth(1).unwrap_or_else(|| { + panic!("this program requires at least one argument") + }); + + // Get a first IP address of the server from the server URL. + let url = url::Url::parse(&connect_addr).unwrap(); + let addr = url.to_socket_addrs().unwrap().next().unwrap(); + + // Create the event loop and initiate the connection to the remote server. + let mut core = Core::new().unwrap(); + let handle = core.handle(); + let tcp = TcpStream::connect(&addr, &handle); + + // 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 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 + + // After the TCP connection has been established, we set up our client to + // start forwarding data. + // + // First we do a WebSocket handshake on a TCP stream, i.e. do the upgrade + // request. + // + // Half of the work we're going to do is to take all data we receive on + // stdin (`stdin_rx`) and send that along the WebSocket stream (`sink`). + // The second half is to take all the data we receive (`stream`) and then + // write that to stdout. Currently we just write to stdout in a synchronous + // fashion. + // + // Finally we set the client to terminate once either half of this work + // finishes. If we don't have any more data to read or we won't receive any + // more work from the remote then we can exit. + let mut stdout = io::stdout(); + let client = tcp.and_then(|stream| { + let req = Request { url: url }; + ClientHandshake::::new_async(stream, req).and_then(|ws_stream| { + println!("WebSocket handshake has been successfully completed"); + + // `sink` is the stream of messages going out. + // `stream` is the stream of incoming messages. + let (sink, stream) = ws_stream.split(); + + // We forward all messages, composed out of the data, entered to + // the stdin, to the `sink`. + let send_stdin = stdin_rx.forward(sink); + let write_stdout = stream.for_each(|message| { + stdout.write_all(&message.into_data()).unwrap(); + Ok(()) + }); + + // Wait for either of futures to complete. + send_stdin.map(|_| ()) + .select(write_stdout.map(|_| ())) + .then(|_| Ok(())) + }).map_err(|e| { + println!("Error during the websocket handshake occurred: {}", e); + io::Error::new(io::ErrorKind::Other, e) + }) + }); + + // And now that we've got our client, we execute it in the event loop! + core.run(client).unwrap(); +} + +// Our helper method which will read data from stdin and send it along the +// sender provided. +fn read_stdin(mut tx: mpsc::Sender) { + 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 = tx.send(Message::binary(buf)).wait().unwrap(); + } +} diff --git a/examples/server.rs b/examples/server.rs new file mode 100644 index 0000000..e23de7b --- /dev/null +++ b/examples/server.rs @@ -0,0 +1,123 @@ +//! A chat server that broadcasts a message to all connections. +//! +//! This is a simple line-based server which accepts WebSocket connections, +//! reads lines from those connections, and broadcasts the lines to all other +//! connected clients. +//! +//! You can test this out by running: +//! +//! cargo run --example server +//! +//! And then in another window run: +//! +//! cargo run --example client ws://127.0.0.1:12345/ +//! +//! You can run the second command in multiple windows and then chat between the +//! two, seeing the messages from the other client as they're received. For all +//! connected clients they'll all join the same room and see everyone else's +//! messages. + +extern crate futures; +extern crate tokio_core; +extern crate tokio_tungstenite; +extern crate tungstenite; + +use std::cell::RefCell; +use std::collections::HashMap; +use std::env; +use std::io::{Error, ErrorKind}; +use std::rc::Rc; + +use futures::stream::Stream; +use futures::{Future}; +use tokio_core::net::{TcpListener, TcpStream}; +use tokio_core::reactor::Core; +use tokio_tungstenite::ServerHandshakeExt; +use tungstenite::handshake::server::ServerHandshake; +use tungstenite::protocol::Message; + +fn main() { + let addr = env::args().nth(1).unwrap_or("127.0.0.1:8080".to_string()); + let addr = addr.parse().unwrap(); + + // Create the event loop and TCP listener we'll accept connections on. + let mut core = Core::new().unwrap(); + let handle = core.handle(); + let socket = TcpListener::bind(&addr, &handle).unwrap(); + println!("Listening on: {}", addr); + + // This is a single-threaded server, so we can just use Rc and RefCell to + // store the map of all connections we know about. + let connections = Rc::new(RefCell::new(HashMap::new())); + + let srv = socket.incoming().for_each(|(stream, addr)| { + + // We have to clone both of these values, because the `and_then` + // function billow constructs a new future, `and_then` requires + // `FnOnce`, so we construct a move closure to move the + // environment inside the future (AndThen future may overlive our + // `for_each` future). + let connections_inner = connections.clone(); + let handle_inner = handle.clone(); + + ServerHandshake::::new_async(stream).and_then(move |ws_stream| { + println!("New WebSocket connection: {}", addr); + + // Create a channel for our stream, which other sockets will use to + // send us messages. Then register our address with the stream to send + // data to us. + let (tx, rx) = futures::sync::mpsc::unbounded(); + connections_inner.borrow_mut().insert(addr, tx); + + // Let's split the WebSocket stream, so we can work with the + // reading and writing halves separately. + let (sink, stream) = ws_stream.split(); + + // Whenever we receive a message from the client, we print it and + // send to other clients, excluding the sender. + let connections = connections_inner.clone(); + let ws_reader = stream.for_each(move |message: Message| { + println!("Received a message from {}: {}", addr, message); + + // For each open connection except the sender, send the + // string via the channel. + let mut conns = connections.borrow_mut(); + let iter = conns.iter_mut() + .filter(|&(&k, _)| k != addr) + .map(|(_, v)| v); + for tx in iter { + tx.send(message.clone()).unwrap(); + } + Ok(()) + }); + + // Whenever we receive a string on the Receiver, we write it to + // `WriteHalf`. + let ws_writer = rx.fold(sink, |mut sink, msg| { + use futures::Sink; + sink.start_send(msg).unwrap(); + Ok(sink) + }); + + // Now that we've got futures representing each half of the socket, we + // use the `select` combinator to wait for either half to be done to + // tear down the other. Then we spawn off the result. + let connection = ws_reader.map(|_| ()).map_err(|_| ()) + .select(ws_writer.map(|_| ()).map_err(|_| ())); + + handle_inner.spawn(connection.then(move |_| { + connections_inner.borrow_mut().remove(&addr); + println!("Connection {} closed.", addr); + Ok(()) + })); + + Ok(()) + }).map_err(|e| { + println!("Error during the websocket handshake occurred: {}", e); + Error::new(ErrorKind::Other, e) + }) + }); + + // Execute server. + core.run(srv).unwrap(); +}