From ac5838f9b78c6894751cfacaf2f238e09a4a7f51 Mon Sep 17 00:00:00 2001
From: Martin Pitt <martin@piware.de>
Date: Sat, 12 Nov 2022 10:54:47 +0100
Subject: [PATCH] axum-server: Add websocket route

---
 axum-server/Cargo.toml  |  2 +-
 axum-server/src/main.rs | 30 +++++++++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/axum-server/Cargo.toml b/axum-server/Cargo.toml
index 978ee06..0444b39 100644
--- a/axum-server/Cargo.toml
+++ b/axum-server/Cargo.toml
@@ -6,7 +6,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-axum = { version = "0.5", features = ["headers"] }
+axum = { version = "0.5", features = ["ws", "headers"] }
 tokio = { version = "1", features = ["full"] }
 tower = "0.4"
 tower-http = { version = "0.3", features = ["trace", "fs"] }
diff --git a/axum-server/src/main.rs b/axum-server/src/main.rs
index 3212d3f..31c0a4f 100644
--- a/axum-server/src/main.rs
+++ b/axum-server/src/main.rs
@@ -2,7 +2,7 @@ use std::io;
 
 use axum::{
     routing::{get, get_service},
-    extract::{Path, TypedHeader},
+    extract::{Path, TypedHeader, ws},
     http::{StatusCode},
     response,
     Router};
@@ -15,6 +15,33 @@ async fn hello(Path(name): Path<String>, user_agent: Option<TypedHeader<axum::he
     }
 }
 
+async fn ws_echo(mut socket: ws::WebSocket) {
+    while let Some(msg) = socket.recv().await {
+        if let Ok(msg) = msg {
+            tracing::debug!("websocket got message: {:?}", msg);
+
+            let reply = match msg  {
+                ws::Message::Text(t) => ws::Message::Text(t),
+                ws::Message::Binary(b) => ws::Message::Binary(b),
+                // axum handles Ping/Pong by itself
+                ws::Message::Ping(_) => { continue },
+                ws::Message::Pong(_) => { continue },
+                ws::Message::Close(_) => { break }
+            };
+
+            if socket.send(reply).await
+                .is_err() {
+                    tracing::info!("websocket client disconnected");
+                    break;
+                }
+        }
+        else {
+            tracing::info!("websocket client disconnected");
+            break;
+        }
+    }
+}
+
 #[tokio::main]
 async fn main() {
     tracing_subscriber::fmt::init();
@@ -26,6 +53,7 @@ async fn main() {
                        (StatusCode::INTERNAL_SERVER_ERROR, format!("Unhandled internal error: {}", e))
                    })
         )
+        .route("/ws-echo", get(|ws: ws::WebSocketUpgrade| async {ws.on_upgrade(ws_echo)}))
         .layer(
             tower::ServiceBuilder::new()
                 .layer(
-- 
2.39.5