├── Cargo.toml ├── README.md ├── tests ├── common.rs ├── test_client.rs └── test_server.rs ├── client.rs ├── LICENSE └── src └── server.rs /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "russmp" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | [dependencies] 7 | tokio = { version = "1", features = ["full"] } 8 | tokio-util = "0.7" # codec 9 | bytes = "1.4" 10 | futures = "0.3" 11 | clap = { version = "4", features = ["derive"] } 12 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # rusmmp 2 | SSMP implementation in Rust 3 | 4 | # How to use 5 | 6 | start server at 127.0.0.1:8080 7 | 8 | ```cargo run -- --server``` 9 | 10 | start client to connect server 127.0.0.1:8080 11 | 12 | ```cargo run -- --client``` 13 | 14 | # Use client only mode 15 | 16 | rust client.rs 17 | 18 | 19 | # run tests 20 | 21 | ```cargo test``` 22 | 23 | -------------------------------------------------------------------------------- /tests/common.rs: -------------------------------------------------------------------------------- 1 | use std::process::{Child, Command}; 2 | use std::time::Duration; 3 | 4 | pub struct ServerHandle { 5 | proc: Child, 6 | } 7 | 8 | impl ServerHandle { 9 | /// start server at random port 10 | pub fn start() -> (Self, u16) { 11 | let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); 12 | let port = listener.local_addr().unwrap().port(); 13 | drop(listener); // release port to reuse server 14 | 15 | let proc = Command::new("cargo") 16 | .args(&["run", "--", "--server", "--port", &port.to_string()]) 17 | .spawn() 18 | .expect("failed to start server"); 19 | 20 | // wait to restart server 21 | std::thread::sleep(Duration::from_millis(300)); 22 | 23 | (Self { proc }, port) 24 | } 25 | } 26 | 27 | impl Drop for ServerHandle { 28 | fn drop(&mut self) { 29 | let _ = self.proc.kill(); 30 | let _ = self.proc.wait(); 31 | } 32 | } 33 | -------------------------------------------------------------------------------- /tests/test_client.rs: -------------------------------------------------------------------------------- 1 | mod common; 2 | use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; 3 | use tokio::net::TcpStream; 4 | use tokio_util::codec::{Framed, LinesCodec}; 5 | 6 | #[tokio::test] 7 | async fn full_flow() { 8 | let (_srv, port) = common::ServerHandle::start(); 9 | let stream = TcpStream::connect(format!("127.0.0.1:{}", port)) 10 | .await 11 | .unwrap(); 12 | let mut framed = Framed::new(stream, LinesCodec::new()); 13 | 14 | // LOGIN 15 | framed.send("LOGIN alice open").await.unwrap(); 16 | let resp = framed.next().await.unwrap().unwrap(); 17 | assert_eq!(resp, "200"); 18 | 19 | // SUBSCRIBE 20 | framed.send("SUBSCRIBE news PRESENCE").await.unwrap(); 21 | assert_eq!(framed.next().await.unwrap().unwrap(), "200"); 22 | 23 | // MCAST 24 | framed.send("MCAST news hello").await.unwrap(); 25 | let event = framed.next().await.unwrap().unwrap(); 26 | assert!(event.starts_with("000 alice MCAST news hello")); 27 | 28 | // DUPLICATE SUBSCRIBE 29 | framed.send("SUBSCRIBE news").await.unwrap(); 30 | assert_eq!(framed.next().await.unwrap().unwrap(), "409"); 31 | 32 | // CLOSE 33 | framed.send("CLOSE").await.unwrap(); 34 | assert_eq!(framed.next().await.unwrap().unwrap(), "200"); 35 | } 36 | 37 | -------------------------------------------------------------------------------- /client.rs: -------------------------------------------------------------------------------- 1 | use tokio::{io::{AsyncBufReadExt, AsyncWriteExt, BufReader}, net::TcpStream}; 2 | use std::io::{self}; 3 | 4 | #[tokio::main] 5 | async fn main() -> io::Result<()> { 6 | let mut stream = TcpStream::connect("127.0.0.1:8080").await?; 7 | let (r, mut w) = stream.split(); 8 | let mut reader = BufReader::new(r); 9 | let mut line = String::new(); 10 | 11 | // LOGIN 12 | w.write_all(b"LOGIN client1 example.org\r\n").await?; 13 | reader.read_line(&mut line).await?; 14 | println!("[RECV] {}", line.trim_end()); 15 | line.clear(); 16 | 17 | // SUBSCRIBE to topic1 18 | w.write_all(b"SUBSCRIBE topic1\r\n").await?; 19 | 20 | // MCAST to topic1 21 | w.write_all(b"MCAST topic1 Hello, topic1!\r\n").await?; 22 | 23 | // BCAST 24 | w.write_all(b"BCAST Hello, all!\r\n").await?; 25 | 26 | // PING 27 | w.write_all(b"PING\r\n").await?; 28 | reader.read_line(&mut line).await?; 29 | println!("[RECV] {}", line.trim_end()); 30 | line.clear(); 31 | 32 | // Loop for incoming server messages (000 lines) 33 | loop { 34 | let bytes = reader.read_line(&mut line).await?; 35 | if bytes == 0 { 36 | println!("[DISCONNECTED]"); 37 | break; 38 | } 39 | println!("[SERVER] {}", line.trim_end()); 40 | line.clear(); 41 | } 42 | 43 | Ok(()) 44 | } 45 | -------------------------------------------------------------------------------- /tests/test_server.rs: -------------------------------------------------------------------------------- 1 | mod common; 2 | use ssmp_rs::{Payload, ServerState}; 3 | use tokio::sync::mpsc; 4 | 5 | #[tokio::test] 6 | async fn login_and_close() { 7 | let state = ServerState::default(); 8 | let (tx, _rx) = mpsc::unbounded_channel(); 9 | 10 | let peer = state.add_peer("user1".into(), tx).await; 11 | assert_eq!(state.peers.read().await.len(), 1); 12 | 13 | state.remove_peer("user1").await; 14 | assert!(state.peers.read().await.is_empty()); 15 | } 16 | 17 | #[tokio::test] 18 | async fn subscribe_multicast() { 19 | let state = ServerState::default(); 20 | let (tx, mut rx) = mpsc::unbounded_channel(); 21 | 22 | let peer = state.add_peer("bob".into(), tx).await; 23 | state.subscribe("news".into(), "bob", true).await; 24 | 25 | // bob should receive it's presence event 26 | let msg = rx.recv().await.unwrap(); 27 | assert!(msg.contains("SUBSCRIBE news PRESENCE")); 28 | 29 | // subscribe again, return 409（interupt when test client later下） 30 | } 31 | 32 | #[tokio::test] 33 | async fn broadcast_payload() { 34 | let state = ServerState::default(); 35 | let (tx1, mut rx1) = mpsc::unbounded_channel(); 36 | let (tx2, _rx2) = mpsc::unbounded_channel(); 37 | 38 | state.add_peer("a".into(), tx1).await; 39 | state.add_peer("b".into(), tx2).await; 40 | 41 | state.subscribe("room".into(), "a", false).await; 42 | state.subscribe("room".into(), "b", false).await; 43 | 44 | state.send_broadcast("a", Payload::Text("hi".into())).await; 45 | 46 | let msg = rx1.recv().await.unwrap(); 47 | assert!(msg.contains("BCAST")); 48 | } 49 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /src/server.rs: -------------------------------------------------------------------------------- 1 | // SSMP – Stupid-Simple Messaging Protocol – Reference Implementation (Rust) 2 | // cargo run -- --server -> listen 127.0.0.1:8080 3 | // cargo run -- --client -> interaction REPL connect to server 4 | 5 | use bytes::Bytes; 6 | use clap::Parser; 7 | use futures::{SinkExt, StreamExt}; 8 | use std::collections::{HashMap, HashSet}; 9 | use std::sync::Arc; 10 | use tokio::net::{TcpListener, TcpStream}; 11 | use tokio::sync::{mpsc, Mutex, RwLock}; 12 | use tokio_util::codec::{Framed, LinesCodec}; 13 | 14 | // ---------- protocol final via ---------- 15 | 16 | const MAX_LINE_LEN: usize = 4096; 17 | const PING_INTERVAL: std::time::Duration = std::time::Duration::from_secs(30); 18 | 19 | // ---------- message defination ---------- 20 | 21 | #[derive(Debug)] 22 | enum Msg { 23 | // client -> server 24 | Login { id: String, scheme: String, cred: Option }, 25 | Close, 26 | Ping, 27 | Pong, 28 | Subscribe { topic: String, presence: bool }, 29 | Unsubscribe { topic: String }, 30 | Ucast { to: String, payload: Payload }, 31 | Mcast { topic: String, payload: Payload }, 32 | Bcast { payload: Payload }, 33 | Unknown(String), // support forward compatibility 34 | } 35 | 36 | #[derive(Debug, Clone)] 37 | enum Payload { 38 | Text(String), 39 | Binary(Bytes), 40 | } 41 | 42 | impl Payload { 43 | fn encode(&self) -> String { 44 | match self { 45 | Payload::Text(s) => s.clone(), 46 | Payload::Binary(b) => { 47 | // binary encode: lead length 2 bytes + content 48 | let mut out = Vec::with_capacity(2 + b.len()); 49 | let len = (b.len() - 1) as u16; // protocol: length-1 50 | out.extend_from_slice(&len.to_be_bytes()); 51 | out.extend_from_slice(b); 52 | // convert to invisible charactors, base64 to use LinesCodec 53 | base64::encode(&out) 54 | } 55 | } 56 | } 57 | 58 | fn decode(raw: &str) -> Result { 59 | if raw.is_empty() { return Ok(Payload::Text("".into())); } 60 | let first = raw.as_bytes()[0]; 61 | if matches!(first, 0 | 1 | 2 | 3) { 62 | // binary 63 | let bytes = base64::decode(raw).map_err(|_| ())?; 64 | if bytes.len() < 2 { return Err(()); } 65 | let len = u16::from_be_bytes([bytes[0], bytes[1]]) as usize + 1; 66 | if len != bytes.len() - 2 { return Err(()); } 67 | Ok(Payload::Binary(Bytes::copy_from_slice(&bytes[2..]))) 68 | } else { 69 | Ok(Payload::Text(raw.into())) 70 | } 71 | } 72 | } 73 | 74 | // ---------- decode ---------- 75 | 76 | fn parse_msg(line: &str) -> Result { 77 | let mut toks = line.splitn(2, ' '); 78 | let verb = toks.next().ok_or(())?; 79 | let rest = toks.next().unwrap_or(""); 80 | 81 | match verb { 82 | "LOGIN" => { 83 | let mut t = rest.splitn(3, ' '); 84 | let id = t.next().ok_or(())?.to_string(); 85 | let scheme = t.next().ok_or(())?.to_string(); 86 | let cred = t.next().map(|s| s.to_string()); 87 | Ok(Msg::Login { id, scheme, cred }) 88 | } 89 | "CLOSE" => Ok(Msg::Close), 90 | "PING" => Ok(Msg::Ping), 91 | "PONG" => Ok(Msg::Pong), 92 | "SUBSCRIBE" => { 93 | let mut t = rest.splitn(2, ' '); 94 | let topic = t.next().ok_or(())?.to_string(); 95 | let presence = t.next() == Some("PRESENCE"); 96 | Ok(Msg::Subscribe { topic, presence }) 97 | } 98 | "UNSUBSCRIBE" => { 99 | let topic = rest.to_string(); 100 | if topic.is_empty() { return Err(()); } 101 | Ok(Msg::Unsubscribe { topic }) 102 | } 103 | "UCAST" => { 104 | let mut t = rest.splitn(2, ' '); 105 | let to = t.next().ok_or(())?.to_string(); 106 | let payload = Payload::decode(t.next().unwrap_or(""))?; 107 | Ok(Msg::Ucast { to, payload }) 108 | } 109 | "MCAST" => { 110 | let mut t = rest.splitn(2, ' '); 111 | let topic = t.next().ok_or(())?.to_string(); 112 | let payload = Payload::decode(t.next().unwrap_or(""))?; 113 | Ok(Msg::Mcast { topic, payload }) 114 | } 115 | "BCAST" => { 116 | let payload = Payload::decode(rest)?; 117 | Ok(Msg::Bcast { payload }) 118 | } 119 | _ => Ok(Msg::Unknown(line.to_string())), 120 | } 121 | } 122 | 123 | // ---------- server ---------- 124 | 125 | type Tx = mpsc::UnboundedSender; 126 | 127 | #[derive(Default)] 128 | struct ServerState { 129 | peers: RwLock>, 130 | topics: RwLock>>, // topic -> set 131 | } 132 | 133 | #[derive(Clone)] 134 | struct Peer { 135 | tx: Tx, 136 | topics: Arc>>, // this peer subscribed topic 137 | presence: Arc>>, // this peer require presence topic 138 | } 139 | 140 | impl ServerState { 141 | async fn add_peer(&self, id: String, tx: Tx) -> Peer { 142 | let peer = Peer { 143 | tx, 144 | topics: Default::default(), 145 | presence: Default::default(), 146 | }; 147 | self.peers.write().await.insert(id.clone(), peer.clone()); 148 | peer 149 | } 150 | 151 | async fn remove_peer(&self, id: &str) { 152 | let mut peers = self.peers.write().await; 153 | if let Some(peer) = peers.remove(id) { 154 | let topics = { peer.topics.lock().await.clone() }; 155 | for topic in topics { 156 | self.unsubscribe(topic, id).await; 157 | } 158 | } 159 | } 160 | 161 | async fn subscribe(&self, topic: String, peer_id: &str, presence: bool) { 162 | { 163 | let mut tmap = self.topics.write().await; 164 | tmap.entry(topic.clone()).or_default().insert(peer_id.to_string()); 165 | } 166 | if let Some(peer) = self.peers.read().await.get(peer_id) { 167 | peer.topics.lock().await.insert(topic.clone()); 168 | if presence { 169 | peer.presence.lock().await.insert(topic.clone()); 170 | } 171 | // send SUBSCRIBE event to other presence clients 172 | let event = format!("000 {} SUBSCRIBE {} {}", peer_id, topic, if presence { "PRESENCE" } else { "" }); 173 | self.broadcast_event(&event, &topic).await; 174 | } 175 | } 176 | 177 | async fn unsubscribe(&self, topic: String, peer_id: &str) { 178 | { 179 | let mut tmap = self.topics.write().await; 180 | if let Some(set) = tmap.get_mut(&topic) { 181 | set.remove(peer_id); 182 | if set.is_empty() { 183 | tmap.remove(&topic); 184 | } 185 | } 186 | } 187 | if let Some(peer) = self.peers.read().await.get(peer_id) { 188 | peer.topics.lock().await.remove(&topic); 189 | peer.presence.lock().await.remove(&topic); 190 | let event = format!("000 {} UNSUBSCRIBE {}", peer_id, topic); 191 | self.broadcast_event(&event, &topic).await; 192 | } 193 | } 194 | 195 | async fn broadcast_event(&self, event: &str, skip_topic: &str) { 196 | let topics = self.topics.read().await; 197 | let peers = self.peers.read().await; 198 | if let Some(subscribers) = topics.get(skip_topic) { 199 | for id in subscribers { 200 | if let Some(peer) = peers.get(id) { 201 | let _ = peer.tx.send(event.to_string()); 202 | } 203 | } 204 | } 205 | } 206 | 207 | async fn send_unicast(&self, to: &str, event: String) -> bool { 208 | if let Some(peer) = self.peers.read().await.get(to) { 209 | let _ = peer.tx.send(event); 210 | true 211 | } else { 212 | false 213 | } 214 | } 215 | 216 | async fn send_multicast(&self, topic: &str, from: &str, payload: Payload) { 217 | let topics = self.topics.read().await; 218 | let peers = self.peers.read().await; 219 | if let Some(subscribers) = topics.get(topic) { 220 | let event = format!("000 {} MCAST {} {}", from, topic, payload.encode()); 221 | for id in subscribers { 222 | if id != from { 223 | if let Some(peer) = peers.get(id) { 224 | let _ = peer.tx.send(event.clone()); 225 | } 226 | } 227 | } 228 | } 229 | } 230 | 231 | async fn send_broadcast(&self, from: &str, payload: Payload) { 232 | let peers = self.peers.read().await; 233 | let mut sent = HashSet::new(); 234 | // find from subscribed topic 235 | if let Some(from_peer) = peers.get(from) { 236 | let from_topics = from_peer.topics.lock().await.clone(); 237 | for topic in from_topics { 238 | if let Some(subscribers) = self.topics.read().await.get(&topic) { 239 | for id in subscribers { 240 | if id != from && !sent.contains(id) { 241 | if let Some(peer) = peers.get(id) { 242 | let event = format!("000 {} BCAST {}", from, payload.encode()); 243 | let _ = peer.tx.send(event); 244 | sent.insert(id.clone()); 245 | } 246 | } 247 | } 248 | } 249 | } 250 | } 251 | } 252 | } 253 | 254 | async fn handle_client(stream: TcpStream, state: Arc) { 255 | let mut framed = Framed::new(stream, LinesCodec::new_with_max_length(MAX_LINE_LEN)); 256 | let (tx, mut rx) = mpsc::unbounded_channel::(); 257 | 258 | // required LOGIN 259 | let login_line = match framed.next().await { 260 | Some(Ok(l)) => l, 261 | _ => return, 262 | }; 263 | let (peer_id, peer) = match parse_msg(&login_line) { 264 | Ok(Msg::Login { id, scheme: _, cred: _ }) => { 265 | // simply ignore scheme/cred 266 | let p = state.add_peer(id.clone(), tx).await; 267 | let _ = framed.send("200").await; 268 | (id, p) 269 | } 270 | _ => { 271 | let _ = framed.send("400").await; 272 | return; 273 | } 274 | }; 275 | 276 | // read async & heartbeat 277 | let (mut sink, mut stream) = framed.split(); 278 | let state_clone = state.clone(); 279 | let peer_id_clone = peer_id.clone(); 280 | 281 | // writeback task 282 | tokio::spawn(async move { 283 | while let Some(line) = rx.recv().await { 284 | if sink.send(line).await.is_err() { 285 | break; 286 | } 287 | } 288 | }); 289 | 290 | // heartbeat 291 | let ping_task = { 292 | let tx = peer.tx.clone(); 293 | tokio::spawn(async move { 294 | loop { 295 | tokio::time::sleep(PING_INTERVAL).await; 296 | if tx.send("000 . PING".into()).is_err() { 297 | break; 298 | } 299 | } 300 | }) 301 | }; 302 | 303 | // main loop 304 | while let Some(Ok(line)) = stream.next().await { 305 | match parse_msg(&line) { 306 | Ok(Msg::Close) => { 307 | let _ = peer.tx.send("200".into()); 308 | break; 309 | } 310 | Ok(Msg::Ping) => { 311 | let _ = peer.tx.send("000 . PONG".into()); 312 | } 313 | Ok(Msg::Pong) => { /* ignore */ } 314 | Ok(Msg::Subscribe { topic, presence }) => { 315 | state.subscribe(topic, &peer_id, presence).await; 316 | let _ = peer.tx.send("200".into()); 317 | } 318 | Ok(Msg::Unsubscribe { topic }) => { 319 | state.unsubscribe(topic, &peer_id).await; 320 | let _ = peer.tx.send("200".into()); 321 | } 322 | Ok(Msg::Ucast { to, payload }) => { 323 | let event = format!("000 {} UCAST {} {}", peer_id, to, payload.encode()); 324 | if !state.send_unicast(&to, event).await { 325 | let _ = peer.tx.send("404".into()); 326 | } 327 | } 328 | Ok(Msg::Mcast { topic, payload }) => { 329 | state.send_multicast(&topic, &peer_id, payload).await; 330 | } 331 | Ok(Msg::Bcast { payload }) => { 332 | state.send_broadcast(&peer_id, payload).await; 333 | } 334 | Ok(Msg::Unknown(line)) => { 335 | let _ = peer.tx.send("501".into()); 336 | } 337 | Err(_) => { 338 | let _ = peer.tx.send("400".into()); 339 | break; 340 | } 341 | } 342 | } 343 | 344 | ping_task.abort(); 345 | state.remove_peer(&peer_id_clone).await; 346 | } 347 | 348 | #[tokio::main] 349 | async fn run_server(port: u16) { 350 | 351 | let state = Arc::new(ServerState::default()); 352 | let listener = TcpListener::bind(format!("127.0.0.1:{}", port)) 353 | .await 354 | .unwrap(); 355 | println!("SSMP server listening on 127.0.0.1:{}", port); 356 | 357 | while let Ok((stream, _)) = listener.accept().await { 358 | tokio::spawn(handle_client(stream, state.clone())); 359 | } 360 | } 361 | 362 | // ---------- client REPL ---------- 363 | 364 | async fn client_repl() { 365 | let mut framed = Framed::new( 366 | TcpStream::connect("127.0.0.1:8080").await.unwrap(), 367 | LinesCodec::new_with_max_length(MAX_LINE_LEN), 368 | ); 369 | 370 | println!("Login identifier:"); 371 | let mut id = String::new(); 372 | std::io::stdin().read_line(&mut id).unwrap(); 373 | let id = id.trim(); 374 | framed.send(format!("LOGIN {} open")).await.unwrap(); 375 | 376 | let resp = framed.next().await.unwrap().unwrap(); 377 | if resp != "200" { 378 | println!("Login failed: {}", resp); 379 | return; 380 | } 381 | println!("Logged in as {}", id); 382 | 383 | // read/write 384 | let (mut sink, mut stream) = framed.split(); 385 | let (tx, mut rx) = mpsc::unbounded_channel::(); 386 | 387 | tokio::spawn(async move { 388 | while let Some(line) = rx.recv().await { 389 | let _ = sink.send(line).await; 390 | } 391 | }); 392 | 393 | // print events 394 | tokio::spawn(async move { 395 | while let Some(Ok(line)) = stream.next().await { 396 | println!("< {}", line); 397 | } 398 | }); 399 | 400 | // REPL 401 | loop { 402 | let mut line = String::new(); 403 | std::io::stdin().read_line(&mut line).unwrap(); 404 | let line = line.trim(); 405 | if line == "quit" { break; } 406 | tx.send(line.into()).unwrap(); 407 | } 408 | } 409 | 410 | // ---------- main ---------- 411 | 412 | #[derive(Parser)] 413 | #[command(version, about)] 414 | struct Args { 415 | #[arg(long)] 416 | server: bool, 417 | #[arg(long)] 418 | client: bool, 419 | #[arg(long)] 420 | port: Option, 421 | } 422 | 423 | #[tokio::main] 424 | async fn main() { 425 | let args = Args::parse(); 426 | if args.server { 427 | run_server().await; 428 | } else if args.client { 429 | client_repl().await; 430 | } else { 431 | println!("Use --server or --client"); 432 | } 433 | } --------------------------------------------------------------------------------