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

wscall_server/
server_runtime.rs
