rama_http/layer/upgrade/
layer.rs

1use crate::io::upgrade::Upgraded;
2
3use super::{UpgradeOutput, UpgradeResponse};
4
5use super::{UpgradeService, service::UpgradeHandler};
6use rama_core::error::BoxError;
7use rama_core::error_sink::{DropErrorSink, ErrorSink, TracingErrorSink};
8use rama_core::{Layer, Service, matcher::Matcher, rt::Executor};
9use rama_http_types::Request;
10use std::{fmt, sync::Arc};
11
12/// UpgradeLayer is a middleware that can be used to upgrade a request.
13///
14/// See [`UpgradeService`] for more details.
15///
16/// [`UpgradeService`]: crate::layer::upgrade::UpgradeService
17pub struct UpgradeLayer<O> {
18    handlers: Vec<Arc<UpgradeHandler<O>>>,
19    exec: Executor,
20    error_sink: Arc<dyn ErrorSink>,
21}
22
23impl<O: Send + 'static> UpgradeLayer<O> {
24    /// Create a new upgrade layer from one service.
25    ///
26    /// The service responds to the upgrade request with an [`UpgradeOutput`]
27    /// containing the response-local handler for the upgraded connection. This
28    /// lets that handler directly own resources established before the response.
29    /// Use [`Self::new_with_services`] when the responder and handler are separate.
30    pub fn new<M, R>(exec: Executor, matcher: M, service: R) -> Self
31    where
32        M: Matcher<Request>,
33        R: Service<Request, Output = UpgradeOutput<Request, O>, Error = O> + Clone,
34    {
35        Self::new_with_error_sink(exec, matcher, service, TracingErrorSink::default())
36    }
37
38    /// Create a single-service upgrade layer with a custom handler error sink.
39    pub fn new_with_error_sink<M, R, Sink>(
40        exec: Executor,
41        matcher: M,
42        service: R,
43        sink: Sink,
44    ) -> Self
45    where
46        M: Matcher<Request>,
47        R: Service<Request, Output = UpgradeOutput<Request, O>, Error = O> + Clone,
48        Sink: ErrorSink,
49    {
50        Self {
51            handlers: vec![Arc::new(UpgradeHandler::new(matcher, service, sink))],
52            exec,
53            error_sink: Arc::new(TracingErrorSink::default()),
54        }
55    }
56
57    /// Create a single-service upgrade layer which silently drops handler errors.
58    pub fn new_dropping_errors<M, R>(exec: Executor, matcher: M, service: R) -> Self
59    where
60        M: Matcher<Request>,
61        R: Service<Request, Output = UpgradeOutput<Request, O>, Error = O> + Clone,
62    {
63        Self::new_with_error_sink(exec, matcher, service, DropErrorSink::new())
64    }
65
66    /// Create an upgrade layer with separate responder and handler services.
67    ///
68    /// This is useful when the handshake can be acknowledged independently and
69    /// the upgraded-connection service performs any deferred setup itself.
70    pub fn new_with_services<M, R, H>(exec: Executor, matcher: M, responder: R, handler: H) -> Self
71    where
72        M: Matcher<Request>,
73        R: Service<Request, Output = UpgradeResponse<Request, O>, Error = O> + Clone,
74        H: Service<Upgraded, Error: Into<BoxError>> + Clone,
75    {
76        Self::new_with_services_and_error_sink(
77            exec,
78            matcher,
79            responder,
80            handler,
81            TracingErrorSink::default(),
82        )
83    }
84
85    /// Create a separate-services upgrade layer with a custom handler error sink.
86    pub fn new_with_services_and_error_sink<M, R, H, Sink>(
87        exec: Executor,
88        matcher: M,
89        responder: R,
90        handler: H,
91        sink: Sink,
92    ) -> Self
93    where
94        M: Matcher<Request>,
95        R: Service<Request, Output = UpgradeResponse<Request, O>, Error = O> + Clone,
96        H: Service<Upgraded> + Clone,
97        Sink: ErrorSink<H::Error>,
98    {
99        Self {
100            handlers: vec![Arc::new(UpgradeHandler::new_with_services(
101                matcher, responder, handler, sink,
102            ))],
103            exec,
104            error_sink: Arc::new(TracingErrorSink::default()),
105        }
106    }
107
108    /// Create a separate-services upgrade layer which silently drops handler errors.
109    pub fn new_with_services_dropping_errors<M, R, H>(
110        exec: Executor,
111        matcher: M,
112        responder: R,
113        handler: H,
114    ) -> Self
115    where
116        M: Matcher<Request>,
117        R: Service<Request, Output = UpgradeResponse<Request, O>, Error = O> + Clone,
118        H: Service<Upgraded> + Clone,
119    {
120        Self::new_with_services_and_error_sink(
121            exec,
122            matcher,
123            responder,
124            handler,
125            DropErrorSink::new(),
126        )
127    }
128
129    /// Add a single-service upgrade handler.
130    #[must_use]
131    pub fn on<M, R>(self, matcher: M, service: R) -> Self
132    where
133        M: Matcher<Request>,
134        R: Service<Request, Output = UpgradeOutput<Request, O>, Error = O> + Clone,
135    {
136        self.on_with_error_sink(matcher, service, TracingErrorSink::default())
137    }
138
139    /// Add a single-service upgrade handler with a custom handler error sink.
140    #[must_use]
141    pub fn on_with_error_sink<M, R, Sink>(mut self, matcher: M, service: R, sink: Sink) -> Self
142    where
143        M: Matcher<Request>,
144        R: Service<Request, Output = UpgradeOutput<Request, O>, Error = O> + Clone,
145        Sink: ErrorSink,
146    {
147        self.handlers
148            .push(Arc::new(UpgradeHandler::new(matcher, service, sink)));
149        self
150    }
151
152    /// Add a single-service upgrade handler which silently drops handler errors.
153    #[must_use]
154    pub fn on_dropping_errors<M, R>(self, matcher: M, service: R) -> Self
155    where
156        M: Matcher<Request>,
157        R: Service<Request, Output = UpgradeOutput<Request, O>, Error = O> + Clone,
158    {
159        self.on_with_error_sink(matcher, service, DropErrorSink::new())
160    }
161
162    /// Add separate responder and upgraded-connection services.
163    #[must_use]
164    pub fn on_with_services<M, R, H>(self, matcher: M, responder: R, handler: H) -> Self
165    where
166        M: Matcher<Request>,
167        R: Service<Request, Output = UpgradeResponse<Request, O>, Error = O> + Clone,
168        H: Service<Upgraded, Error: Into<BoxError>> + Clone,
169    {
170        self.on_with_services_and_error_sink(
171            matcher,
172            responder,
173            handler,
174            TracingErrorSink::default(),
175        )
176    }
177
178    /// Add separate services with a custom handler error sink.
179    #[must_use]
180    pub fn on_with_services_and_error_sink<M, R, H, Sink>(
181        mut self,
182        matcher: M,
183        responder: R,
184        handler: H,
185        sink: Sink,
186    ) -> Self
187    where
188        M: Matcher<Request>,
189        R: Service<Request, Output = UpgradeResponse<Request, O>, Error = O> + Clone,
190        H: Service<Upgraded> + Clone,
191        Sink: ErrorSink<H::Error>,
192    {
193        self.handlers
194            .push(Arc::new(UpgradeHandler::new_with_services(
195                matcher, responder, handler, sink,
196            )));
197        self
198    }
199
200    /// Add separate services while silently dropping handler errors.
201    #[must_use]
202    pub fn on_with_services_dropping_errors<M, R, H>(
203        self,
204        matcher: M,
205        responder: R,
206        handler: H,
207    ) -> Self
208    where
209        M: Matcher<Request>,
210        R: Service<Request, Output = UpgradeResponse<Request, O>, Error = O> + Clone,
211        H: Service<Upgraded> + Clone,
212    {
213        self.on_with_services_and_error_sink(matcher, responder, handler, DropErrorSink::new())
214    }
215
216    /// Set the [`ErrorSink`] used for errors that occur while *establishing*
217    /// the upgraded connection (i.e. the HTTP upgrade itself fails, before any
218    /// handler runs). Per-handler errors are routed to their own sink instead;
219    /// see [`Self::on_with_error_sink`] and
220    /// [`Self::on_with_services_and_error_sink`].
221    ///
222    /// Defaults to [`TracingErrorSink::default`] (traces at DEBUG level).
223    #[must_use]
224    pub fn with_upgrade_error_sink(mut self, sink: impl ErrorSink) -> Self {
225        self.error_sink = Arc::new(sink);
226        self
227    }
228}
229
230impl<O> fmt::Debug for UpgradeLayer<O> {
231    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
232        f.debug_struct("UpgradeLayer")
233            .field("handlers", &self.handlers)
234            .finish()
235    }
236}
237
238impl<O> Clone for UpgradeLayer<O> {
239    fn clone(&self) -> Self {
240        Self {
241            handlers: self.handlers.clone(),
242            exec: self.exec.clone(),
243            error_sink: self.error_sink.clone(),
244        }
245    }
246}
247
248impl<S, O> Layer<S> for UpgradeLayer<O> {
249    type Service = UpgradeService<S, O>;
250
251    fn layer(&self, inner: S) -> Self::Service {
252        UpgradeService::new(
253            self.handlers.clone(),
254            inner,
255            self.exec.clone(),
256            self.error_sink.clone(),
257        )
258    }
259
260    fn into_layer(self, inner: S) -> Self::Service {
261        UpgradeService::new(self.handlers, inner, self.exec, self.error_sink)
262    }
263}
rama_http/layer/upgrade/layer.rs

rama_http/layer/upgrade/
layer.rs
