blatherskite

commit d100662bbdaa4d128905007b45388cd92a6440a7
parent c8446e0d548ae1c1241d3ac34fe9051f726aa700
Author: quantumish <freifeld.david@gmail.com>
Date:   Sat, 24 Sep 2022 18:14:22 -0700

Add barebones websocket implementation from Poem examples

Diffstat:
M
chatterbox/Cargo.toml
 | 
4
++++
M
chatterbox/src/main.rs
 | 
60
++++++++++++++++++++++++++++++++++++++++++++++++++++++++++--

2 files changed, 62 insertions(+), 2 deletions(-)
diff --git a/chatterbox/Cargo.toml b/chatterbox/Cargo.toml
@@ -6,3 +6,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+futures-util = "0.3.24"
+poem = { version = "1.3.43", features = ["websocket"] }
+tokio = { version = "1.21.1", features = ["full"] }
+tracing-subscriber = "0.3.15"
diff --git a/chatterbox/src/main.rs b/chatterbox/src/main.rs
@@ -1,3 +1,59 @@
-fn main() {
-    println!("Hello, world!");
+/// Currently a modified version of `poem`'s default websocket-chat example
+use futures_util::{SinkExt, StreamExt};
+use poem::{
+    get, handler,
+    listener::TcpListener,
+    web::{
+        websocket::{Message, WebSocket},
+        Data, Path,
+    },
+    EndpointExt, IntoResponse, Route, Server,
+};
+
+#[handler]
+fn ws(
+    Path(name): Path<String>,
+    ws: WebSocket,
+    sender: Data<&tokio::sync::broadcast::Sender<String>>,
+) -> impl IntoResponse {
+    let sender = sender.clone();
+    let mut receiver = sender.subscribe();
+    ws.on_upgrade(move |socket| async move {
+        let (mut sink, mut stream) = socket.split();
+
+        tokio::spawn(async move {
+            while let Some(Ok(msg)) = stream.next().await {
+                if let Message::Text(text) = msg {
+                    if sender.send(format!("{}: {}", name, text)).is_err() {
+                        break;
+                    }
+                }
+            }
+        });
+
+        tokio::spawn(async move {
+            while let Ok(msg) = receiver.recv().await {
+                if sink.send(Message::Text(msg)).await.is_err() {
+                    break;
+                }
+            }
+        });
+    })
+}
+
+#[tokio::main]
+async fn main() -> Result<(), std::io::Error> {
+    if std::env::var_os("RUST_LOG").is_none() {
+        std::env::set_var("RUST_LOG", "poem=debug");
+    }
+    tracing_subscriber::fmt::init();
+
+    let app = Route::new().at(
+        "/ws/:name",
+        get(ws.data(tokio::sync::broadcast::channel::<String>(32).0)),
+    );
+
+    Server::new(TcpListener::bind("127.0.0.1:3001"))
+        .run(app)
+        .await
 }