a3s_boot/websocket/
mod.rs

1use crate::routing::path::{
2    join_paths, match_path_params, match_path_shape, route_shape_key, validate_route_path,
3};
4use crate::{
5    BootError, BootRequest, BoxFuture, ExecutionContext, ExecutionInterceptor, Guard, HttpMethod,
6    Result,
7};
8use serde::{Deserialize, Serialize};
9use serde_json::Value;
10use std::collections::BTreeMap;
11use std::future::Future;
12use std::ops::Deref;
13use std::sync::Arc;
14
15/// Adapter-neutral WebSocket message used by gateways and adapters.
16#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
17pub struct WebSocketMessage {
18    pub event: String,
19    #[serde(default)]
20    pub data: Value,
21}
22
23impl WebSocketMessage {
24    pub fn new(event: impl Into<String>, data: impl Into<Value>) -> Self {
25        Self {
26            event: event.into(),
27            data: data.into(),
28        }
29    }
30
31    pub fn event(&self) -> &str {
32        &self.event
33    }
34
35    pub fn data(&self) -> &Value {
36        &self.data
37    }
38
39    pub fn text(event: impl Into<String>, data: impl Into<String>) -> Self {
40        Self::new(event, Value::String(data.into()))
41    }
42
43    pub fn json<T>(event: impl Into<String>, data: &T) -> Result<Self>
44    where
45        T: Serialize,
46    {
47        Ok(Self::new(
48            event,
49            serde_json::to_value(data).map_err(|err| BootError::Internal(err.to_string()))?,
50        ))
51    }
52}
53
54/// Return value accepted by WebSocket gateway handlers.
55pub trait IntoWebSocketReply {
56    fn into_websocket_reply(self) -> Option<WebSocketMessage>;
57}
58
59impl IntoWebSocketReply for WebSocketMessage {
60    fn into_websocket_reply(self) -> Option<WebSocketMessage> {
61        Some(self)
62    }
63}
64
65impl IntoWebSocketReply for Option<WebSocketMessage> {
66    fn into_websocket_reply(self) -> Option<WebSocketMessage> {
67        self
68    }
69}
70
71impl IntoWebSocketReply for () {
72    fn into_websocket_reply(self) -> Option<WebSocketMessage> {
73        None
74    }
75}
76
77/// Context available to WebSocket guards and interceptors.
78#[derive(Debug, Clone)]
79pub struct WebSocketContext {
80    pub request: BootRequest,
81    pub gateway_path: String,
82    pub event: String,
83    pub module_name: Option<String>,
84    execution_context: ExecutionContext,
85}
86
87impl WebSocketContext {
88    fn new(gateway: &WebSocketGatewayDefinition, request: BootRequest, event: &str) -> Self {
89        let gateway_path = gateway.path.clone();
90        let event = event.to_string();
91        let module_name = gateway.module_name.clone();
92        let execution_context = ExecutionContext::websocket(
93            request.clone(),
94            gateway_path.clone(),
95            event.clone(),
96            module_name.clone(),
97        );
98        Self {
99            request,
100            gateway_path,
101            event,
102            module_name,
103            execution_context,
104        }
105    }
106
107    pub fn execution_context(&self) -> &ExecutionContext {
108        &self.execution_context
109    }
110
111    pub fn into_execution_context(self) -> ExecutionContext {
112        self.execution_context
113    }
114}
115
116impl Deref for WebSocketContext {
117    type Target = ExecutionContext;
118
119    fn deref(&self) -> &Self::Target {
120        self.execution_context()
121    }
122}
123
124/// Message transformation hook for WebSocket gateways.
125pub trait WebSocketPipe: Send + Sync + 'static {
126    fn transform(&self, message: WebSocketMessage) -> BoxFuture<'static, Result<WebSocketMessage>>;
127}
128
129impl<F, Fut> WebSocketPipe for F
130where
131    F: Fn(WebSocketMessage) -> Fut + Send + Sync + 'static,
132    Fut: Future<Output = Result<WebSocketMessage>> + Send + 'static,
133{
134    fn transform(&self, message: WebSocketMessage) -> BoxFuture<'static, Result<WebSocketMessage>> {
135        Box::pin(self(message))
136    }
137}
138
139/// Authorization hook for WebSocket gateway messages.
140pub trait WebSocketGuard: Send + Sync + 'static {
141    fn can_activate(&self, context: WebSocketContext) -> BoxFuture<'static, Result<bool>>;
142}
143
144impl<F, Fut> WebSocketGuard for F
145where
146    F: Fn(WebSocketContext) -> Fut + Send + Sync + 'static,
147    Fut: Future<Output = Result<bool>> + Send + 'static,
148{
149    fn can_activate(&self, context: WebSocketContext) -> BoxFuture<'static, Result<bool>> {
150        Box::pin(self(context))
151    }
152}
153
154struct ExecutionWebSocketGuard<G> {
155    inner: G,
156}
157
158impl<G> WebSocketGuard for ExecutionWebSocketGuard<G>
159where
160    G: Guard,
161{
162    fn can_activate(&self, context: WebSocketContext) -> BoxFuture<'static, Result<bool>> {
163        self.inner.can_activate(context.into_execution_context())
164    }
165}
166
167/// Around-handler hook for WebSocket gateway messages.
168pub trait WebSocketInterceptor: Send + Sync + 'static {
169    fn before(&self, _context: WebSocketContext) -> BoxFuture<'static, Result<()>> {
170        Box::pin(async { Ok(()) })
171    }
172
173    fn after(
174        &self,
175        _context: WebSocketContext,
176        reply: Option<WebSocketMessage>,
177    ) -> BoxFuture<'static, Result<Option<WebSocketMessage>>> {
178        Box::pin(async move { Ok(reply) })
179    }
180}
181
182struct ExecutionWebSocketInterceptor<I> {
183    inner: I,
184}
185
186impl<I> WebSocketInterceptor for ExecutionWebSocketInterceptor<I>
187where
188    I: ExecutionInterceptor,
189{
190    fn before(&self, context: WebSocketContext) -> BoxFuture<'static, Result<()>> {
191        self.inner.before(context.into_execution_context())
192    }
193
194    fn after(
195        &self,
196        context: WebSocketContext,
197        reply: Option<WebSocketMessage>,
198    ) -> BoxFuture<'static, Result<Option<WebSocketMessage>>> {
199        let future = self.inner.after(context.into_execution_context());
200        Box::pin(async move {
201            future.await?;
202            Ok(reply)
203        })
204    }
205}
206
207type WebSocketHandlerFuture = BoxFuture<'static, Result<Option<WebSocketMessage>>>;
208
209trait WebSocketMessageHandler: Send + Sync + 'static {
210    fn call(&self, message: WebSocketMessage) -> WebSocketHandlerFuture;
211}
212
213struct WebSocketHandlerAdapter<H> {
214    handler: H,
215}
216
217impl<H, Fut, R> WebSocketMessageHandler for WebSocketHandlerAdapter<H>
218where
219    H: Fn(WebSocketMessage) -> Fut + Send + Sync + 'static,
220    Fut: Future<Output = Result<R>> + Send + 'static,
221    R: IntoWebSocketReply + Send + 'static,
222{
223    fn call(&self, message: WebSocketMessage) -> WebSocketHandlerFuture {
224        let future = (self.handler)(message);
225        Box::pin(async move { Ok(future.await?.into_websocket_reply()) })
226    }
227}
228
229/// Framework-neutral WebSocket gateway definition.
230#[derive(Clone)]
231pub struct WebSocketGatewayDefinition {
232    path: String,
233    handlers: BTreeMap<String, Arc<dyn WebSocketMessageHandler>>,
234    pipes: Vec<Arc<dyn WebSocketPipe>>,
235    guards: Vec<Arc<dyn WebSocketGuard>>,
236    interceptors: Vec<Arc<dyn WebSocketInterceptor>>,
237    module_name: Option<String>,
238}
239
240impl WebSocketGatewayDefinition {
241    pub fn new(path: impl Into<String>) -> Result<Self> {
242        let path = path.into();
243        validate_route_path(&path)?;
244        Ok(Self {
245            path,
246            handlers: BTreeMap::new(),
247            pipes: Vec::new(),
248            guards: Vec::new(),
249            interceptors: Vec::new(),
250            module_name: None,
251        })
252    }
253
254    pub fn path(&self) -> &str {
255        &self.path
256    }
257
258    pub fn path_shape(&self) -> String {
259        route_shape_key(&self.path)
260    }
261
262    pub fn module_name(&self) -> Option<&str> {
263        self.module_name.as_deref()
264    }
265
266    pub fn events(&self) -> Vec<&str> {
267        self.handlers.keys().map(String::as_str).collect()
268    }
269
270    pub fn matches_path(&self, path: &str) -> bool {
271        match_path_shape(&self.path, path)
272    }
273
274    pub fn path_params(&self, path: &str) -> Result<Option<BTreeMap<String, String>>> {
275        match_path_params(&self.path, path)
276    }
277
278    pub fn subscribe<H, Fut, R>(mut self, event: impl Into<String>, handler: H) -> Result<Self>
279    where
280        H: Fn(WebSocketMessage) -> Fut + Send + Sync + 'static,
281        Fut: Future<Output = Result<R>> + Send + 'static,
282        R: IntoWebSocketReply + Send + 'static,
283    {
284        let event = event.into();
285        if event.trim().is_empty() {
286            return Err(BootError::BadRequest(
287                "websocket event cannot be empty".to_string(),
288            ));
289        }
290        if self.handlers.contains_key(&event) {
291            return Err(BootError::DuplicateRoute(format!(
292                "{} {}",
293                self.path, event
294            )));
295        }
296        self.handlers
297            .insert(event, Arc::new(WebSocketHandlerAdapter { handler }));
298        Ok(self)
299    }
300
301    pub fn with_pipe<P>(mut self, pipe: P) -> Self
302    where
303        P: WebSocketPipe,
304    {
305        self.pipes.push(Arc::new(pipe));
306        self
307    }
308
309    pub fn with_guard<G>(mut self, guard: G) -> Self
310    where
311        G: WebSocketGuard,
312    {
313        self.guards.push(Arc::new(guard));
314        self
315    }
316
317    pub fn with_execution_guard<G>(mut self, guard: G) -> Self
318    where
319        G: Guard,
320    {
321        self.guards
322            .push(Arc::new(ExecutionWebSocketGuard { inner: guard }));
323        self
324    }
325
326    pub(crate) fn with_execution_pipeline_prefix(
327        mut self,
328        guards: &[Arc<dyn Guard>],
329        interceptors: &[Arc<dyn ExecutionInterceptor>],
330    ) -> Self {
331        self.guards = prepend_execution_guards(guards, self.guards);
332        self.interceptors = prepend_execution_interceptors(interceptors, self.interceptors);
333        self
334    }
335
336    pub fn with_interceptor<I>(mut self, interceptor: I) -> Self
337    where
338        I: WebSocketInterceptor,
339    {
340        self.interceptors.push(Arc::new(interceptor));
341        self
342    }
343
344    pub fn with_execution_interceptor<I>(mut self, interceptor: I) -> Self
345    where
346        I: ExecutionInterceptor,
347    {
348        self.interceptors
349            .push(Arc::new(ExecutionWebSocketInterceptor {
350                inner: interceptor,
351            }));
352        self
353    }
354
355    pub fn connect(&self, request: BootRequest) -> Result<WebSocketGatewayConnection> {
356        if request.method() != HttpMethod::Get {
357            return Err(BootError::MethodNotAllowed(format!(
358                "{} {}",
359                request.method().as_str(),
360                request.path()
361            )));
362        }
363        let Some(params) = self.path_params(request.path())? else {
364            return Err(BootError::NotFound(format!(
365                "{} {}",
366                request.method().as_str(),
367                request.path()
368            )));
369        };
370        Ok(WebSocketGatewayConnection {
371            gateway: self.clone(),
372            request: request.with_path_params(params),
373        })
374    }
375
376    pub async fn dispatch(
377        &self,
378        request: BootRequest,
379        message: WebSocketMessage,
380    ) -> Result<Option<WebSocketMessage>> {
381        self.connect(request)?.dispatch(message).await
382    }
383
384    pub(crate) fn with_path_prefix(mut self, prefix: &str) -> Result<Self> {
385        self.path = join_paths(prefix, &self.path)?;
386        Ok(self)
387    }
388
389    pub(crate) fn with_module_name(mut self, module_name: &str) -> Self {
390        self.module_name = Some(module_name.to_string());
391        self
392    }
393}
394
395fn prepend_execution_guards(
396    prefix: &[Arc<dyn Guard>],
397    values: Vec<Arc<dyn WebSocketGuard>>,
398) -> Vec<Arc<dyn WebSocketGuard>> {
399    let mut merged = prefix
400        .iter()
401        .cloned()
402        .map(|guard| Arc::new(ExecutionWebSocketGuard { inner: guard }) as Arc<dyn WebSocketGuard>)
403        .collect::<Vec<_>>();
404    merged.extend(values);
405    merged
406}
407
408fn prepend_execution_interceptors(
409    prefix: &[Arc<dyn ExecutionInterceptor>],
410    values: Vec<Arc<dyn WebSocketInterceptor>>,
411) -> Vec<Arc<dyn WebSocketInterceptor>> {
412    let mut merged = prefix
413        .iter()
414        .cloned()
415        .map(|interceptor| {
416            Arc::new(ExecutionWebSocketInterceptor { inner: interceptor })
417                as Arc<dyn WebSocketInterceptor>
418        })
419        .collect::<Vec<_>>();
420    merged.extend(values);
421    merged
422}
423
424/// Adapter-neutral WebSocket connection.
425pub trait WebSocketConnection: Send + Sync {
426    fn request(&self) -> &BootRequest;
427
428    fn dispatch(
429        &self,
430        message: WebSocketMessage,
431    ) -> BoxFuture<'static, Result<Option<WebSocketMessage>>>;
432}
433
434/// In-process WebSocket gateway connection used by adapters and tests.
435#[derive(Clone)]
436pub struct WebSocketGatewayConnection {
437    gateway: WebSocketGatewayDefinition,
438    request: BootRequest,
439}
440
441impl WebSocketGatewayConnection {
442    pub fn request(&self) -> &BootRequest {
443        &self.request
444    }
445
446    pub async fn dispatch(
447        &self,
448        mut message: WebSocketMessage,
449    ) -> Result<Option<WebSocketMessage>> {
450        let event = message.event.clone();
451        let handler = self.gateway.handlers.get(&event).cloned().ok_or_else(|| {
452            BootError::NotFound(format!("websocket event {} {}", self.gateway.path, event))
453        })?;
454
455        let context = WebSocketContext::new(&self.gateway, self.request.clone(), &message.event);
456        for guard in &self.gateway.guards {
457            let can_activate = guard.can_activate(context.clone()).await?;
458            if !can_activate {
459                return Err(BootError::Forbidden(format!(
460                    "websocket event {} {}",
461                    self.gateway.path, message.event
462                )));
463            }
464        }
465
466        for interceptor in &self.gateway.interceptors {
467            interceptor.before(context.clone()).await?;
468        }
469
470        for pipe in &self.gateway.pipes {
471            message = pipe.transform(message).await?;
472        }
473
474        let mut reply = handler.call(message).await?;
475        for interceptor in self.gateway.interceptors.iter().rev() {
476            reply = interceptor.after(context.clone(), reply).await?;
477        }
478        Ok(reply)
479    }
480}
481
482impl WebSocketConnection for WebSocketGatewayConnection {
483    fn request(&self) -> &BootRequest {
484        self.request()
485    }
486
487    fn dispatch(
488        &self,
489        message: WebSocketMessage,
490    ) -> BoxFuture<'static, Result<Option<WebSocketMessage>>> {
491        let connection = self.clone();
492        Box::pin(async move { connection.dispatch(message).await })
493    }
494}
a3s_boot/websocket/mod.rs

a3s_boot/websocket/
mod.rs
