surreal_client/engines/
ws_cbor.rs

1//! WebSocket CBOR Engine for SurrealDB
2//!
3//! Connects with the `cbor` WebSocket subprotocol and exchanges CBOR-encoded
4//! binary frames. Preserves native types (datetime, duration, recordid, bytes)
5//! that JSON cannot carry. Accepts `ws://`, `wss://`, or `cbor://` URLs;
6//! `cbor://` is treated as `ws://`.
7
8use async_trait::async_trait;
9use ciborium::Value as CborValue;
10use futures_util::stream::{SplitSink, SplitStream};
11use futures_util::{SinkExt, StreamExt};
12use std::collections::HashMap;
13use std::sync::Arc;
14use std::sync::atomic::{AtomicU64, Ordering::SeqCst};
15use tokio::net::TcpStream;
16use tokio::sync::{Mutex, mpsc, oneshot};
17use tokio_tungstenite::MaybeTlsStream;
18use tokio_tungstenite::tungstenite::client::IntoClientRequest;
19use tokio_tungstenite::tungstenite::http::HeaderValue;
20use tokio_tungstenite::tungstenite::http::header::SEC_WEBSOCKET_PROTOCOL;
21use tokio_tungstenite::{WebSocketStream, connect_async, tungstenite::Message};
22use tracing::{Instrument as _, warn};
23
24use crate::SurrealConnection;
25use crate::live::{Action, Notification};
26use crate::{
27    engine::Engine,
28    error::{Result, SurrealError},
29};
30
31/// Live-query subscribers, keyed by the server's live-query id.
32type LiveSubscribers = Arc<Mutex<HashMap<String, mpsc::UnboundedSender<Notification>>>>;
33
34/// Normalise a CBOR id to a stable string key.
35///
36/// Request ids arrive as numeric text; live-query ids arrive as CBOR UUIDs
37/// (tag 37 wrapping a 16-byte string). Both the `live` RPC result and the
38/// unsolicited notification frames encode the query id the same way, so
39/// normalising both through here makes them compare equal.
40fn cbor_id_key(v: &CborValue) -> Option<String> {
41    match v {
42        CborValue::Text(t) => Some(t.clone()),
43        CborValue::Integer(i) => {
44            let n: i128 = (*i).into();
45            Some(n.to_string())
46        }
47        CborValue::Tag(_, inner) => cbor_id_key(inner),
48        CborValue::Bytes(b) if b.len() == 16 => Some(
49            uuid::Uuid::from_slice(b)
50                .map(|u| u.to_string())
51                .unwrap_or_else(|_| hex::encode(b)),
52        ),
53        CborValue::Bytes(b) => Some(hex::encode(b)),
54        _ => None,
55    }
56}
57
58/// Pull a [`Notification`] out of a live-query frame's inner `result` map.
59///
60/// The map SurrealDB delivers is `{ action, id, record, result, session }`:
61/// `action` is `CREATE`/`UPDATE`/`DELETE`, `id` is the live-query uuid, and
62/// the nested `result` is the affected record. Returns `None` if it doesn't
63/// look like a notification (so ordinary map-shaped responses fall through).
64fn parse_notification(inner: &[(CborValue, CborValue)]) -> Option<Notification> {
65    let mut action = None;
66    let mut query_id = None;
67    let mut record_id = None;
68    let mut data = None;
69    for (k, v) in inner {
70        if let CborValue::Text(k) = k {
71            match k.as_str() {
72                "action" => {
73                    if let CborValue::Text(a) = v {
74                        action = Action::parse(a);
75                    }
76                }
77                "id" => query_id = cbor_id_key(v),
78                "record" => record_id = Some(v.clone()),
79                "result" => data = Some(v.clone()),
80                _ => {}
81            }
82        }
83    }
84    Some(Notification {
85        query_id: query_id?,
86        action: action?,
87        record_id: record_id.unwrap_or(CborValue::Null),
88        data: data.unwrap_or(CborValue::Null),
89    })
90}
91
92/// Request structure for CBOR WebSocket protocol
93#[derive(Debug, Clone)]
94struct RouterRequest {
95    id: String,
96    method: String,
97    params: Option<CborValue>,
98}
99
100type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
101
102/// WebSocket engine using native CBOR for SurrealDB.
103///
104/// Connects with `Sec-WebSocket-Protocol: cbor` and sends/receives binary
105/// CBOR frames. Authentication and `use ns/db` are performed during
106/// `from_connection` so the returned engine is ready to issue queries.
107pub struct WsCborEngine {
108    sink: Arc<Mutex<SplitSink<WsStream, Message>>>,
109    stream: Arc<Mutex<SplitStream<WsStream>>>,
110    msg_id: AtomicU64,
111    pending_requests: Arc<Mutex<HashMap<String, oneshot::Sender<CborValue>>>>,
112    live_subscribers: LiveSubscribers,
113    task_handle: Option<tokio::task::JoinHandle<()>>,
114}
115
116impl WsCborEngine {
117    pub async fn from_connection(connect: &SurrealConnection) -> Result<Self> {
118        // rustls 0.23 needs a process-wide crypto provider before any `wss://`
119        // handshake. Install ring once; a competing prior install is fine.
120        static TLS_PROVIDER: std::sync::Once = std::sync::Once::new();
121        TLS_PROVIDER.call_once(|| {
122            let _ = rustls::crypto::ring::default_provider().install_default();
123        });
124
125        let base_url = connect
126            .url
127            .as_ref()
128            .ok_or_else(|| SurrealError::Connection("URL is required to connect".to_string()))?;
129
130        let mut ws_url = if let Some(rest) = base_url.strip_prefix("cbor://") {
131            format!("ws://{}", rest)
132        } else {
133            base_url.clone()
134        };
135        if !ws_url.ends_with("/rpc") {
136            if ws_url.ends_with('/') {
137                ws_url.push_str("rpc");
138            } else {
139                ws_url.push_str("/rpc");
140            }
141        }
142
143        let mut request = ws_url
144            .as_str()
145            .into_client_request()
146            .map_err(|e| SurrealError::Connection(format!("Invalid WebSocket URL: {}", e)))?;
147
148        request
149            .headers_mut()
150            .insert(SEC_WEBSOCKET_PROTOCOL, HeaderValue::from_static("cbor"));
151
152        let (stream, _response) = connect_async(request).await.map_err(|e| {
153            SurrealError::Connection(format!("Failed to connect to WebSocket: {}", e))
154        })?;
155
156        let (sink, stream) = stream.split();
157
158        let mut engine = Self {
159            sink: Arc::new(Mutex::new(sink)),
160            stream: Arc::new(Mutex::new(stream)),
161            msg_id: AtomicU64::new(0),
162            pending_requests: Arc::new(Mutex::new(HashMap::new())),
163            live_subscribers: Arc::new(Mutex::new(HashMap::new())),
164            task_handle: None,
165        };
166
167        let task_handle = engine.handle_messages();
168        engine.task_handle = Some(task_handle);
169
170        connect.init_engine(&mut engine).await?;
171
172        Ok(engine)
173    }
174
175    fn handle_messages(&self) -> tokio::task::JoinHandle<()> {
176        let stream = Arc::clone(&self.stream);
177        let pending_requests = Arc::clone(&self.pending_requests);
178        let live_subscribers = Arc::clone(&self.live_subscribers);
179
180        tokio::spawn(
181            async move {
182                loop {
183                    let msg = {
184                        let mut stream_guard = stream.lock().await;
185                        stream_guard.next().await
186                    };
187
188                    let msg = match msg {
189                        None => break,
190                        Some(Err(e)) => {
191                            warn!(error = %e, "CBOR ws receive error");
192                            break;
193                        }
194                        Some(Ok(msg)) => msg,
195                    };
196
197                    match msg {
198                        Message::Text(_text) => {
199                            // Ignore text messages - we only use CBOR binary
200                        }
201                        Message::Binary(binary) => {
202                            // Two frame shapes share this channel:
203                            //   response:     {id, result|error}         (id matches a request)
204                            //   notification: {id, action, result}        (id is a live-query uuid)
205                            match ciborium::from_reader::<CborValue, _>(binary.as_ref()) {
206                                Ok(CborValue::Map(map)) => {
207                                    let mut top_id = None;
208                                    let mut result = None;
209                                    let mut error = None;
210
211                                    for (key, value) in &map {
212                                        if let CborValue::Text(k) = key {
213                                            match k.as_str() {
214                                                "id" => top_id = Some(value.clone()),
215                                                "result" => result = Some(value.clone()),
216                                                "error" => error = Some(value.clone()),
217                                                _ => {}
218                                            }
219                                        }
220                                    }
221
222                                    // Live-query notifications carry no top-level id;
223                                    // the live-query id and action live inside `result`:
224                                    //   { result: { action, id: <uuid>, result: <record> } }
225                                    if top_id.is_none()
226                                        && let Some(CborValue::Map(inner)) = &result
227                                        && let Some(n) = parse_notification(inner)
228                                    {
229                                        let tx = {
230                                            let subs = live_subscribers.lock().await;
231                                            subs.get(&n.query_id).cloned()
232                                        };
233                                        if let Some(tx) = tx {
234                                            let _ = tx.send(n);
235                                        }
236                                        continue;
237                                    }
238
239                                    // Otherwise it is a reply — route it to the waiter.
240                                    if let Some(id) = top_id.as_ref().and_then(cbor_id_key) {
241                                        let tx = {
242                                            let mut pending = pending_requests.lock().await;
243                                            pending.remove(&id)
244                                        };
245
246                                        if let Some(tx) = tx {
247                                            if let Some(err) = error {
248                                                let _ = tx.send(CborValue::Map(vec![(
249                                                    CborValue::Text("error".to_string()),
250                                                    err,
251                                                )]));
252                                            } else if let Some(res) = result {
253                                                let _ = tx.send(res);
254                                            } else {
255                                                let _ = tx.send(CborValue::Null);
256                                            }
257                                        }
258                                    }
259                                }
260                                Ok(_) => {}
261                                Err(e) => {
262                                    warn!(error = %e, bytes = binary.len(), "CBOR parse failed");
263                                }
264                            }
265                        }
266                        Message::Ping(_) => {}
267                        Message::Pong(_) => {}
268                        Message::Close(_) => break,
269                        _ => {}
270                    }
271                }
272            }
273            .in_current_span(),
274        )
275    }
276}
277
278#[async_trait]
279impl Engine for WsCborEngine {
280    async fn send_message_cbor(&mut self, method: &str, params: CborValue) -> Result<CborValue> {
281        let (tx, rx) = oneshot::channel();
282        let id = self.msg_id.fetch_add(1, SeqCst).to_string();
283
284        {
285            let mut pending = self.pending_requests.lock().await;
286            pending.insert(id.clone(), tx);
287        }
288
289        let request = RouterRequest {
290            id: id.clone(),
291            method: method.to_string(),
292            params: Some(params),
293        };
294
295        let mut request_map = vec![
296            (
297                CborValue::Text("id".to_string()),
298                CborValue::Text(request.id),
299            ),
300            (
301                CborValue::Text("method".to_string()),
302                CborValue::Text(request.method),
303            ),
304        ];
305
306        if let Some(params) = request.params {
307            request_map.push((CborValue::Text("params".to_string()), params));
308        }
309
310        let rpc_message = CborValue::Map(request_map);
311
312        let mut payload = Vec::new();
313        ciborium::into_writer(&rpc_message, &mut payload)
314            .map_err(|e| SurrealError::Protocol(format!("CBOR encoding failed: {}", e)))?;
315
316        {
317            let mut sink = self.sink.lock().await;
318            sink.send(Message::Binary(payload.into()))
319                .await
320                .map_err(|e| SurrealError::Connection(format!("WS send failed: {}", e)))?;
321        }
322
323        let response = rx
324            .await
325            .map_err(|_| SurrealError::Protocol("Response channel closed".to_string()))?;
326
327        if let CborValue::Map(map) = &response {
328            for (key, value) in map {
329                if let CborValue::Text(k) = key
330                    && k == "error"
331                {
332                    if let CborValue::Map(error_map) = value {
333                        let mut code = -1;
334                        let mut message = String::new();
335
336                        for (error_key, error_value) in error_map {
337                            if let CborValue::Text(error_k) = error_key {
338                                match error_k.as_str() {
339                                    "code" => {
340                                        if let CborValue::Integer(c) = error_value {
341                                            code = (*c).try_into().unwrap_or(-1);
342                                        }
343                                    }
344                                    "message" => {
345                                        if let CborValue::Text(m) = error_value {
346                                            message = m.clone();
347                                        }
348                                    }
349                                    _ => {}
350                                }
351                            }
352                        }
353
354                        if !message.is_empty() {
355                            return Err(SurrealError::ServerError { code, message });
356                        }
357                    }
358
359                    return Err(SurrealError::Protocol(format!("Server error: {:?}", value)));
360                }
361            }
362        }
363
364        Ok(response)
365    }
366
367    async fn register_live(
368        &mut self,
369        query_id: &str,
370    ) -> Result<mpsc::UnboundedReceiver<Notification>> {
371        let (tx, rx) = mpsc::unbounded_channel();
372        let mut subs = self.live_subscribers.lock().await;
373        subs.insert(query_id.to_string(), tx);
374        Ok(rx)
375    }
376
377    async fn unregister_live(&mut self, query_id: &str) {
378        let mut subs = self.live_subscribers.lock().await;
379        subs.remove(query_id);
380    }
381}
382
383impl Drop for WsCborEngine {
384    fn drop(&mut self) {
385        if let Some(handle) = self.task_handle.take() {
386            handle.abort();
387        }
388    }
389}
surreal_client/engines/ws_cbor.rs

surreal_client/engines/
ws_cbor.rs
