a3s_boot/websocket/
gateway.rs

1use super::connection::WebSocketGatewayConnection;
2use super::context::WebSocketGatewayInitContext;
3use super::handler::WebSocketSubscriptionDefinition;
4use super::hooks::{
5    WebSocketGatewayConnectionHook, WebSocketGatewayDisconnectHook, WebSocketGatewayInitHook,
6};
7use super::message::{send_to_outbounds, IntoWebSocketReply, WebSocketMessage, WebSocketOutbound};
8use super::pipeline::{
9    prepend_execution_guards, prepend_execution_interceptors, ExecutionWebSocketGuard,
10    ExecutionWebSocketInterceptor, WebSocketGuard, WebSocketInterceptor, WebSocketPipe,
11};
12use super::server::WebSocketGatewayServer;
13use super::state::{normalize_namespace, normalize_room, WebSocketGatewayState};
14use crate::pipeline::{PipelineComponent, PipelineOverrides};
15use crate::routing::path::{
16    join_paths, match_path_params, match_path_shape, route_shape_key, validate_route_path,
17};
18use crate::{
19    catch_errors, BootError, BootErrorKind, BootRequest, ExecutionInterceptor, Guard, HttpMethod,
20    Result, ValidationOptions, WebSocketExceptionFilter,
21};
22use serde::Serialize;
23use serde_json::Value;
24use std::collections::BTreeMap;
25use std::future::Future;
26use std::sync::atomic::AtomicBool;
27use std::sync::Arc;
28
29/// Framework-neutral WebSocket gateway definition.
30#[derive(Clone)]
31pub struct WebSocketGatewayDefinition {
32    pub(crate) path: String,
33    pub(crate) namespace: Option<String>,
34    pub(crate) handlers: BTreeMap<String, WebSocketSubscriptionDefinition>,
35    pub(crate) init_hooks: Vec<Arc<dyn WebSocketGatewayInitHook>>,
36    pub(crate) connection_hooks: Vec<Arc<dyn WebSocketGatewayConnectionHook>>,
37    pub(crate) disconnect_hooks: Vec<Arc<dyn WebSocketGatewayDisconnectHook>>,
38    pub(crate) pipes: Vec<PipelineComponent<dyn WebSocketPipe>>,
39    pub(crate) guards: Vec<PipelineComponent<dyn WebSocketGuard>>,
40    pub(crate) interceptors: Vec<PipelineComponent<dyn WebSocketInterceptor>>,
41    pub(crate) filters: Vec<PipelineComponent<dyn WebSocketExceptionFilter>>,
42    pub(crate) metadata: BTreeMap<String, Value>,
43    pub(crate) module_name: Option<String>,
44    pub(crate) state: Arc<WebSocketGatewayState>,
45}
46
47impl WebSocketGatewayDefinition {
48    pub fn new(path: impl Into<String>) -> Result<Self> {
49        let path = path.into();
50        validate_route_path(&path)?;
51        Ok(Self {
52            path,
53            namespace: None,
54            handlers: BTreeMap::new(),
55            init_hooks: Vec::new(),
56            connection_hooks: Vec::new(),
57            disconnect_hooks: Vec::new(),
58            pipes: Vec::new(),
59            guards: Vec::new(),
60            interceptors: Vec::new(),
61            filters: Vec::new(),
62            metadata: BTreeMap::new(),
63            module_name: None,
64            state: Arc::new(WebSocketGatewayState::default()),
65        })
66    }
67
68    pub fn path(&self) -> &str {
69        &self.path
70    }
71
72    pub fn path_shape(&self) -> String {
73        route_shape_key(&self.path)
74    }
75
76    pub fn module_name(&self) -> Option<&str> {
77        self.module_name.as_deref()
78    }
79
80    pub fn namespace(&self) -> Option<&str> {
81        self.namespace.as_deref()
82    }
83
84    pub fn metadata(&self) -> &BTreeMap<String, Value> {
85        &self.metadata
86    }
87
88    pub fn metadata_value(&self, key: &str) -> Option<&Value> {
89        self.metadata.get(key)
90    }
91
92    pub fn event_metadata(&self, event: &str) -> Option<&BTreeMap<String, Value>> {
93        self.handlers
94            .get(event)
95            .map(WebSocketSubscriptionDefinition::metadata)
96    }
97
98    pub fn with_metadata<V>(self, key: impl Into<String>, value: V) -> Result<Self>
99    where
100        V: Serialize,
101    {
102        let key = key.into();
103        let value = serde_json::to_value(value).map_err(|error| {
104            BootError::Internal(format!(
105                "failed to serialize websocket gateway metadata `{key}`: {error}"
106            ))
107        })?;
108        Ok(self.with_metadata_value(key, value))
109    }
110
111    pub fn with_metadata_value(mut self, key: impl Into<String>, value: Value) -> Self {
112        let key = key.into();
113        self.metadata.insert(key.clone(), value.clone());
114        self.handlers = self
115            .handlers
116            .into_iter()
117            .map(|(event, handler)| {
118                (
119                    event,
120                    handler.with_metadata_default_value(key.clone(), value.clone()),
121                )
122            })
123            .collect();
124        self
125    }
126
127    pub fn with_namespace(mut self, namespace: impl Into<String>) -> Result<Self> {
128        self.namespace = Some(normalize_namespace(namespace)?);
129        Ok(self)
130    }
131
132    pub fn events(&self) -> Vec<&str> {
133        self.handlers.keys().map(String::as_str).collect()
134    }
135
136    pub fn server(&self) -> WebSocketGatewayServer {
137        WebSocketGatewayServer::new(self.clone())
138    }
139
140    pub fn active_connection_count(&self) -> Result<usize> {
141        self.state.connection_count()
142    }
143
144    pub fn active_connection_ids(&self) -> Result<Vec<u64>> {
145        self.state.connection_ids()
146    }
147
148    pub fn rooms(&self) -> Result<Vec<String>> {
149        self.state.rooms()
150    }
151
152    pub fn room_members(&self, room: impl Into<String>) -> Result<Vec<u64>> {
153        self.state.room_members(room)
154    }
155
156    pub async fn broadcast(&self, message: WebSocketMessage) -> Result<usize> {
157        let outbounds = self.state.broadcast_targets(None, None)?;
158        send_to_outbounds(outbounds, message).await
159    }
160
161    pub async fn broadcast_to_room(
162        &self,
163        room: impl Into<String>,
164        message: WebSocketMessage,
165    ) -> Result<usize> {
166        let room = normalize_room(room)?;
167        let outbounds = self.state.broadcast_targets(Some(&room), None)?;
168        send_to_outbounds(outbounds, message).await
169    }
170
171    /// Run gateway initialization hooks.
172    pub async fn after_init(&self) -> Result<()> {
173        let context = WebSocketGatewayInitContext::new(self);
174        for hook in &self.init_hooks {
175            hook.after_init(context.clone()).await?;
176        }
177        Ok(())
178    }
179
180    pub fn matches_path(&self, path: &str) -> bool {
181        match_path_shape(&self.path, path)
182    }
183
184    pub fn path_params(&self, path: &str) -> Result<Option<BTreeMap<String, String>>> {
185        match_path_params(&self.path, path)
186    }
187
188    pub fn subscribe<H, Fut, R>(mut self, event: impl Into<String>, handler: H) -> Result<Self>
189    where
190        H: Fn(WebSocketMessage) -> Fut + Send + Sync + 'static,
191        Fut: Future<Output = Result<R>> + Send + 'static,
192        R: IntoWebSocketReply + Send + 'static,
193    {
194        let event = event.into();
195        if event.trim().is_empty() {
196            return Err(BootError::BadRequest(
197                "websocket event cannot be empty".to_string(),
198            ));
199        }
200        if self.handlers.contains_key(&event) {
201            return Err(BootError::DuplicateRoute(format!(
202                "{} {}",
203                self.path, event
204            )));
205        }
206        self.handlers.insert(
207            event,
208            WebSocketSubscriptionDefinition::new(handler).with_metadata_defaults(&self.metadata),
209        );
210        Ok(self)
211    }
212
213    pub fn subscribe_with_connection<H, Fut, R>(
214        mut self,
215        event: impl Into<String>,
216        handler: H,
217    ) -> Result<Self>
218    where
219        H: Fn(WebSocketGatewayConnection, WebSocketMessage) -> Fut + Send + Sync + 'static,
220        Fut: Future<Output = Result<R>> + Send + 'static,
221        R: IntoWebSocketReply + Send + 'static,
222    {
223        let event = event.into();
224        if event.trim().is_empty() {
225            return Err(BootError::BadRequest(
226                "websocket event cannot be empty".to_string(),
227            ));
228        }
229        if self.handlers.contains_key(&event) {
230            return Err(BootError::DuplicateRoute(format!(
231                "{} {}",
232                self.path, event
233            )));
234        }
235        self.handlers.insert(
236            event,
237            WebSocketSubscriptionDefinition::new_with_connection(handler)
238                .with_metadata_defaults(&self.metadata),
239        );
240        Ok(self)
241    }
242
243    pub fn subscribe_with_server<H, Fut, R>(
244        mut self,
245        event: impl Into<String>,
246        handler: H,
247    ) -> Result<Self>
248    where
249        H: Fn(WebSocketGatewayServer, WebSocketMessage) -> Fut + Send + Sync + 'static,
250        Fut: Future<Output = Result<R>> + Send + 'static,
251        R: IntoWebSocketReply + Send + 'static,
252    {
253        let event = event.into();
254        if event.trim().is_empty() {
255            return Err(BootError::BadRequest(
256                "websocket event cannot be empty".to_string(),
257            ));
258        }
259        if self.handlers.contains_key(&event) {
260            return Err(BootError::DuplicateRoute(format!(
261                "{} {}",
262                self.path, event
263            )));
264        }
265        self.handlers.insert(
266            event,
267            WebSocketSubscriptionDefinition::new_with_server(handler)
268                .with_metadata_defaults(&self.metadata),
269        );
270        Ok(self)
271    }
272
273    pub fn subscribe_definition(
274        mut self,
275        event: impl Into<String>,
276        subscription: WebSocketSubscriptionDefinition,
277    ) -> Result<Self> {
278        let event = event.into();
279        if event.trim().is_empty() {
280            return Err(BootError::BadRequest(
281                "websocket event cannot be empty".to_string(),
282            ));
283        }
284        if self.handlers.contains_key(&event) {
285            return Err(BootError::DuplicateRoute(format!(
286                "{} {}",
287                self.path, event
288            )));
289        }
290        self.handlers
291            .insert(event, subscription.with_metadata_defaults(&self.metadata));
292        Ok(self)
293    }
294
295    pub fn with_after_init<H>(mut self, hook: H) -> Self
296    where
297        H: WebSocketGatewayInitHook,
298    {
299        self.init_hooks.push(Arc::new(hook));
300        self
301    }
302
303    pub fn with_connection_hook<H>(mut self, hook: H) -> Self
304    where
305        H: WebSocketGatewayConnectionHook,
306    {
307        self.connection_hooks.push(Arc::new(hook));
308        self
309    }
310
311    pub fn with_disconnect_hook<H>(mut self, hook: H) -> Self
312    where
313        H: WebSocketGatewayDisconnectHook,
314    {
315        self.disconnect_hooks.push(Arc::new(hook));
316        self
317    }
318
319    pub fn with_pipe<P>(mut self, pipe: P) -> Self
320    where
321        P: WebSocketPipe,
322    {
323        self.pipes
324            .push(PipelineComponent::<dyn WebSocketPipe>::new(pipe));
325        self
326    }
327
328    pub fn with_guard<G>(mut self, guard: G) -> Self
329    where
330        G: WebSocketGuard,
331    {
332        self.guards
333            .push(PipelineComponent::<dyn WebSocketGuard>::new(guard));
334        self
335    }
336
337    pub fn with_execution_guard<G>(mut self, guard: G) -> Self
338    where
339        G: Guard,
340    {
341        self.guards
342            .push(PipelineComponent::<dyn WebSocketGuard>::new(
343                ExecutionWebSocketGuard { inner: guard },
344            ));
345        self
346    }
347
348    pub(crate) fn with_execution_pipeline_prefix(
349        mut self,
350        guards: &[Arc<dyn Guard>],
351        interceptors: &[Arc<dyn ExecutionInterceptor>],
352    ) -> Self {
353        self.guards = prepend_execution_guards(guards, self.guards);
354        self.interceptors = prepend_execution_interceptors(interceptors, self.interceptors);
355        self
356    }
357
358    pub(crate) fn with_guard_prefix(mut self, guards: &[Arc<dyn WebSocketGuard>]) -> Self {
359        let mut merged = guards
360            .iter()
361            .cloned()
362            .map(PipelineComponent::<dyn WebSocketGuard>::from_arc)
363            .collect::<Vec<_>>();
364        merged.extend(self.guards);
365        self.guards = merged;
366        self
367    }
368
369    pub(crate) fn with_interceptor_prefix(
370        mut self,
371        interceptors: &[Arc<dyn WebSocketInterceptor>],
372    ) -> Self {
373        let mut merged = interceptors
374            .iter()
375            .cloned()
376            .map(PipelineComponent::<dyn WebSocketInterceptor>::from_arc)
377            .collect::<Vec<_>>();
378        merged.extend(self.interceptors);
379        self.interceptors = merged;
380        self
381    }
382
383    pub(crate) fn with_pipe_prefix(mut self, pipes: &[Arc<dyn WebSocketPipe>]) -> Self {
384        let mut merged = pipes
385            .iter()
386            .cloned()
387            .map(PipelineComponent::<dyn WebSocketPipe>::from_arc)
388            .collect::<Vec<_>>();
389        merged.extend(self.pipes);
390        self.pipes = merged;
391        self
392    }
393
394    pub(crate) fn with_filter_prefix(
395        mut self,
396        filters: &[Arc<dyn WebSocketExceptionFilter>],
397    ) -> Self {
398        let mut merged = filters
399            .iter()
400            .cloned()
401            .map(PipelineComponent::<dyn WebSocketExceptionFilter>::from_arc)
402            .collect::<Vec<_>>();
403        merged.extend(self.filters);
404        self.filters = merged;
405        self
406    }
407
408    pub(crate) fn with_pipeline_overrides(mut self, overrides: &PipelineOverrides) -> Self {
409        overrides.apply_to_websocket_pipes(&mut self.pipes);
410        overrides.apply_to_websocket_guards(&mut self.guards);
411        overrides.apply_to_websocket_interceptors(&mut self.interceptors);
412        overrides.apply_to_websocket_filters(&mut self.filters);
413        self.handlers = self
414            .handlers
415            .into_iter()
416            .map(|(event, subscription)| (event, subscription.with_pipeline_overrides(overrides)))
417            .collect();
418        self
419    }
420
421    pub(crate) fn with_validation_prefix(
422        mut self,
423        validation_enabled: bool,
424        validation_options: ValidationOptions,
425    ) -> Self {
426        self.handlers = self
427            .handlers
428            .into_iter()
429            .map(|(event, subscription)| {
430                (
431                    event,
432                    subscription.with_validation_prefix(validation_enabled, validation_options),
433                )
434            })
435            .collect();
436        self
437    }
438
439    pub fn with_interceptor<I>(mut self, interceptor: I) -> Self
440    where
441        I: WebSocketInterceptor,
442    {
443        self.interceptors
444            .push(PipelineComponent::<dyn WebSocketInterceptor>::new(
445                interceptor,
446            ));
447        self
448    }
449
450    pub fn with_execution_interceptor<I>(mut self, interceptor: I) -> Self
451    where
452        I: ExecutionInterceptor,
453    {
454        self.interceptors
455            .push(PipelineComponent::<dyn WebSocketInterceptor>::new(
456                ExecutionWebSocketInterceptor { inner: interceptor },
457            ));
458        self
459    }
460
461    pub fn with_filter<F>(mut self, filter: F) -> Self
462    where
463        F: WebSocketExceptionFilter,
464    {
465        self.filters
466            .push(PipelineComponent::<dyn WebSocketExceptionFilter>::new(
467                filter,
468            ));
469        self
470    }
471
472    pub fn with_catch_filter<I, F>(self, kinds: I, filter: F) -> Self
473    where
474        I: IntoIterator<Item = BootErrorKind>,
475        F: WebSocketExceptionFilter,
476    {
477        self.with_filter(catch_errors(kinds, filter))
478    }
479
480    pub fn connect(&self, request: BootRequest) -> Result<WebSocketGatewayConnection> {
481        if request.method() != HttpMethod::Get {
482            return Err(BootError::MethodNotAllowed(format!(
483                "{} {}",
484                request.method().as_str(),
485                request.path()
486            )));
487        }
488        let Some(params) = self.path_params(request.path())? else {
489            return Err(BootError::NotFound(format!(
490                "{} {}",
491                request.method().as_str(),
492                request.path()
493            )));
494        };
495        Ok(WebSocketGatewayConnection {
496            gateway: self.clone(),
497            id: self.state.next_connection_id(),
498            request: request.with_path_params(params),
499            outbound: None,
500            opened: Arc::new(AtomicBool::new(false)),
501        })
502    }
503
504    pub fn connect_with_outbound<O>(
505        &self,
506        request: BootRequest,
507        outbound: O,
508    ) -> Result<WebSocketGatewayConnection>
509    where
510        O: WebSocketOutbound,
511    {
512        let mut connection = self.connect(request)?;
513        connection.outbound = Some(Arc::new(outbound));
514        Ok(connection)
515    }
516
517    pub async fn connect_async(&self, request: BootRequest) -> Result<WebSocketGatewayConnection> {
518        let connection = self.connect(request)?;
519        connection.open().await?;
520        Ok(connection)
521    }
522
523    pub async fn connect_async_with_outbound<O>(
524        &self,
525        request: BootRequest,
526        outbound: O,
527    ) -> Result<WebSocketGatewayConnection>
528    where
529        O: WebSocketOutbound,
530    {
531        let connection = self.connect_with_outbound(request, outbound)?;
532        connection.open().await?;
533        Ok(connection)
534    }
535
536    pub async fn emit_to_connection(
537        &self,
538        connection_id: u64,
539        message: WebSocketMessage,
540    ) -> Result<bool> {
541        let outbounds = self
542            .state
543            .outbound_for_connection(connection_id)?
544            .into_iter()
545            .collect();
546        Ok(send_to_outbounds(outbounds, message).await? > 0)
547    }
548
549    pub async fn dispatch(
550        &self,
551        request: BootRequest,
552        message: WebSocketMessage,
553    ) -> Result<Option<WebSocketMessage>> {
554        self.connect(request)?.dispatch(message).await
555    }
556
557    pub(crate) fn with_path_prefix(mut self, prefix: &str) -> Result<Self> {
558        self.path = join_paths(prefix, &self.path)?;
559        Ok(self)
560    }
561
562    pub(crate) fn with_module_name(mut self, module_name: &str) -> Self {
563        self.module_name = Some(module_name.to_string());
564        self
565    }
566}
a3s_boot/websocket/gateway.rs

a3s_boot/websocket/
gateway.rs
