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

a3s_boot/websocket/
gateway.rs
