├── .dockerignore ├── dockerfile ├── .gitignore ├── Cargo.toml ├── Makefile ├── README.md ├── LICENSE ├── .vscode └── launch.json └── src ├── ws.rs ├── main.rs └── handler.rs /.dockerignore: -------------------------------------------------------------------------------- 1 | Cargo.lock -------------------------------------------------------------------------------- /dockerfile: -------------------------------------------------------------------------------- 1 | FROM rust:1.73 AS build 2 | WORKDIR /usr/src/wsserver 3 | 4 | COPY . . 5 | 6 | RUN make 7 | 8 | FROM debian:buster-slim 9 | WORKDIR /usr/local/bin 10 | 11 | COPY --from=build /usr/src/wsserver/target/release/wsserver . 12 | 13 | EXPOSE 8080 14 | 15 | CMD ["./wsserver"] -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Generated by Cargo 2 | # will have compiled files and executables 3 | /target/ 4 | 5 | # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries 6 | # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html 7 | Cargo.lock 8 | 9 | # These are backup files generated by rustfmt 10 | **/*.rs.bk 11 | 12 | 13 | #Added by cargo 14 | 15 | /target 16 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | 2 | [package] 3 | name = "wsserver" 4 | version = "0.1.1" 5 | authors = ["Mario Zupan "] 6 | edition = "2021" 7 | 8 | [dependencies] 9 | tokio = { version = "1.19.2", features = ["macros", "sync", "rt-multi-thread"] } 10 | tokio-stream = "0.1.9" 11 | warp = "0.3" 12 | serde = {version = "1.0", features = ["derive"] } 13 | serde_json = "1.0" 14 | futures = { version = "0.3", default-features = false } 15 | uuid = { version = "1.1.2", features = ["serde", "v4"] } 16 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | build: 2 | @cargo build --release 3 | 4 | clean: 5 | @cargo clean 6 | 7 | TESTS = "" 8 | test: 9 | @cargo test $(TESTS) --offline --lib -- --color=always --nocapture 10 | 11 | docs: build 12 | @cargo doc --no-deps 13 | 14 | style-check: 15 | @rustup component add rustfmt 2> /dev/null 16 | cargo fmt --all -- --check 17 | 18 | lint: 19 | @rustup component add clippy 2> /dev/null 20 | cargo clippy --all-targets --all-features -- -D warnings 21 | 22 | dev: 23 | cargo run 24 | 25 | .PHONY: build test docs style-check lint 26 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # warp-websockets-example 2 | Basic example using websockets with warp in Rust 3 | 4 | Run with 5 | 6 | ```bash 7 | make dev 8 | ``` 9 | 10 | Then, you can register/unregister a client: 11 | 12 | ```bash 13 | curl -X POST 'http://localhost:8000/register' -H 'Content-Type: application/json' -d '{ "user_id": 1 }' 14 | 15 | curl -X DELETE 'http://localhost:8000/register/e2fa90682255472b9221709566dbceba' 16 | ``` 17 | 18 | Or connect to the WebSocket using the returned URL: `ws://127.0.0.1:8000/ws/625ac78b88e047a1bc7b3f8459702078`. 19 | 20 | Then, you can publish messages using 21 | 22 | ```bash 23 | curl -X POST 'http://localhost:8000/publish' \ 24 | -H 'Content-Type: application/json' \ 25 | -d '{"user_id": 1, "topic": "cats", "message": "are awesome"}' 26 | ``` 27 | 28 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 mario 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /.vscode/launch.json: -------------------------------------------------------------------------------- 1 | { 2 | // Use IntelliSense to learn about possible attributes. 3 | // Hover to view descriptions of existing attributes. 4 | // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 5 | "version": "0.2.0", 6 | "configurations": [ 7 | { 8 | "type": "lldb", 9 | "request": "launch", 10 | "name": "Debug executable 'warp-websockets-example'", 11 | "cargo": { 12 | "args": [ 13 | "build", 14 | "--bin=warp-websockets-example", 15 | "--package=warp-websockets-example" 16 | ], 17 | "filter": { 18 | "name": "warp-websockets-example", 19 | "kind": "bin" 20 | } 21 | }, 22 | "args": [], 23 | "cwd": "${workspaceFolder}" 24 | }, 25 | { 26 | "type": "lldb", 27 | "request": "launch", 28 | "name": "Debug unit tests in executable 'warp-websockets-example'", 29 | "cargo": { 30 | "args": [ 31 | "test", 32 | "--no-run", 33 | "--bin=warp-websockets-example", 34 | "--package=warp-websockets-example" 35 | ], 36 | "filter": { 37 | "name": "warp-websockets-example", 38 | "kind": "bin" 39 | } 40 | }, 41 | "args": [], 42 | "cwd": "${workspaceFolder}" 43 | } 44 | ] 45 | } -------------------------------------------------------------------------------- /src/ws.rs: -------------------------------------------------------------------------------- 1 | use crate::{Client, Clients}; 2 | use futures::{FutureExt, StreamExt}; 3 | use serde::Deserialize; 4 | use serde_json::from_str; 5 | use tokio::sync::mpsc; 6 | use tokio_stream::wrappers::UnboundedReceiverStream; 7 | use warp::ws::{Message, WebSocket}; 8 | 9 | #[derive(Deserialize, Debug)] 10 | pub struct TopicsRequest { 11 | topics: Vec, 12 | } 13 | 14 | pub async fn client_connection(ws: WebSocket, id: String, clients: Clients, mut client: Client) { 15 | let (client_ws_sender, mut client_ws_rcv) = ws.split(); 16 | let (client_sender, client_rcv) = mpsc::unbounded_channel(); 17 | 18 | let client_rcv = UnboundedReceiverStream::new(client_rcv); 19 | tokio::task::spawn(client_rcv.forward(client_ws_sender).map(|result| { 20 | if let Err(e) = result { 21 | eprintln!("error sending websocket msg: {}", e); 22 | } 23 | })); 24 | 25 | client.sender = Some(client_sender); 26 | clients.write().await.insert(id.clone(), client); 27 | 28 | println!("{} connected", id); 29 | 30 | while let Some(result) = client_ws_rcv.next().await { 31 | let msg = match result { 32 | Ok(msg) => msg, 33 | Err(e) => { 34 | eprintln!("error receiving ws message for id: {}): {}", id.clone(), e); 35 | break; 36 | } 37 | }; 38 | client_msg(&id, msg, &clients).await; 39 | } 40 | 41 | clients.write().await.remove(&id); 42 | println!("{} disconnected", id); 43 | } 44 | 45 | async fn client_msg(id: &str, msg: Message, clients: &Clients) { 46 | println!("received message from {}: {:?}", id, msg); 47 | let message = match msg.to_str() { 48 | Ok(v) => v, 49 | Err(_) => return, 50 | }; 51 | 52 | if message == "ping" || message == "ping\n" { 53 | return; 54 | } 55 | 56 | let topics_req: TopicsRequest = match from_str(&message) { 57 | Ok(v) => v, 58 | Err(e) => { 59 | eprintln!("error while parsing message to topics request: {}", e); 60 | return; 61 | } 62 | }; 63 | 64 | let mut locked = clients.write().await; 65 | if let Some(v) = locked.get_mut(id) { 66 | v.topics = topics_req.topics; 67 | } 68 | } 69 | -------------------------------------------------------------------------------- /src/main.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | use std::convert::Infallible; 3 | use std::sync::Arc; 4 | use handler::TopicActionRequest; 5 | use tokio::sync::{mpsc, RwLock}; 6 | use warp::{ws::Message, Filter, Rejection}; 7 | use crate::handler::{add_topic, remove_topic}; 8 | 9 | mod handler; 10 | mod ws; 11 | 12 | type Result = std::result::Result; 13 | type Clients = Arc>>; 14 | 15 | #[derive(Debug, Clone)] 16 | pub struct Client { 17 | pub user_id: usize, 18 | pub topics: Vec, 19 | pub sender: Option>>, 20 | } 21 | 22 | #[tokio::main] 23 | async fn main() { 24 | let clients: Clients = Arc::new(RwLock::new(HashMap::new())); 25 | 26 | let health_route = warp::path!("health").and_then(handler::health_handler); 27 | 28 | let register = warp::path("register"); 29 | let register_routes = register 30 | .and(warp::post()) 31 | .and(warp::body::json()) 32 | .and(with_clients(clients.clone())) 33 | .and_then(handler::register_handler) 34 | .or(register 35 | .and(warp::delete()) 36 | .and(warp::path::param()) 37 | .and(with_clients(clients.clone())) 38 | .and_then(handler::unregister_handler)); 39 | 40 | let publish = warp::path!("publish") 41 | .and(warp::body::json()) 42 | .and(with_clients(clients.clone())) 43 | .and_then(handler::publish_handler); 44 | 45 | let ws_route = warp::path("ws") 46 | .and(warp::ws()) 47 | .and(warp::path::param()) 48 | .and(with_clients(clients.clone())) 49 | .and_then(handler::ws_handler); 50 | 51 | let clients_for_add = clients.clone(); 52 | let add_topic_route = warp::post() 53 | .and(warp::path("add_topic")) 54 | .and(warp::body::json::()) 55 | .and(warp::any().map(move || clients_for_add.clone())) 56 | .and_then(add_topic); 57 | 58 | let clients_for_remove = clients.clone(); 59 | let remove_topic_route = warp::delete() 60 | .and(warp::path("remove_topic")) 61 | .and(warp::body::json::()) 62 | .and(warp::any().map(move || clients_for_remove.clone())) 63 | .and_then(remove_topic); 64 | 65 | let routes = health_route 66 | .or(register_routes) 67 | .or(ws_route) 68 | .or(publish) 69 | .or(add_topic_route) 70 | .or(remove_topic_route) 71 | .with(warp::cors().allow_any_origin()); 72 | 73 | 74 | 75 | warp::serve(routes).run(([127, 0, 0, 1], 8000)).await; 76 | } 77 | 78 | fn with_clients(clients: Clients) -> impl Filter + Clone { 79 | warp::any().map(move || clients.clone()) 80 | } 81 | -------------------------------------------------------------------------------- /src/handler.rs: -------------------------------------------------------------------------------- 1 | use crate::{ws, Client, Clients, Result}; 2 | use serde::{Deserialize, Serialize}; 3 | use uuid::Uuid; 4 | use warp::{http::StatusCode, reply::json, ws::Message, Reply}; 5 | #[derive(Deserialize, Debug)] 6 | pub struct RegisterRequest { 7 | user_id: usize, 8 | topic: String, 9 | } 10 | 11 | #[derive(Deserialize)] 12 | pub struct TopicActionRequest { 13 | topic: String, 14 | client_id: String, 15 | } 16 | 17 | 18 | #[derive(Serialize, Debug)] 19 | pub struct RegisterResponse { 20 | url: String, 21 | } 22 | 23 | #[derive(Deserialize, Debug)] 24 | pub struct Event { 25 | topic: String, 26 | user_id: Option, 27 | message: String, 28 | } 29 | 30 | pub async fn publish_handler(body: Event, clients: Clients) -> Result { 31 | clients 32 | .read() 33 | .await 34 | .iter() 35 | .filter(|(_, client)| match body.user_id { 36 | Some(v) => client.user_id == v, 37 | None => true, 38 | }) 39 | .filter(|(_, client)| client.topics.contains(&body.topic)) 40 | .for_each(|(_, client)| { 41 | if let Some(sender) = &client.sender { 42 | let _ = sender.send(Ok(Message::text(body.message.clone()))); 43 | } 44 | }); 45 | 46 | Ok(StatusCode::OK) 47 | } 48 | 49 | pub async fn register_handler(body: RegisterRequest, clients: Clients) -> Result { 50 | let user_id = body.user_id; 51 | let topic = body.topic; // Capture the entry topic 52 | let uuid = Uuid::new_v4().as_simple().to_string(); 53 | 54 | register_client(uuid.clone(), user_id, topic, clients).await; // Pass the entry topic 55 | Ok(json(&RegisterResponse { 56 | url: format!("ws://127.0.0.1:8000/ws/{}", uuid), 57 | })) 58 | } 59 | 60 | async fn register_client(id: String, user_id: usize, topic: String, clients: Clients) { 61 | clients.write().await.insert( 62 | id, 63 | Client { 64 | user_id, 65 | topics: vec![topic], 66 | sender: None, 67 | }, 68 | ); 69 | } 70 | 71 | pub async fn unregister_handler(id: String, clients: Clients) -> Result { 72 | clients.write().await.remove(&id); 73 | Ok(StatusCode::OK) 74 | } 75 | 76 | pub async fn ws_handler(ws: warp::ws::Ws, id: String, clients: Clients) -> Result { 77 | let client = clients.read().await.get(&id).cloned(); 78 | match client { 79 | Some(c) => Ok(ws.on_upgrade(move |socket| ws::client_connection(socket, id, clients, c))), 80 | None => Err(warp::reject::not_found()), 81 | } 82 | } 83 | 84 | pub async fn health_handler() -> Result { 85 | Ok(StatusCode::OK) 86 | } 87 | 88 | 89 | 90 | pub async fn add_topic(body: TopicActionRequest, clients: Clients) -> Result { 91 | let mut clients_write = clients.write().await; 92 | if let Some(client) = clients_write.get_mut(&body.client_id) { 93 | client.topics.push(body.topic); 94 | } 95 | Ok(warp::reply::with_status("Added topic successfully", StatusCode::OK)) 96 | } 97 | 98 | pub async fn remove_topic(body: TopicActionRequest, clients: Clients) -> Result { 99 | let mut clients_write = clients.write().await; 100 | if let Some(client) = clients_write.get_mut(&body.client_id) { 101 | client.topics.retain(|t| t != &body.topic); 102 | } 103 | Ok(warp::reply::with_status("Removed topic successfully", StatusCode::OK)) 104 | } --------------------------------------------------------------------------------