Ahmed Charles
2 years ago
committed by
Sebastian Dröge
parent
31dac03929
commit
2e61ecbbb1
@ -0,0 +1,190 @@ |
||||
//! 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 127.0.0.1:12345
|
||||
//!
|
||||
//! And then in another window run:
|
||||
//!
|
||||
//! cargo run --example client ws://127.0.0.1:12345/socket
|
||||
//!
|
||||
//! 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.
|
||||
|
||||
use std::{ |
||||
collections::HashMap, |
||||
convert::Infallible, |
||||
env, |
||||
net::SocketAddr, |
||||
sync::{Arc, Mutex}, |
||||
}; |
||||
|
||||
use hyper::{ |
||||
header::{ |
||||
HeaderValue, CONNECTION, SEC_WEBSOCKET_ACCEPT, SEC_WEBSOCKET_KEY, SEC_WEBSOCKET_VERSION, |
||||
UPGRADE, |
||||
}, |
||||
server::conn::AddrStream, |
||||
service::{make_service_fn, service_fn}, |
||||
upgrade::Upgraded, |
||||
Body, Method, Request, Response, Server, StatusCode, Version, |
||||
}; |
||||
|
||||
use futures_channel::mpsc::{unbounded, UnboundedSender}; |
||||
use futures_util::{future, pin_mut, stream::TryStreamExt, StreamExt}; |
||||
|
||||
use async_tungstenite::{tokio::TokioAdapter, WebSocketStream}; |
||||
use tungstenite::{ |
||||
handshake::derive_accept_key, |
||||
protocol::{Message, Role}, |
||||
}; |
||||
|
||||
type Tx = UnboundedSender<Message>; |
||||
type PeerMap = Arc<Mutex<HashMap<SocketAddr, Tx>>>; |
||||
|
||||
async fn handle_connection( |
||||
peer_map: PeerMap, |
||||
ws_stream: WebSocketStream<TokioAdapter<Upgraded>>, |
||||
addr: SocketAddr, |
||||
) { |
||||
println!("WebSocket connection established: {}", addr); |
||||
|
||||
// Insert the write part of this peer to the peer map.
|
||||
let (tx, rx) = unbounded(); |
||||
peer_map.lock().unwrap().insert(addr, tx); |
||||
|
||||
let (outgoing, incoming) = ws_stream.split(); |
||||
|
||||
let broadcast_incoming = incoming.try_for_each(|msg| { |
||||
println!( |
||||
"Received a message from {}: {}", |
||||
addr, |
||||
msg.to_text().unwrap() |
||||
); |
||||
let peers = peer_map.lock().unwrap(); |
||||
|
||||
// We want to broadcast the message to everyone except ourselves.
|
||||
let broadcast_recipients = peers |
||||
.iter() |
||||
.filter(|(peer_addr, _)| peer_addr != &&addr) |
||||
.map(|(_, ws_sink)| ws_sink); |
||||
|
||||
for recp in broadcast_recipients { |
||||
recp.unbounded_send(msg.clone()).unwrap(); |
||||
} |
||||
|
||||
future::ok(()) |
||||
}); |
||||
|
||||
let receive_from_others = rx.map(Ok).forward(outgoing); |
||||
|
||||
pin_mut!(broadcast_incoming, receive_from_others); |
||||
future::select(broadcast_incoming, receive_from_others).await; |
||||
|
||||
println!("{} disconnected", &addr); |
||||
peer_map.lock().unwrap().remove(&addr); |
||||
} |
||||
|
||||
async fn handle_request( |
||||
peer_map: PeerMap, |
||||
mut req: Request<Body>, |
||||
addr: SocketAddr, |
||||
) -> Result<Response<Body>, Infallible> { |
||||
println!("Received a new, potentially ws handshake"); |
||||
println!("The request's path is: {}", req.uri().path()); |
||||
println!("The request's headers are:"); |
||||
for (ref header, _value) in req.headers() { |
||||
println!("* {}", header); |
||||
} |
||||
let upgrade = HeaderValue::from_static("Upgrade"); |
||||
let websocket = HeaderValue::from_static("websocket"); |
||||
let headers = req.headers(); |
||||
let key = headers.get(SEC_WEBSOCKET_KEY); |
||||
let derived = key.map(|k| derive_accept_key(k.as_bytes())); |
||||
if req.method() != Method::GET |
||||
|| req.version() < Version::HTTP_11 |
||||
|| !headers |
||||
.get(CONNECTION) |
||||
.and_then(|h| h.to_str().ok()) |
||||
.map(|h| { |
||||
h.split(|c| c == ' ' || c == ',') |
||||
.any(|p| p.eq_ignore_ascii_case(upgrade.to_str().unwrap())) |
||||
}) |
||||
.unwrap_or(false) |
||||
|| !headers |
||||
.get(UPGRADE) |
||||
.and_then(|h| h.to_str().ok()) |
||||
.map(|h| h.eq_ignore_ascii_case("websocket")) |
||||
.unwrap_or(false) |
||||
|| !headers |
||||
.get(SEC_WEBSOCKET_VERSION) |
||||
.map(|h| h == "13") |
||||
.unwrap_or(false) |
||||
|| key.is_none() |
||||
|| req.uri() != "/socket" |
||||
{ |
||||
return Ok(Response::new(Body::from("Hello World!"))); |
||||
} |
||||
let ver = req.version(); |
||||
tokio::task::spawn(async move { |
||||
match hyper::upgrade::on(&mut req).await { |
||||
Ok(upgraded) => { |
||||
handle_connection( |
||||
peer_map, |
||||
WebSocketStream::from_raw_socket( |
||||
TokioAdapter::new(upgraded), |
||||
Role::Server, |
||||
None, |
||||
) |
||||
.await, |
||||
addr, |
||||
) |
||||
.await; |
||||
} |
||||
Err(e) => println!("upgrade error: {}", e), |
||||
} |
||||
}); |
||||
let mut res = Response::new(Body::empty()); |
||||
*res.status_mut() = StatusCode::SWITCHING_PROTOCOLS; |
||||
*res.version_mut() = ver; |
||||
res.headers_mut().append(CONNECTION, upgrade); |
||||
res.headers_mut().append(UPGRADE, websocket); |
||||
res.headers_mut() |
||||
.append(SEC_WEBSOCKET_ACCEPT, derived.unwrap().parse().unwrap()); |
||||
// Let's add an additional header to our response to the client.
|
||||
res.headers_mut() |
||||
.append("MyCustomHeader", ":)".parse().unwrap()); |
||||
res.headers_mut() |
||||
.append("SOME_TUNGSTENITE_HEADER", "header_value".parse().unwrap()); |
||||
Ok(res) |
||||
} |
||||
|
||||
#[tokio::main] |
||||
async fn main() -> Result<(), hyper::Error> { |
||||
let state = PeerMap::new(Mutex::new(HashMap::new())); |
||||
|
||||
let addr = env::args() |
||||
.nth(1) |
||||
.unwrap_or_else(|| "127.0.0.1:8080".to_string()) |
||||
.parse() |
||||
.unwrap(); |
||||
|
||||
let make_svc = make_service_fn(move |conn: &AddrStream| { |
||||
let remote_addr = conn.remote_addr(); |
||||
let state = state.clone(); |
||||
let service = service_fn(move |req| handle_request(state.clone(), req, remote_addr)); |
||||
async { Ok::<_, Infallible>(service) } |
||||
}); |
||||
|
||||
let server = Server::bind(&addr).serve(make_svc); |
||||
|
||||
server.await?; |
||||
|
||||
Ok::<_, hyper::Error>(()) |
||||
} |
||||
Loading…
Reference in new issue