rama_http/layer/
set_status.rs

1//! Middleware to override status codes.
2//!
3//! # Example
4//!
5//! ```
6//! use std::{iter::once, convert::Infallible};
7//! use rama_core::bytes::Bytes;
8//! use rama_http::layer::set_status::SetStatusLayer;
9//! use rama_http::{Body, Request, Response, StatusCode};
10//! use rama_core::service::service_fn;
11//! use rama_core::{Layer, Service};
12//! use rama_core::error::BoxError;
13//!
14//! async fn handle(req: Request) -> Result<Response, Infallible> {
15//!     // ...
16//!     # Ok(Response::new(Body::empty()))
17//! }
18//!
19//! # #[tokio::main]
20//! # async fn main() -> Result<(), BoxError> {
21//! let mut service = (
22//!     // change the status to `404 Not Found` regardless what the inner service returns
23//!     SetStatusLayer::new(StatusCode::NOT_FOUND),
24//! ).into_layer(service_fn(handle));
25//!
26//! // Call the service.
27//! let request = Request::builder().body(Body::empty())?;
28//!
29//! let response = service.serve(request).await?;
30//!
31//! assert_eq!(response.status(), StatusCode::NOT_FOUND);
32//! #
33//! # Ok(())
34//! # }
35//! ```
36
37use crate::{Request, Response, StatusCode};
38use rama_core::{Layer, Service};
39use rama_utils::macros::define_inner_service_accessors;
40
41/// Layer that applies [`SetStatus`] which overrides the status codes.
42#[derive(Debug, Clone)]
43pub struct SetStatusLayer {
44    status: StatusCode,
45}
46
47impl SetStatusLayer {
48    /// Create a new [`SetStatusLayer`].
49    ///
50    /// The response status code will be `status` regardless of what the inner service returns.
51    #[must_use]
52    pub const fn new(status: StatusCode) -> Self {
53        Self { status }
54    }
55
56    /// Create a new [`SetStatusLayer`] layer which will create
57    /// a service that will always set the status code at [`StatusCode::OK`].
58    #[inline]
59    #[must_use]
60    pub const fn ok() -> Self {
61        Self::new(StatusCode::OK)
62    }
63}
64
65impl<S> Layer<S> for SetStatusLayer {
66    type Service = SetStatus<S>;
67
68    fn layer(&self, inner: S) -> Self::Service {
69        SetStatus::new(inner, self.status)
70    }
71}
72
73/// Middleware to override status codes.
74///
75/// See the [module docs](self) for more details.
76#[derive(Debug, Clone)]
77pub struct SetStatus<S> {
78    inner: S,
79    status: StatusCode,
80}
81
82impl<S> SetStatus<S> {
83    /// Create a new [`SetStatus`].
84    ///
85    /// The response status code will be `status` regardless of what the inner service returns.
86    pub const fn new(inner: S, status: StatusCode) -> Self {
87        Self { status, inner }
88    }
89
90    /// Create a new [`SetStatus`] service which will always set the
91    /// status code at [`StatusCode::OK`].
92    #[inline]
93    pub const fn ok(inner: S) -> Self {
94        Self::new(inner, StatusCode::OK)
95    }
96
97    define_inner_service_accessors!();
98}
99
100impl<S: Copy> Copy for SetStatus<S> {}
101
102impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for SetStatus<S>
103where
104    S: Service<Request<ReqBody>, Output = Response<ResBody>>,
105    ReqBody: Send + 'static,
106    ResBody: Send + 'static,
107{
108    type Output = S::Output;
109    type Error = S::Error;
110
111    async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
112        let mut response = self.inner.serve(req).await?;
113        *response.status_mut() = self.status;
114        Ok(response)
115    }
116}
rama_http/layer/set_status.rs

rama_http/layer/
set_status.rs
