Skip to main content

a3s_boot/websocket/
pipeline.rs

1use super::context::WebSocketContext;
2use super::message::WebSocketMessage;
3use crate::pipeline::PipelineComponent;
4use crate::{BoxFuture, 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    fn before(&self, _context: WebSocketContext) -> BoxFuture<'static, Result<()>> {
54        Box::pin(async { Ok(()) })
55    }
56
57    fn after(
58        &self,
59        _context: WebSocketContext,
60        reply: Option<WebSocketMessage>,
61    ) -> BoxFuture<'static, Result<Option<WebSocketMessage>>> {
62        Box::pin(async move { Ok(reply) })
63    }
64}
65
66pub(crate) struct ExecutionWebSocketInterceptor<I> {
67    pub(crate) inner: I,
68}
69
70impl<I> WebSocketInterceptor for ExecutionWebSocketInterceptor<I>
71where
72    I: ExecutionInterceptor,
73{
74    fn before(&self, context: WebSocketContext) -> BoxFuture<'static, Result<()>> {
75        self.inner.before(context.into_execution_context())
76    }
77
78    fn after(
79        &self,
80        context: WebSocketContext,
81        reply: Option<WebSocketMessage>,
82    ) -> BoxFuture<'static, Result<Option<WebSocketMessage>>> {
83        let future = self.inner.after(context.into_execution_context());
84        Box::pin(async move {
85            future.await?;
86            Ok(reply)
87        })
88    }
89}
90
91pub(crate) fn prepend_execution_guards(
92    prefix: &[Arc<dyn Guard>],
93    values: Vec<PipelineComponent<dyn WebSocketGuard>>,
94) -> Vec<PipelineComponent<dyn WebSocketGuard>> {
95    let mut merged = prefix
96        .iter()
97        .cloned()
98        .map(|guard| {
99            PipelineComponent::<dyn WebSocketGuard>::new(ExecutionWebSocketGuard { inner: guard })
100        })
101        .collect::<Vec<_>>();
102    merged.extend(values);
103    merged
104}
105
106pub(crate) fn prepend_execution_interceptors(
107    prefix: &[Arc<dyn ExecutionInterceptor>],
108    values: Vec<PipelineComponent<dyn WebSocketInterceptor>>,
109) -> Vec<PipelineComponent<dyn WebSocketInterceptor>> {
110    let mut merged = prefix
111        .iter()
112        .cloned()
113        .map(|interceptor| {
114            PipelineComponent::<dyn WebSocketInterceptor>::new(ExecutionWebSocketInterceptor {
115                inner: interceptor,
116            })
117        })
118        .collect::<Vec<_>>();
119    merged.extend(values);
120    merged
121}