a3s_boot/websocket/
pipeline.rs

1use super::context::WebSocketContext;
2use super::message::WebSocketMessage;
3use crate::pipeline::PipelineComponent;
4use crate::{BoxFuture, CallHandler, ExecutionInterceptor, Guard, Result};
5use std::future::Future;
6use std::sync::Arc;
7
8/// Message transformation hook for WebSocket gateways.
9pub trait WebSocketPipe: Send + Sync + 'static {
10    fn transform(&self, message: WebSocketMessage) -> BoxFuture<'static, Result<WebSocketMessage>>;
11}
12
13impl<F, Fut> WebSocketPipe for F
14where
15    F: Fn(WebSocketMessage) -> Fut + Send + Sync + 'static,
16    Fut: Future<Output = Result<WebSocketMessage>> + Send + 'static,
17{
18    fn transform(&self, message: WebSocketMessage) -> BoxFuture<'static, Result<WebSocketMessage>> {
19        Box::pin(self(message))
20    }
21}
22
23/// Authorization hook for WebSocket gateway messages.
24pub trait WebSocketGuard: Send + Sync + 'static {
25    fn can_activate(&self, context: WebSocketContext) -> BoxFuture<'static, Result<bool>>;
26}
27
28impl<F, Fut> WebSocketGuard for F
29where
30    F: Fn(WebSocketContext) -> Fut + Send + Sync + 'static,
31    Fut: Future<Output = Result<bool>> + Send + 'static,
32{
33    fn can_activate(&self, context: WebSocketContext) -> BoxFuture<'static, Result<bool>> {
34        Box::pin(self(context))
35    }
36}
37
38pub(crate) struct ExecutionWebSocketGuard<G> {
39    pub(crate) inner: G,
40}
41
42impl<G> WebSocketGuard for ExecutionWebSocketGuard<G>
43where
44    G: Guard,
45{
46    fn can_activate(&self, context: WebSocketContext) -> BoxFuture<'static, Result<bool>> {
47        self.inner.can_activate(context.into_execution_context())
48    }
49}
50
51/// Around-handler hook for WebSocket gateway messages.
52pub trait WebSocketInterceptor: Send + Sync + 'static {
53    /// Run around the remaining WebSocket pipeline.
54    ///
55    /// Override this method to recover downstream errors, retry the handler,
56    /// or return a reply without calling `next`. The default implementation
57    /// preserves the legacy `before` and `after` hook behavior.
58    fn intercept<'a>(
59        &'a self,
60        context: WebSocketContext,
61        next: CallHandler<'a, Option<WebSocketMessage>>,
62    ) -> BoxFuture<'a, Result<Option<WebSocketMessage>>> {
63        Box::pin(async move {
64            self.before(context.clone()).await?;
65            let reply = next.handle().await?;
66            self.after(context, reply).await
67        })
68    }
69
70    fn before(&self, _context: WebSocketContext) -> BoxFuture<'static, Result<()>> {
71        Box::pin(async { Ok(()) })
72    }
73
74    fn after(
75        &self,
76        _context: WebSocketContext,
77        reply: Option<WebSocketMessage>,
78    ) -> BoxFuture<'static, Result<Option<WebSocketMessage>>> {
79        Box::pin(async move { Ok(reply) })
80    }
81}
82
83pub(crate) struct ExecutionWebSocketInterceptor<I> {
84    pub(crate) inner: I,
85}
86
87impl<I> WebSocketInterceptor for ExecutionWebSocketInterceptor<I>
88where
89    I: ExecutionInterceptor,
90{
91    fn before(&self, context: WebSocketContext) -> BoxFuture<'static, Result<()>> {
92        self.inner.before(context.into_execution_context())
93    }
94
95    fn after(
96        &self,
97        context: WebSocketContext,
98        reply: Option<WebSocketMessage>,
99    ) -> BoxFuture<'static, Result<Option<WebSocketMessage>>> {
100        let future = self.inner.after(context.into_execution_context());
101        Box::pin(async move {
102            future.await?;
103            Ok(reply)
104        })
105    }
106}
107
108pub(crate) fn prepend_execution_guards(
109    prefix: &[Arc<dyn Guard>],
110    values: Vec<PipelineComponent<dyn WebSocketGuard>>,
111) -> Vec<PipelineComponent<dyn WebSocketGuard>> {
112    let mut merged = prefix
113        .iter()
114        .cloned()
115        .map(|guard| {
116            PipelineComponent::<dyn WebSocketGuard>::new(ExecutionWebSocketGuard { inner: guard })
117        })
118        .collect::<Vec<_>>();
119    merged.extend(values);
120    merged
121}
122
123pub(crate) fn prepend_execution_interceptors(
124    prefix: &[Arc<dyn ExecutionInterceptor>],
125    values: Vec<PipelineComponent<dyn WebSocketInterceptor>>,
126) -> Vec<PipelineComponent<dyn WebSocketInterceptor>> {
127    let mut merged = prefix
128        .iter()
129        .cloned()
130        .map(|interceptor| {
131            PipelineComponent::<dyn WebSocketInterceptor>::new(ExecutionWebSocketInterceptor {
132                inner: interceptor,
133            })
134        })
135        .collect::<Vec<_>>();
136    merged.extend(values);
137    merged
138}
a3s_boot/websocket/pipeline.rs

a3s_boot/websocket/
pipeline.rs
