1use std::sync::Arc;
2
3use n0_future::{SinkExt, StreamExt, split};
4use tokio::sync::broadcast;
5use tokio_tungstenite::tungstenite::Message;
6
7use crate::error::{Result, SyncwebError};
8
9use super::{encoding, service::BridgeService};
10
11const FRAME_HEADER_LEN: usize = 9;
12
13pub struct WsSession {
14 service: Arc<BridgeService>,
15 sender: split::SplitSink<tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>, Message>,
16 shutdown: broadcast::Receiver<()>,
17}
18
19impl WsSession {
20 #[must_use]
21 pub const fn new(
22 service: Arc<BridgeService>,
23 sender: split::SplitSink<tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>, Message>,
24 shutdown: broadcast::Receiver<()>,
25 ) -> Self {
26 Self {
27 service,
28 sender,
29 shutdown,
30 }
31 }
32
33 pub async fn run(
39 mut self,
40 receiver: &mut split::SplitStream<tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>>,
41 ) -> Result<()> {
42 self.send_node_status("connected").await;
43
44 loop {
45 tokio::select! {
46 _ = self.shutdown.recv() => {
47 self.send_node_status("disconnected").await;
48 break;
49 }
50 msg = receiver.next() => {
51 match msg {
52 Some(Ok(Message::Binary(data))) => {
53 if let Err(e) = self.handle_frame(data.to_vec()).await {
54 tracing::warn!("frame error (continuing): {e}");
55 }
56 }
57 Some(Ok(Message::Close(_))) => {
58 break;
59 }
60 Some(Ok(Message::Ping(payload))) => {
61 let _ = self.sender.send(Message::Pong(payload)).await;
62 }
63 Some(Err(e)) => {
64 tracing::warn!("websocket error: {e}");
65 break;
66 }
67 None => break,
68 _ => {}
69 }
70 }
71 }
72 }
73
74 Ok(())
75 }
76
77 async fn handle_frame(&mut self, bytes: Vec<u8>) -> Result<()> {
78 let header = Self::parse_frame_header(&bytes)?;
79 let payload = Self::extract_payload(&bytes, header.payload_len)?;
80
81 let result = match header.tag {
82 0x01 => Self::handle_dial_peer(payload),
83 0x02 => Box::pin(self.handle_append_event(payload)).await,
84 0x03 => Box::pin(self.handle_get_events(payload)).await,
85 0x04 => Box::pin(self.handle_get_events_paged(payload)).await,
86 0x05 => Box::pin(self.handle_share_collection(payload)).await,
87 0x06 => Box::pin(self.handle_import_collection(payload)).await,
88 0x07 => Box::pin(self.handle_join_gossip_topic(payload, false)).await,
89 0x08 => Box::pin(self.handle_join_gossip_topic(payload, true)).await,
90 0x09 => Box::pin(self.handle_leave_gossip_topic(payload)).await,
91 0x0A => Box::pin(self.handle_send_gossip_message(payload)).await,
92 0x0B => Box::pin(self.handle_get_connected_peers(payload)).await,
93 0x0C => self.handle_block_peer(payload),
94 0x0D => self.handle_unblock_peer(payload),
95 0x0E => Ok(self.handle_get_blocked_peers(payload)),
96 0x0F => Ok(self.handle_get_node_id(payload)),
97 _ => {
98 let msg = format!("unknown tag: 0x{:02X}", header.tag);
99 self.send_error(header.seq, &msg).await;
100 return Err(SyncwebError::operation("unknown tag", format!("0x{:02X}", header.tag)));
101 }
102 };
103
104 match result {
105 Ok(response_bytes) => {
106 self.send_ok(header.seq, &response_bytes).await?;
107 }
108 Err(e) => {
109 self.send_error(header.seq, &e.to_string()).await;
110 }
111 }
112
113 Ok(())
114 }
115
116 fn parse_frame_header(bytes: &[u8]) -> Result<FrameHeader> {
117 if bytes.len() < FRAME_HEADER_LEN {
118 return Err(SyncwebError::operation("frame parse error", "frame too short"));
119 }
120 let tag = *bytes
121 .first()
122 .ok_or_else(|| SyncwebError::operation("frame parse error", "empty frame"))?;
123 let seq_arr: [u8; 4] = bytes
124 .get(1..5)
125 .ok_or_else(|| SyncwebError::operation("frame parse error", "missing seq"))?
126 .try_into()
127 .map_err(|error: std::array::TryFromSliceError| SyncwebError::operation("frame parse error", error))?;
128 let len_arr: [u8; 4] = bytes
129 .get(5..9)
130 .ok_or_else(|| SyncwebError::operation("frame parse error", "missing payload_len"))?
131 .try_into()
132 .map_err(|error: std::array::TryFromSliceError| SyncwebError::operation("frame parse error", error))?;
133 Ok(FrameHeader {
134 tag,
135 seq: u32::from_be_bytes(seq_arr),
136 payload_len: u32::from_be_bytes(len_arr),
137 })
138 }
139
140 fn extract_payload(bytes: &[u8], payload_len: u32) -> Result<&[u8]> {
141 let len = usize::try_from(payload_len).unwrap_or(usize::MAX);
142 let end = FRAME_HEADER_LEN.saturating_add(len);
143 bytes
144 .get(FRAME_HEADER_LEN..end)
145 .ok_or_else(|| SyncwebError::operation("frame parse error", "truncated payload"))
146 }
147
148 fn handle_dial_peer(payload: &[u8]) -> Result<Vec<u8>> {
149 let mut offset = 0;
150 let node_id = encoding::read_string(payload, &mut offset)?;
151 let _: iroh::PublicKey = node_id
152 .parse()
153 .map_err(|e| SyncwebError::operation("invalid peer id", e))?;
154 Ok(Vec::new())
155 }
156
157 async fn handle_append_event(&self, payload: &[u8]) -> Result<Vec<u8>> {
158 let mut offset = 0;
159 let collection_id = encoding::read_string(payload, &mut offset)?;
160 let event_payload = encoding::read_bytes(payload, &mut offset)?;
161
162 let namespace = self.service.get_or_create_collection(&collection_id).await?;
163 let doc_opt = self.service.node().docs_engine().open(namespace).await?;
164 let doc = doc_opt
165 .ok_or_else(|| SyncwebError::operation("collection namespace not available", namespace.to_string()))?;
166 let author = self.service.node().docs_engine().author().await?;
167
168 let key = format!("evt/{}", uuid::Uuid::new_v4());
169 self.service
170 .node()
171 .docs_engine()
172 .set(&doc, author, &key, &event_payload)
173 .await?;
174
175 let node_id = self.service.node().endpoint().id().to_string();
176 let mut response = Vec::new();
177 encoding::write_string(&mut response, &node_id);
178 encoding::write_u64(&mut response, 0);
179 Ok(response)
180 }
181
182 async fn handle_get_events(&self, payload: &[u8]) -> Result<Vec<u8>> {
183 let mut offset = 0;
184 let collection_id = encoding::read_string(payload, &mut offset)?;
185 let entries = self.list_entries(&collection_id).await?;
186
187 let node_id = self.service.node().endpoint().id().to_string();
188 let mut response = Vec::new();
189 let count = u32::try_from(entries.len()).unwrap_or(0);
190 encoding::write_u32(&mut response, count);
191 for content in entries {
192 encoding::write_string(&mut response, &node_id);
193 encoding::write_bytes(&mut response, &content);
194 encoding::write_u64(&mut response, 0);
195 }
196 Ok(response)
197 }
198
199 async fn handle_get_events_paged(&self, payload: &[u8]) -> Result<Vec<u8>> {
200 let mut offset = 0;
201 let collection_id = encoding::read_string(payload, &mut offset)?;
202 let page_raw = encoding::read_u64(payload, &mut offset)?;
203 let page_size_raw = encoding::read_u64(payload, &mut offset)?;
204 let entries = self.list_entries(&collection_id).await?;
205
206 let page_idx = usize::try_from(page_raw).unwrap_or(0);
207 let limit = usize::try_from(page_size_raw).unwrap_or(0);
208 let start = page_idx.saturating_sub(1).saturating_mul(limit);
209 let paged: Vec<Vec<u8>> = entries.into_iter().skip(start).take(limit).collect();
210
211 let node_id = self.service.node().endpoint().id().to_string();
212 let mut response = Vec::new();
213 let count = u32::try_from(paged.len()).unwrap_or(0);
214 encoding::write_u32(&mut response, count);
215 for content in paged {
216 encoding::write_string(&mut response, &node_id);
217 encoding::write_bytes(&mut response, &content);
218 encoding::write_u64(&mut response, 0);
219 }
220 Ok(response)
221 }
222
223 async fn list_entries(&self, collection_id: &str) -> Result<Vec<Vec<u8>>> {
224 let namespace = self.service.get_collection(collection_id).await?;
225 let doc_opt = self.service.node().docs_engine().open(namespace).await?;
226 let doc = doc_opt.ok_or_else(|| SyncwebError::operation("collection not found", collection_id.to_owned()))?;
227 let entries = self.service.node().docs_engine().list_latest(&doc).await?;
228 let mut result = Vec::with_capacity(entries.len());
229 for entry in entries {
230 let content = self.service.node().blob_store().get(entry.content_hash()).await?;
231 result.push(content.to_vec());
232 }
233 Ok(result)
234 }
235
236 async fn handle_share_collection(&self, payload: &[u8]) -> Result<Vec<u8>> {
237 let mut offset = 0;
238 let collection_id = encoding::read_string(payload, &mut offset)?;
239
240 let namespace = self.service.get_collection(&collection_id).await?;
241 let doc_opt = self.service.node().docs_engine().open(namespace).await?;
242 let doc = doc_opt.ok_or_else(|| SyncwebError::operation("collection not found", collection_id))?;
243 let ticket = self
244 .service
245 .node()
246 .docs_engine()
247 .share(
248 &doc,
249 iroh_docs::api::protocol::ShareMode::Read,
250 self.service.node().endpoint().addr(),
251 )
252 .await?;
253
254 let mut response = Vec::new();
255 encoding::write_string(&mut response, &ticket.to_string());
256 Ok(response)
257 }
258
259 async fn handle_import_collection(&self, payload: &[u8]) -> Result<Vec<u8>> {
260 let mut offset = 0;
261 let collection_id = encoding::read_string(payload, &mut offset)?;
262 let ticket_str = encoding::read_string(payload, &mut offset)?;
263
264 let ticket: iroh_docs::DocTicket = ticket_str
265 .parse()
266 .map_err(|e| SyncwebError::operation("invalid doc ticket", e))?;
267 let doc = self.service.node().docs_engine().import_ticket(ticket).await?;
268 let namespace = self.service.node().docs_engine().namespace_id(&doc);
269 self.service.node().docs_engine().start_sync(&doc, Vec::new()).await?;
270
271 self.service.insert_collection(&collection_id, namespace).await;
272
273 Ok(Vec::new())
274 }
275
276 async fn handle_join_gossip_topic(&self, payload: &[u8], with_discovery: bool) -> Result<Vec<u8>> {
277 let mut offset = 0;
278 let topic_str = encoding::read_string(payload, &mut offset)?;
279
280 let _sender = self.service.subscribe_gossip(&topic_str, with_discovery).await?;
281 Ok(Vec::new())
282 }
283
284 async fn handle_leave_gossip_topic(&self, payload: &[u8]) -> Result<Vec<u8>> {
285 let mut offset = 0;
286 let topic_str = encoding::read_string(payload, &mut offset)?;
287
288 self.service.leave_gossip(&topic_str).await;
289 Ok(Vec::new())
290 }
291
292 async fn handle_send_gossip_message(&self, payload: &[u8]) -> Result<Vec<u8>> {
293 let mut offset = 0;
294 let topic_str = encoding::read_string(payload, &mut offset)?;
295 let message = encoding::read_bytes(payload, &mut offset)?;
296
297 self.service.send_gossip(&topic_str, &message).await?;
298 Ok(Vec::new())
299 }
300
301 async fn handle_get_connected_peers(&self, _payload: &[u8]) -> Result<Vec<u8>> {
302 let peers = self
303 .service
304 .connected_peers
305 .read()
306 .await
307 .values()
308 .cloned()
309 .collect::<Vec<_>>();
310 let mut response = Vec::new();
311 encoding::write_peer_list(&mut response, &peers);
312 Ok(response)
313 }
314
315 fn handle_block_peer(&self, payload: &[u8]) -> Result<Vec<u8>> {
316 let mut offset = 0;
317 let node_id = encoding::read_string(payload, &mut offset)?;
318 self.service.block_peer(&node_id)?;
319 Ok(Vec::new())
320 }
321
322 fn handle_unblock_peer(&self, payload: &[u8]) -> Result<Vec<u8>> {
323 let mut offset = 0;
324 let node_id = encoding::read_string(payload, &mut offset)?;
325 self.service.unblock_peer(&node_id)?;
326 Ok(Vec::new())
327 }
328
329 fn handle_get_blocked_peers(&self, _payload: &[u8]) -> Vec<u8> {
330 let blocked = self.service.get_blocked_peers();
331 let mut response = Vec::new();
332 encoding::write_string_list(&mut response, &blocked);
333 response
334 }
335
336 fn handle_get_node_id(&self, _payload: &[u8]) -> Vec<u8> {
337 let node_id = self.service.node().endpoint().id().to_string();
338 let mut response = Vec::new();
339 encoding::write_string(&mut response, &node_id);
340 response
341 }
342
343 async fn send_ok(&mut self, seq: u32, payload: &[u8]) -> Result<()> {
344 let frame = build_response(0x80, seq, payload);
345 self.sender
346 .send(Message::Binary(frame.into()))
347 .await
348 .map_err(|e| SyncwebError::operation("websocket send failed", e))
349 }
350
351 async fn send_error(&mut self, seq: u32, message: &str) {
352 let mut payload = Vec::new();
353 encoding::write_string(&mut payload, message);
354 let frame = build_response(0x81, seq, &payload);
355 if let Err(e) = self.sender.send(Message::Binary(frame.into())).await {
356 tracing::warn!("failed to send error frame: {e}");
357 }
358 }
359
360 async fn send_node_status(&mut self, status: &str) {
361 let node_id = self.service.node().endpoint().id().to_string();
362 let mut payload = Vec::new();
363 encoding::write_string(&mut payload, status);
364 encoding::write_string(&mut payload, &node_id);
365 let frame = build_response(0x83, 0, &payload);
366 let _ = self.sender.send(Message::Binary(frame.into())).await;
367 }
368}
369
370struct FrameHeader {
371 tag: u8,
372 seq: u32,
373 payload_len: u32,
374}
375
376fn build_response(tag: u8, seq: u32, payload: &[u8]) -> Vec<u8> {
377 let payload_len = u32::try_from(payload.len()).unwrap_or(0);
378 let len = FRAME_HEADER_LEN
379 .checked_add(usize::try_from(payload_len).unwrap_or(0))
380 .unwrap_or(FRAME_HEADER_LEN);
381 let mut frame = Vec::with_capacity(len);
382 frame.push(tag);
383 frame.extend_from_slice(&seq.to_be_bytes());
384 frame.extend_from_slice(&payload_len.to_be_bytes());
385 frame.extend_from_slice(payload);
386 frame
387}