create server-headers example (#269)

* create server-headers example

* fixed fmt errors

Cargo.toml

@@ -95,6 +95,10 @@ required-features = ["handshake"]
 name = "server"
 required-features = ["handshake"]
 
+[[example]]
+name = "server-headers"
+required-features = ["handshake"]
+
 [[example]]
 name = "interval-server"
 required-features = ["handshake"]

examples/README.md

@@ -1,8 +1,9 @@
 Examples
 
-* [autobahn-client.rs](https://github.com/snapview/tokio-tungstenite/blob/master/examples/autobahn-client.rs)
-* [autobahn-server.rs](https://github.com/snapview/tokio-tungstenite/blob/master/examples/autobahn-server.rs)
-* [client.rs](https://github.com/snapview/tokio-tungstenite/blob/master/examples/client.rs)
-* [echo-server.rs](https://github.com/snapview/tokio-tungstenite/blob/master/examples/echo-server.rs)
-* [server.rs](https://github.com/snapview/tokio-tungstenite/blob/master/examples/server.rs)
-* [interval-server.rs](https://github.com/snapview/tokio-tungstenite/blob/master/examples/interval-server.rs)
+- [autobahn-client.rs](https://github.com/snapview/tokio-tungstenite/blob/master/examples/autobahn-client.rs)
+- [autobahn-server.rs](https://github.com/snapview/tokio-tungstenite/blob/master/examples/autobahn-server.rs)
+- [client.rs](https://github.com/snapview/tokio-tungstenite/blob/master/examples/client.rs)
+- [echo-server.rs](https://github.com/snapview/tokio-tungstenite/blob/master/examples/echo-server.rs)
+- [server.rs](https://github.com/snapview/tokio-tungstenite/blob/master/examples/server.rs)
+- [server-headers.rs](https://github.com/snapview/tokio-tungstenite/blob/master/examples/server-headers.rs)
+- [interval-server.rs](https://github.com/snapview/tokio-tungstenite/blob/master/examples/interval-server.rs)

examples/server-headers.rs

@@ -0,0 +1,86 @@
+///! Read/Write headers on server example
+///!
+///! Run with logs:
+///! Linux:
+///! ```sh
+///! RUST_LOG=debug cargo run --example server-headers
+///! ```
+///! Windows
+///! ```sh
+///! cmd /c "set RUST_LOG=debug && cargo run --example server-headers"
+///! ```
+use tokio::net::{TcpListener, TcpStream};
+use tokio_tungstenite::{
+    accept_hdr_async,
+    tungstenite::{
+        connect,
+        handshake::server::{Request, Response},
+        Message,
+    },
+};
+use url::Url;
+#[macro_use]
+extern crate log;
+use futures_util::{SinkExt, StreamExt};
+
+#[tokio::main]
+async fn main() {
+    env_logger::builder().format_timestamp(None).init();
+
+    tokio::spawn(async move {
+        server().await;
+    });
+    client();
+}
+
+async fn server() {
+    let server = TcpListener::bind("127.0.0.1:8080").await.unwrap();
+
+    while let Ok((stream, _)) = server.accept().await {
+        tokio::spawn(accept_connection(stream));
+    }
+}
+
+async fn accept_connection(stream: TcpStream) {
+    let callback = |req: &Request, mut response: Response| {
+        debug!("Received a new ws handshake");
+        debug!("The request's path is: {}", req.uri().path());
+        debug!("The request's headers are:");
+        for (ref header, _value) in req.headers() {
+            debug!("* {}: {:?}", header, _value);
+        }
+
+        let headers = response.headers_mut();
+        headers.append("MyCustomHeader", ":)".parse().unwrap());
+
+        Ok(response)
+    };
+    let mut ws_stream = accept_hdr_async(stream, callback)
+        .await
+        .expect("Error during the websocket handshake occurred");
+
+    while let Some(msg) = ws_stream.next().await {
+        let msg = msg.unwrap();
+        if msg.is_text() || msg.is_binary() {
+            debug!("Server on message: {:?}", &msg);
+            ws_stream.send(msg).await.unwrap();
+        }
+    }
+}
+
+fn client() {
+    let (mut socket, response) =
+        connect(Url::parse("ws://localhost:8080/socket").unwrap()).expect("Can't connect");
+    debug!("Connected to the server");
+    debug!("Response HTTP code: {}", response.status());
+    debug!("Response contains the following headers:");
+    for (ref header, _value) in response.headers() {
+        debug!("* {}: {:?}", header, _value);
+    }
+
+    socket.write_message(Message::Text("Hello WebSocket".into())).unwrap();
+    loop {
+        let msg = socket.read_message().expect("Error reading message");
+        debug!("Received: {}", msg);
+    }
+}