From 883c4d01f6138070672c08118e2be9facdcd2e15 Mon Sep 17 00:00:00 2001
From: Martin Pitt <martin@piware.de>
Date: Sat, 10 Dec 2022 12:50:51 +0100
Subject: [PATCH] actix-server: Add echo websocket route

---
 actix-server/Cargo.toml  |  6 +++++
 actix-server/src/main.rs | 55 ++++++++++++++++++++++++++++++++++++++--
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/actix-server/Cargo.toml b/actix-server/Cargo.toml
index 803f53e..4208b99 100644
--- a/actix-server/Cargo.toml
+++ b/actix-server/Cargo.toml
@@ -6,7 +6,13 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+actix = "0.13"
 actix-web = "4"
+actix-web-actors = "4.1"
 actix-files = "0.6"
 env_logger = "0.9"
 log = "0.4"
+
+[dev-dependencies]
+actix-test = "0.1"
+futures-util = "0.3"
diff --git a/actix-server/src/main.rs b/actix-server/src/main.rs
index 4a43ef5..58756de 100644
--- a/actix-server/src/main.rs
+++ b/actix-server/src/main.rs
@@ -1,9 +1,11 @@
 use std::path::Path;
 
-use actix_web::{get, route, web, App, HttpRequest, HttpServer, Responder, Result};
+use actix::{Actor, ActorContext, StreamHandler};
+use actix_web::{get, route, web, App, Error, HttpRequest, HttpResponse, HttpServer, Responder, Result};
 use actix_web::http::header;
 use actix_web::middleware::Logger;
 use actix_files::{Files, NamedFile};
+use actix_web_actors::ws;
 
 #[route("/hello/{name}", method="GET", method="HEAD")]
 async fn hello(params: web::Path<String>, req: HttpRequest) -> Result<String> {
@@ -37,6 +39,33 @@ async fn static_file(params: web::Path<String>, req: HttpRequest) -> Result<impl
     Ok(NamedFile::open_async(disk_path).await?.customize())
 }
 
+struct WsEcho;
+
+impl Actor for WsEcho {
+    type Context = ws::WebsocketContext<Self>;
+}
+
+impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WsEcho {
+    fn handle(&mut self, msg: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
+        log::info!("WsEcho got message {:?}", msg);
+        match msg {
+            Ok(ws::Message::Ping(msg)) => ctx.pong(&msg),
+            Ok(ws::Message::Text(text)) => ctx.text(text),
+            Ok(ws::Message::Binary(bin)) => ctx.binary(bin),
+            Ok(ws::Message::Close(reason)) => {
+                ctx.close(reason);
+                ctx.stop();
+            },
+            _ => ctx.stop(),
+        }
+    }
+}
+
+#[get("/ws-echo")]
+async fn ws_echo(req: HttpRequest, stream: web::Payload) -> Result<HttpResponse, Error> {
+    ws::start(WsEcho {}, &req, stream)
+}
+
 #[actix_web::main]
 async fn main() -> std::io::Result<()> {
     env_logger::init_from_env(env_logger::Env::default().default_filter_or("info"));
@@ -46,6 +75,7 @@ async fn main() -> std::io::Result<()> {
             .service(hello)
             .service(static_file)
             .service(Files::new("/dir", "../static"))
+            .service(ws_echo)
             .wrap(Logger::default())
     })
         .bind(("127.0.0.1", 3030))?
@@ -57,8 +87,12 @@ async fn main() -> std::io::Result<()> {
 mod tests {
     use actix_web::{App, body, test, web};
     use actix_web::http::{header, StatusCode};
+    use actix_web_actors::ws;
 
-    use super::{hello, static_file};
+    use futures_util::sink::SinkExt;
+    use futures_util::StreamExt;
+
+    use super::{hello, static_file, ws_echo};
 
     #[actix_web::test]
     async fn test_hello() {
@@ -132,4 +166,21 @@ mod tests {
         assert_eq!(res_bytes.len(), 63); // file size of ../static/dir1/optzip.txt.gz
         assert_eq!(res_bytes[0], 31);
     }
+
+    #[actix_web::test]
+    async fn test_ws_echo() {
+        // FIXME: duplicating the .service() call from main() here is super ugly, but it's hard to move that into a fn
+        let mut srv = actix_test::start(|| App::new().service(ws_echo));
+        let mut client = srv.ws_at("/ws-echo").await.unwrap();
+
+        // text echo
+        client.send(ws::Message::Text("hello".into())).await.unwrap();
+        let received = client.next().await.unwrap().unwrap();
+        assert_eq!(received, ws::Frame::Text("hello".into()));
+
+        // binary echo
+        client.send(ws::Message::Binary(web::Bytes::from_static(&[42, 99]))).await.unwrap();
+        let received = client.next().await.unwrap().unwrap();
+        assert_eq!(received, ws::Frame::Binary(web::Bytes::from_static(&[42, 99])));
+    }
 }
-- 
2.39.2