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