Skip to main content

wscall_server/
server_runtime.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::panic::AssertUnwindSafe;
4use std::sync::Arc;
5use std::sync::atomic::Ordering;
6use std::time::Duration;
7
8use bytes::Bytes;
9use futures_util::{FutureExt, SinkExt, StreamExt, future::BoxFuture};
10use serde::de::DeserializeOwned;
11use serde_json::{Map, Value, json};
12use tokio::net::{TcpListener, TcpStream};
13use tokio::sync::{Semaphore, mpsc};
14use tokio::time::{MissedTickBehavior, interval, timeout};
15use tokio_tungstenite::{
16    WebSocketStream, accept_async_with_config, tungstenite::Message,
17    tungstenite::protocol::WebSocketConfig,
18};
19use uuid::Uuid;
20use validator::Validate;
21use wscall_protocol::{
22    DEFAULT_MAX_FRAME_BYTES, EcdhKeypair, EncryptionKind, ErrorPayload, FileAttachment, FrameCodec,
23    PacketBody, PacketEnvelope, ProtocolError, parse_peer_public,
24};
25
26use crate::server_types::{
27    ApiContext, ApiError, ClientEntry, EventContext, ExceptionContext, ServerConnectionContext,
28    ServerDisconnectContext, ServerError, ServerHandle, ServerOutbound, ServerState,
29};
30
31const SERVER_IDLE_TIMEOUT: Duration = Duration::from_secs(45);
32const SERVER_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15);
33const SERVER_OUTBOUND_QUEUE_CAPACITY: usize = 256;
34/// Default cap on concurrently in-flight request/event handlers per connection.
35const SERVER_DEFAULT_MAX_IN_FLIGHT: usize = 64;
36/// Grace window given to the writer task to flush a close frame before abort.
37const SERVER_CLOSE_GRACE: Duration = Duration::from_millis(500);
38
39type ApiHandler =
40    Arc<dyn Fn(ApiContext) -> BoxFuture<'static, Result<Value, ApiError>> + Send + Sync>;
41type Filter =
42    Arc<dyn Fn(ApiContext) -> BoxFuture<'static, Result<ApiContext, ApiError>> + Send + Sync>;
43type EventHandler =
44    Arc<dyn Fn(EventContext) -> BoxFuture<'static, Result<Value, ApiError>> + Send + Sync>;
45type ConnectionHandler =
46    Arc<dyn Fn(ServerConnectionContext) -> BoxFuture<'static, ()> + Send + Sync>;
47type DisconnectHandler =
48    Arc<dyn Fn(ServerDisconnectContext) -> BoxFuture<'static, ()> + Send + Sync>;
49type ExceptionHandler =
50    Arc<dyn Fn(ExceptionContext) -> BoxFuture<'static, ErrorPayload> + Send + Sync>;
51
52struct ApiRequestInput {
53    request_id: u64,
54    route: String,
55    params: Value,
56    attachments: Vec<FileAttachment>,
57    metadata: Value,
58}
59
60struct EventEmitInput {
61    event_id: u64,
62    name: String,
63    data: Map<String, Value>,
64    attachments: Vec<FileAttachment>,
65    metadata: Value,
66}
67
68impl ServerHandle {
69    /// Broadcasts an event to every live connection.
70    ///
71    /// The frame is encoded exactly once and shared as [`Bytes`] across all
72    /// recipients, so the cost no longer scales with the connection count times
73    /// the payload size. A full outbound queue for a single connection drops
74    /// that one delivery instead of failing the whole broadcast.
75    pub async fn broadcast_event(
76        &self,
77        name: impl Into<String>,
78        data: Map<String, Value>,
79        attachments: Vec<FileAttachment>,
80    ) -> Result<(), ApiError> {
81        let event_id = self.state.next_event_id.fetch_add(1, Ordering::Relaxed) + 1;
82        let packet = PacketEnvelope::with_encryption(
83            PacketBody::EventEmit {
84                event_id,
85                name: name.into(),
86                data,
87                attachments,
88                metadata: json!({ "source": "server" }),
89                expect_ack: true,
90            },
91            self.default_encryption,
92        );
93
94        if self.is_ecdh {
95            // ECDH mode: each connection has a unique session key, so the
96            // writer task must encode per-connection.
97            for entry in self.state.clients.iter() {
98                if entry
99                    .value()
100                    .sender
101                    .try_send(ServerOutbound::Packet(packet.clone()))
102                    .is_err()
103                {
104                    tracing::warn!(
105                        "broadcast: outbound queue full, dropping event for a connection"
106                    );
107                }
108            }
109        } else {
110            // PSK mode: all connections share the same key, so encode once.
111            let encoded = self
112                .codec
113                .encode(&packet)
114                .map_err(|_| ApiError::internal("failed to encode broadcast event"))?;
115            let encoded = Bytes::from(encoded);
116
117            for entry in self.state.clients.iter() {
118                if entry
119                    .value()
120                    .sender
121                    .try_send(ServerOutbound::PreEncoded(encoded.clone()))
122                    .is_err()
123                {
124                    tracing::warn!(
125                        "broadcast: outbound queue full, dropping event for a connection"
126                    );
127                }
128            }
129        }
130        Ok(())
131    }
132
133    /// Sends an event to a single connection by id.
134    ///
135    /// The frame is pre-encoded once so the writer task only ships bytes.
136    pub async fn send_event_to(
137        &self,
138        connection_id: &str,
139        name: impl Into<String>,
140        data: Map<String, Value>,
141        attachments: Vec<FileAttachment>,
142    ) -> Result<(), ApiError> {
143        let event_id = self.state.next_event_id.fetch_add(1, Ordering::Relaxed) + 1;
144        let packet = PacketEnvelope::with_encryption(
145            PacketBody::EventEmit {
146                event_id,
147                name: name.into(),
148                data,
149                attachments,
150                metadata: json!({ "source": "server" }),
151                expect_ack: true,
152            },
153            self.default_encryption,
154        );
155
156        let entry = self
157            .state
158            .clients
159            .get(connection_id)
160            .ok_or_else(|| ApiError::not_found("target connection not found"))?;
161
162        if self.is_ecdh {
163            // ECDH mode: send unencoded packet; the writer encodes it.
164            entry
165                .sender
166                .try_send(ServerOutbound::Packet(packet))
167                .map_err(|_| ApiError::internal("failed to queue direct event"))
168        } else {
169            let encoded = self
170                .codec
171                .encode(&packet)
172                .map_err(|_| ApiError::internal("failed to encode direct event"))?;
173            let encoded = Bytes::from(encoded);
174            entry
175                .sender
176                .try_send(ServerOutbound::PreEncoded(encoded))
177                .map_err(|_| ApiError::internal("failed to queue direct event"))
178        }
179    }
180
181    /// Returns the current number of live connections.
182    pub async fn connection_count(&self) -> usize {
183        self.state.clients.len()
184    }
185}
186
187pub struct WscallServer {
188    state: Arc<ServerState>,
189    routes: HashMap<String, ApiHandler>,
190    filters: Vec<Filter>,
191    event_handlers: HashMap<String, EventHandler>,
192    connection_handlers: Vec<ConnectionHandler>,
193    disconnect_handlers: Vec<DisconnectHandler>,
194    exception_handler: Option<ExceptionHandler>,
195    codec: FrameCodec,
196    default_encryption: EncryptionKind,
197    /// Whether ECDH dynamic key agreement is enabled (no pre-shared key).
198    is_ecdh: bool,
199    /// Optional global cap on concurrent accepted connections.
200    max_connections: Option<usize>,
201    /// Per-connection cap on concurrently running request/event handlers.
202    max_in_flight: usize,
203    /// Maximum total frame size (including 4-byte length prefix).
204    max_frame_bytes: usize,
205}
206
207impl Default for WscallServer {
208    fn default() -> Self {
209        Self::new()
210    }
211}
212
213impl WscallServer {
214    pub fn new() -> Self {
215        Self {
216            state: Arc::new(ServerState::new()),
217            routes: HashMap::new(),
218            filters: Vec::new(),
219            event_handlers: HashMap::new(),
220            connection_handlers: Vec::new(),
221            disconnect_handlers: Vec::new(),
222            exception_handler: None,
223            codec: FrameCodec::plaintext(),
224            default_encryption: EncryptionKind::None,
225            is_ecdh: false,
226            max_connections: None,
227            max_in_flight: SERVER_DEFAULT_MAX_IN_FLIGHT,
228            max_frame_bytes: DEFAULT_MAX_FRAME_BYTES,
229        }
230    }
231
232    /// Enables ECDH dynamic key agreement.
233    ///
234    /// Instead of a pre-shared symmetric key, each connection performs an
235    /// X25519 handshake right after the WebSocket upgrade. The negotiated
236    /// 32-byte session key is used with ChaCha20-Poly1305 for all subsequent
237    /// frames. This is mutually exclusive with `with_chacha20_key` /
238    /// `with_aes256_key`.
239    pub fn with_ecdh(mut self) -> Self {
240        self.is_ecdh = true;
241        // The global codec stays plaintext; per-connection codecs carry the
242        // negotiated session keys.
243        self.default_encryption = EncryptionKind::ChaCha20;
244        self
245    }
246
247    pub fn with_chacha20_key(mut self, key: [u8; 32]) -> Self {
248        self.codec = self.codec.clone().with_chacha20_key(key);
249        self.default_encryption = EncryptionKind::ChaCha20;
250        self
251    }
252
253    pub fn with_aes256_key(mut self, key: [u8; 32]) -> Self {
254        self.codec = self.codec.clone().with_aes256_key(key);
255        self.default_encryption = EncryptionKind::Aes256;
256        self
257    }
258
259    /// Caps the number of concurrently accepted connections.
260    ///
261    /// When the cap is reached, new `accept` calls block until an existing
262    /// connection drops, providing natural backpressure against connection
263    /// flooding instead of spawning unbounded tasks.
264    pub fn with_max_connections(mut self, max: usize) -> Self {
265        self.max_connections = Some(max);
266        self
267    }
268
269    /// Caps the number of concurrently in-flight request/event handlers per
270    /// single connection, bounding per-connection CPU and memory usage.
271    pub fn with_max_in_flight(mut self, max: usize) -> Self {
272        self.max_in_flight = max;
273        self
274    }
275
276    /// Sets the maximum total frame size (including the 4-byte length prefix).
277    ///
278    /// Frames exceeding this limit cause the server to send a 413 error response
279    /// frame (`request_id=0`, `code="frame_too_large"`); the connection stays
280    /// open. Default: 100 MiB.
281    pub fn with_max_frame_bytes(mut self, max: usize) -> Self {
282        self.max_frame_bytes = max;
283        self
284    }
285
286    pub fn handle(&self) -> ServerHandle {
287        ServerHandle {
288            state: Arc::clone(&self.state),
289            codec: self.codec.clone(),
290            default_encryption: self.default_encryption,
291            is_ecdh: self.is_ecdh,
292        }
293    }
294
295    pub fn route<F, Fut>(&mut self, route: impl Into<String>, handler: F)
296    where
297        F: Fn(ApiContext) -> Fut + Send + Sync + 'static,
298        Fut: Future<Output = Result<Value, ApiError>> + Send + 'static,
299    {
300        let handler = Arc::new(move |ctx: ApiContext| {
301            Box::pin(handler(ctx)) as BoxFuture<'static, Result<Value, ApiError>>
302        });
303        self.routes.insert(route.into(), handler);
304    }
305
306    pub fn typed_route<T, F, Fut>(&mut self, route: impl Into<String>, handler: F)
307    where
308        T: DeserializeOwned + Send + 'static,
309        F: Fn(ApiContext, T) -> Fut + Send + Sync + 'static,
310        Fut: Future<Output = Result<Value, ApiError>> + Send + 'static,
311    {
312        let handler = Arc::new(handler);
313        self.route(route, move |ctx| {
314            let handler = Arc::clone(&handler);
315            let params = ctx.bind::<T>();
316            async move {
317                let params = params?;
318                handler(ctx, params).await
319            }
320        });
321    }
322
323    pub fn validated_route<T, F, Fut>(&mut self, route: impl Into<String>, handler: F)
324    where
325        T: DeserializeOwned + Validate + Send + 'static,
326        F: Fn(ApiContext, T) -> Fut + Send + Sync + 'static,
327        Fut: Future<Output = Result<Value, ApiError>> + Send + 'static,
328    {
329        let handler = Arc::new(handler);
330        self.route(route, move |ctx| {
331            let handler = Arc::clone(&handler);
332            let params = ctx.bind_validated::<T>();
333            async move {
334                let params = params?;
335                handler(ctx, params).await
336            }
337        });
338    }
339
340    pub fn filter<F, Fut>(&mut self, filter: F)
341    where
342        F: Fn(ApiContext) -> Fut + Send + Sync + 'static,
343        Fut: Future<Output = Result<ApiContext, ApiError>> + Send + 'static,
344    {
345        let filter = Arc::new(move |ctx: ApiContext| {
346            Box::pin(filter(ctx)) as BoxFuture<'static, Result<ApiContext, ApiError>>
347        });
348        self.filters.push(filter);
349    }
350
351    pub fn event_handler<F, Fut>(&mut self, name: impl Into<String>, handler: F)
352    where
353        F: Fn(EventContext) -> Fut + Send + Sync + 'static,
354        Fut: Future<Output = Result<Value, ApiError>> + Send + 'static,
355    {
356        let handler = Arc::new(move |ctx: EventContext| {
357            Box::pin(handler(ctx)) as BoxFuture<'static, Result<Value, ApiError>>
358        });
359        self.event_handlers.insert(name.into(), handler);
360    }
361
362    pub fn on_connected<F, Fut>(&mut self, handler: F)
363    where
364        F: Fn(ServerConnectionContext) -> Fut + Send + Sync + 'static,
365        Fut: Future<Output = ()> + Send + 'static,
366    {
367        let handler = Arc::new(move |ctx: ServerConnectionContext| {
368            Box::pin(handler(ctx)) as BoxFuture<'static, ()>
369        });
370        self.connection_handlers.push(handler);
371    }
372
373    pub fn on_disconnected<F, Fut>(&mut self, handler: F)
374    where
375        F: Fn(ServerDisconnectContext) -> Fut + Send + Sync + 'static,
376        Fut: Future<Output = ()> + Send + 'static,
377    {
378        let handler = Arc::new(move |ctx: ServerDisconnectContext| {
379            Box::pin(handler(ctx)) as BoxFuture<'static, ()>
380        });
381        self.disconnect_handlers.push(handler);
382    }
383
384    pub fn exception_handler<F, Fut>(&mut self, handler: F)
385    where
386        F: Fn(ExceptionContext) -> Fut + Send + Sync + 'static,
387        Fut: Future<Output = ErrorPayload> + Send + 'static,
388    {
389        self.exception_handler = Some(Arc::new(move |ctx: ExceptionContext| {
390            Box::pin(handler(ctx)) as BoxFuture<'static, ErrorPayload>
391        }));
392    }
393
394    pub async fn listen(self, address: &str) -> Result<(), ServerError> {
395        let listener = TcpListener::bind(address).await?;
396        tracing::info!(%address, "wscall server listening on ws://{address}/socket");
397
398        let shared = Arc::new(self);
399        let conn_sem = shared
400            .max_connections
401            .map(|max| Arc::new(Semaphore::new(max)));
402
403        loop {
404            // Acquire a connection permit (when capped) before accepting so a
405            // connection flood turns into backpressure rather than unbounded
406            // task spawning.
407            let permit = match &conn_sem {
408                Some(sem) => match sem.clone().acquire_owned().await {
409                    Ok(permit) => Some(permit),
410                    Err(_) => continue,
411                },
412                None => None,
413            };
414
415            let (stream, peer) = listener.accept().await?;
416            // Disable Nagle's algorithm for low-latency RPC frames.
417            let _ = TcpStream::set_nodelay(&stream, true);
418
419            let server = Arc::clone(&shared);
420            tokio::spawn(async move {
421                let _permit = permit;
422                if let Err(error) = server.serve_connection(stream, peer).await {
423                    tracing::warn!(?error, ?peer, "connection failed");
424                }
425            });
426        }
427    }
428
429    async fn serve_connection(
430        self: Arc<Self>,
431        stream: TcpStream,
432        peer: std::net::SocketAddr,
433    ) -> Result<(), ServerError> {
434        // Configure WebSocket-level message size limit slightly above the
435        // WSCALL limit so that oversized frames can still be received and
436        // answered with a 413 error response instead of an abrupt close.
437        let ws_config = WebSocketConfig {
438            max_message_size: Some(self.max_frame_bytes + 1024 * 1024),
439            max_frame_size: Some(self.max_frame_bytes + 1024 * 1024),
440            ..Default::default()
441        };
442        let mut websocket = accept_async_with_config(stream, Some(ws_config)).await?;
443        let connection_id = Uuid::now_v7().to_string();
444        let peer_addr = Some(peer);
445
446        // ECDH handshake: exchange X25519 public keys and derive a per-
447        // connection session key before splitting the stream.
448        let session_codec = if self.is_ecdh {
449            let (codec, _encryption) = self.perform_ecdh_handshake(&mut websocket).await?;
450            codec.with_max_frame_bytes(self.max_frame_bytes)
451        } else {
452            self.codec
453                .clone()
454                .with_max_frame_bytes(self.max_frame_bytes)
455        };
456
457        let (mut sink, mut stream) = websocket.split();
458        let (tx, mut rx) = mpsc::channel::<ServerOutbound>(SERVER_OUTBOUND_QUEUE_CAPACITY);
459
460        // Store the per-connection entry.
461        self.state
462            .clients
463            .insert(connection_id.clone(), ClientEntry { sender: tx.clone() });
464
465        self.notify_connected(&connection_id, peer_addr).await;
466
467        // The writer needs the per-connection codec to encode `Packet`
468        // variants (ECDH mode).
469        let writer_codec = session_codec.clone();
470        let mut writer = tokio::spawn(async move {
471            let mut ticker = interval(SERVER_HEARTBEAT_INTERVAL);
472            ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
473            loop {
474                tokio::select! {
475                    maybe_outbound = rx.recv() => {
476                        let Some(outbound) = maybe_outbound else {
477                            break Ok::<(), ServerError>(());
478                        };
479                        match outbound {
480                            ServerOutbound::PreEncoded(bytes) => {
481                                sink.send(Message::Binary(bytes.to_vec())).await?;
482                            }
483                            ServerOutbound::Packet(packet) => {
484                                let encoded = writer_codec.encode(&packet)?;
485                                sink.send(Message::Binary(encoded)).await?;
486                            }
487                            ServerOutbound::Pong(payload) => {
488                                sink.send(Message::Pong(payload)).await?;
489                            }
490                            ServerOutbound::Close => {
491                                let _ = sink.send(Message::Close(None)).await;
492                                break Ok(());
493                            }
494                        }
495                    }
496                    _ = ticker.tick() => {
497                        if let Err(error) = sink.send(Message::Ping(Vec::new())).await {
498                            break Err(ServerError::WebSocket(error));
499                        }
500                    }
501                }
502            }
503        });
504
505        // The reader uses the per-connection codec (session key in ECDH mode,
506        // global codec in PSK mode).
507        let reader_codec = session_codec.clone();
508
509        // Per-connection concurrency limit for handler execution. The reader
510        // stays non-blocking under normal load; under overload it pauses reading
511        // new frames, which is the desired backpressure signal.
512        let in_flight = Arc::new(Semaphore::new(self.max_in_flight));
513
514        let result = async {
515            self.handle()
516                .send_event_to(
517                    &connection_id,
518                    "system.notice",
519                    json!({ "message": "connected", "connection_id": connection_id })
520                        .as_object()
521                        .unwrap()
522                        .clone(),
523                    Vec::new(),
524                )
525                .await
526                .map_err(ServerError::Api)?;
527
528            loop {
529                let next_message = timeout(SERVER_IDLE_TIMEOUT, stream.next()).await;
530                let Some(message) =
531                    next_message.map_err(|_| ServerError::IdleTimeout(connection_id.clone()))?
532                else {
533                    break Ok(());
534                };
535
536                match message? {
537                    Message::Binary(bytes) => {
538                        // WSCALL-level frame size check: reject oversized frames
539                        // with a 413 error response; the connection stays open.
540                        if bytes.len() > self.max_frame_bytes {
541                            let error_packet = PacketEnvelope::with_encryption(
542                                PacketBody::ApiResponse {
543                                    request_id: 0,
544                                    ok: false,
545                                    status: 413,
546                                    data: json!({}),
547                                    error: Some(ErrorPayload {
548                                        code: "frame_too_large".to_string(),
549                                        message: format!(
550                                            "frame size {} exceeds limit {}",
551                                            bytes.len(),
552                                            self.max_frame_bytes
553                                        ),
554                                        status: 413,
555                                        details: None,
556                                    }),
557                                    metadata: json!({}),
558                                },
559                                self.default_encryption,
560                            );
561                            let _ = self.queue_for(&connection_id, error_packet).await;
562                            continue;
563                        }
564
565                        let packet = reader_codec.decode(&bytes)?;
566
567                        let spawn_handler = matches!(
568                            &packet.body,
569                            PacketBody::ApiRequest { .. } | PacketBody::EventEmit { .. }
570                        );
571
572                        if spawn_handler {
573                            let permit = in_flight.clone().acquire_owned().await.map_err(|_| {
574                                ServerError::Api(ApiError::internal("in-flight semaphore closed"))
575                            })?;
576                            let server = Arc::clone(&self);
577                            let conn_id = connection_id.clone();
578                            tokio::spawn(async move {
579                                let _permit = permit;
580                                if let Err(error) =
581                                    server.process_packet(&conn_id, peer_addr, packet).await
582                                {
583                                    tracing::warn!(%error, "packet processing failed");
584                                }
585                            });
586                        } else {
587                            // Trivial control messages (acks / responses from
588                            // clients) are handled inline without spawning.
589                            self.process_packet(&connection_id, peer_addr, packet)
590                                .await?;
591                        }
592                    }
593                    Message::Close(_) => break Ok(()),
594                    Message::Ping(payload) => {
595                        if tx
596                            .send(ServerOutbound::Pong(payload.to_vec()))
597                            .await
598                            .is_err()
599                        {
600                            break Ok(());
601                        }
602                    }
603                    Message::Pong(_) => {}
604                    Message::Text(_) => {}
605                    Message::Frame(_) => {}
606                }
607            }
608        }
609        .await;
610
611        // Teardown: remove from the live table, attempt a graceful close, then
612        // abort the writer. Aborting the writer drops the channel receiver so
613        // any in-flight handler `queue_for` calls fail fast instead of hanging.
614        self.state.clients.remove(&connection_id);
615        let _ = tx.try_send(ServerOutbound::Close);
616        let _ = timeout(SERVER_CLOSE_GRACE, &mut writer).await;
617        writer.abort();
618        self.notify_disconnected(&connection_id, peer_addr, Self::disconnect_reason(&result))
619            .await;
620        result
621    }
622
623    /// Performs the X25519 ECDH handshake over the combined WebSocket stream.
624    ///
625    /// 1. Generate a server keypair.
626    /// 2. Wait for the client's 32-byte public key (raw binary message).
627    /// 3. Send the server's 32-byte public key back.
628    /// 4. Derive the 32-byte ChaCha20-Poly1305 session key.
629    ///
630    /// Must be called before `websocket.split()` because it uses both
631    /// the read and write halves of the stream.
632    async fn perform_ecdh_handshake(
633        &self,
634        websocket: &mut WebSocketStream<TcpStream>,
635    ) -> Result<(FrameCodec, EncryptionKind), ServerError> {
636        // 1. Generate server keypair.
637        let keypair = EcdhKeypair::generate()?;
638
639        // 2. Read the client's public key (first binary WebSocket message).
640        let next = timeout(Duration::from_secs(10), websocket.next()).await;
641        let client_public = match next {
642            Ok(Some(Ok(Message::Binary(bytes)))) => parse_peer_public(&bytes).map_err(|_| {
643                ServerError::Protocol(ProtocolError::EcdhHandshake(
644                    "client sent invalid public key length".to_string(),
645                ))
646            })?,
647            _ => {
648                return Err(ProtocolError::EcdhHandshake(
649                    "client did not send a valid public key".to_string(),
650                )
651                .into());
652            }
653        };
654
655        // 3. Send the server's public key.
656        websocket
657            .send(Message::Binary(keypair.public_bytes().to_vec()))
658            .await?;
659
660        // 4. Derive the session key and build the per-connection codec.
661        let session_key = keypair.derive_session_key(&client_public);
662        Ok((
663            FrameCodec::plaintext().with_chacha20_key(session_key),
664            EncryptionKind::ChaCha20,
665        ))
666    }
667
668    async fn process_packet(
669        &self,
670        connection_id: &str,
671        peer_addr: Option<std::net::SocketAddr>,
672        packet: PacketEnvelope,
673    ) -> Result<(), ServerError> {
674        match packet.body {
675            PacketBody::ApiRequest {
676                request_id,
677                route,
678                params,
679                attachments,
680                metadata,
681            } => {
682                let response = self
683                    .run_api_request(
684                        connection_id,
685                        peer_addr,
686                        ApiRequestInput {
687                            request_id,
688                            route,
689                            params,
690                            attachments,
691                            metadata,
692                        },
693                    )
694                    .await;
695                self.queue_for(connection_id, response).await?;
696            }
697            PacketBody::EventEmit {
698                event_id,
699                name,
700                data,
701                attachments,
702                metadata,
703                ..
704            } => {
705                let ack = self
706                    .run_event(
707                        connection_id,
708                        peer_addr,
709                        EventEmitInput {
710                            event_id,
711                            name,
712                            data,
713                            attachments,
714                            metadata,
715                        },
716                    )
717                    .await;
718                self.queue_for(connection_id, ack).await?;
719            }
720            PacketBody::EventAck {
721                event_id,
722                ok,
723                receipt,
724                error,
725            } => {
726                tracing::debug!(
727                    connection_id,
728                    event_id,
729                    ok,
730                    ?receipt,
731                    ?error,
732                    "received event ack from client",
733                );
734            }
735            PacketBody::ApiResponse { .. } => {}
736        }
737        Ok(())
738    }
739
740    /// Queues a response packet for a connection.
741    ///
742    /// The frame is pre-encoded here (offloading the writer task and enabling
743    /// parallel encoding across concurrent handlers) and sent with backpressure
744    /// semantics: a slow consumer pauses the producer rather than being dropped.
745    async fn queue_for(
746        &self,
747        connection_id: &str,
748        packet: PacketEnvelope,
749    ) -> Result<(), ServerError> {
750        // Clone the sender out of the DashMap guard before awaiting so the
751        // shard lock is never held across an await point.
752        let sender = match self.state.clients.get(connection_id) {
753            Some(entry) => entry.sender.clone(),
754            None => {
755                return Err(ServerError::Api(ApiError::not_found(
756                    "connection is closed",
757                )));
758            }
759        };
760
761        if self.is_ecdh {
762            // ECDH mode: the writer task encodes with the per-connection codec.
763            sender
764                .send(ServerOutbound::Packet(packet))
765                .await
766                .map_err(|_| {
767                    ServerError::Api(ApiError::internal("failed to queue outbound packet"))
768                })
769        } else {
770            // PSK mode: pre-encode with the shared global codec.
771            let encoded = self.codec.encode(&packet)?;
772            let encoded = Bytes::from(encoded);
773            sender
774                .send(ServerOutbound::PreEncoded(encoded))
775                .await
776                .map_err(|_| {
777                    ServerError::Api(ApiError::internal("failed to queue outbound packet"))
778                })
779        }
780    }
781
782    async fn run_api_request(
783        &self,
784        connection_id: &str,
785        peer_addr: Option<std::net::SocketAddr>,
786        request: ApiRequestInput,
787    ) -> PacketEnvelope {
788        let ApiRequestInput {
789            request_id,
790            route,
791            params,
792            attachments,
793            metadata,
794        } = request;
795
796        let mut ctx = ApiContext {
797            connection_id: connection_id.to_string(),
798            peer_addr,
799            request_id,
800            route: route.clone(),
801            params,
802            attachments,
803            metadata,
804            server: self.handle(),
805        };
806
807        for filter in &self.filters {
808            match filter(ctx).await {
809                Ok(next_ctx) => ctx = next_ctx,
810                Err(error) => {
811                    return self
812                        .api_error_packet(connection_id, Some(request_id), route, error)
813                        .await;
814                }
815            }
816        }
817
818        let Some(handler) = self.routes.get(&ctx.route) else {
819            return self
820                .api_error_packet(
821                    connection_id,
822                    Some(request_id),
823                    route,
824                    ApiError::not_found("route not found"),
825                )
826                .await;
827        };
828
829        match AssertUnwindSafe(handler(ctx)).catch_unwind().await {
830            Ok(Ok(data)) => PacketEnvelope::with_encryption(
831                PacketBody::ApiResponse {
832                    request_id,
833                    ok: true,
834                    status: 200,
835                    data,
836                    error: None,
837                    metadata: json!({}),
838                },
839                self.default_encryption,
840            ),
841            Ok(Err(error)) => {
842                self.api_error_packet(connection_id, Some(request_id), route, error)
843                    .await
844            }
845            Err(_) => {
846                self.api_error_packet(
847                    connection_id,
848                    Some(request_id),
849                    route,
850                    ApiError::internal("handler panicked"),
851                )
852                .await
853            }
854        }
855    }
856
857    async fn run_event(
858        &self,
859        connection_id: &str,
860        peer_addr: Option<std::net::SocketAddr>,
861        event: EventEmitInput,
862    ) -> PacketEnvelope {
863        let EventEmitInput {
864            event_id,
865            name,
866            data,
867            attachments,
868            metadata,
869        } = event;
870
871        let ctx = EventContext {
872            connection_id: connection_id.to_string(),
873            peer_addr,
874            event_id,
875            name: name.clone(),
876            data,
877            attachments,
878            metadata,
879            server: self.handle(),
880        };
881
882        let Some(handler) = self.event_handlers.get(&name) else {
883            return PacketEnvelope::with_encryption(
884                PacketBody::EventAck {
885                    event_id,
886                    ok: false,
887                    receipt: json!({}),
888                    error: Some(ApiError::not_found("event handler not found").into_payload()),
889                },
890                self.default_encryption,
891            );
892        };
893
894        match AssertUnwindSafe(handler(ctx)).catch_unwind().await {
895            Ok(Ok(receipt)) => PacketEnvelope::with_encryption(
896                PacketBody::EventAck {
897                    event_id,
898                    ok: true,
899                    receipt,
900                    error: None,
901                },
902                self.default_encryption,
903            ),
904            Ok(Err(error)) => PacketEnvelope::with_encryption(
905                PacketBody::EventAck {
906                    event_id,
907                    ok: false,
908                    receipt: json!({}),
909                    error: Some(
910                        self.map_exception(ExceptionContext {
911                            connection_id: connection_id.to_string(),
912                            request_id: Some(event_id),
913                            target: name,
914                            message_kind: "event",
915                            error,
916                        })
917                        .await,
918                    ),
919                },
920                self.default_encryption,
921            ),
922            Err(_) => PacketEnvelope::with_encryption(
923                PacketBody::EventAck {
924                    event_id,
925                    ok: false,
926                    receipt: json!({}),
927                    error: Some(
928                        self.map_exception(ExceptionContext {
929                            connection_id: connection_id.to_string(),
930                            request_id: Some(event_id),
931                            target: name,
932                            message_kind: "event",
933                            error: ApiError::internal("event handler panicked"),
934                        })
935                        .await,
936                    ),
937                },
938                self.default_encryption,
939            ),
940        }
941    }
942
943    async fn api_error_packet(
944        &self,
945        connection_id: &str,
946        request_id: Option<u64>,
947        route: String,
948        error: ApiError,
949    ) -> PacketEnvelope {
950        let request_id = request_id.unwrap_or(0);
951        let status = error.status;
952        let payload = self
953            .map_exception(ExceptionContext {
954                connection_id: connection_id.to_string(),
955                request_id: Some(request_id),
956                target: route,
957                message_kind: "api",
958                error,
959            })
960            .await;
961
962        PacketEnvelope::with_encryption(
963            PacketBody::ApiResponse {
964                request_id,
965                ok: false,
966                status,
967                data: json!({}),
968                error: Some(payload),
969                metadata: json!({}),
970            },
971            self.default_encryption,
972        )
973    }
974
975    async fn notify_connected(&self, connection_id: &str, peer_addr: Option<std::net::SocketAddr>) {
976        let handlers = self.connection_handlers.clone();
977        for handler in handlers {
978            let context = ServerConnectionContext {
979                connection_id: connection_id.to_string(),
980                peer_addr,
981                server: self.handle(),
982            };
983
984            if AssertUnwindSafe(handler(context))
985                .catch_unwind()
986                .await
987                .is_err()
988            {
989                tracing::error!("server connected handler panicked");
990            }
991        }
992    }
993
994    async fn notify_disconnected(
995        &self,
996        connection_id: &str,
997        peer_addr: Option<std::net::SocketAddr>,
998        reason: String,
999    ) {
1000        let handlers = self.disconnect_handlers.clone();
1001        for handler in handlers {
1002            let context = ServerDisconnectContext {
1003                connection_id: connection_id.to_string(),
1004                peer_addr,
1005                reason: reason.clone(),
1006                server: self.handle(),
1007            };
1008
1009            if AssertUnwindSafe(handler(context))
1010                .catch_unwind()
1011                .await
1012                .is_err()
1013            {
1014                tracing::error!("server disconnected handler panicked");
1015            }
1016        }
1017    }
1018
1019    fn disconnect_reason(result: &Result<(), ServerError>) -> String {
1020        match result {
1021            Ok(()) => "connection closed".to_string(),
1022            Err(ServerError::IdleTimeout(_)) => "idle timeout".to_string(),
1023            Err(error) => error.to_string(),
1024        }
1025    }
1026
1027    async fn map_exception(&self, context: ExceptionContext) -> ErrorPayload {
1028        match &self.exception_handler {
1029            Some(handler) => handler(context).await,
1030            None => context.error.into_payload(),
1031        }
1032    }
1033}