rama_http/layer/
set_status.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
//! Middleware to override status codes.
//!
//! # Example
//!
//! ```
//! use std::{iter::once, convert::Infallible};
//! use bytes::Bytes;
//! use rama_http::layer::set_status::SetStatusLayer;
//! use rama_http::{Body, Request, Response, StatusCode};
//! use rama_core::service::service_fn;
//! use rama_core::{Context, Layer, Service};
//! use rama_core::error::BoxError;
//!
//! async fn handle(req: Request) -> Result<Response, Infallible> {
//!     // ...
//!     # Ok(Response::new(Body::empty()))
//! }
//!
//! # #[tokio::main]
//! # async fn main() -> Result<(), BoxError> {
//! let mut service = (
//!     // change the status to `404 Not Found` regardless what the inner service returns
//!     SetStatusLayer::new(StatusCode::NOT_FOUND),
//! ).layer(service_fn(handle));
//!
//! // Call the service.
//! let request = Request::builder().body(Body::empty())?;
//!
//! let response = service.serve(Context::default(), request).await?;
//!
//! assert_eq!(response.status(), StatusCode::NOT_FOUND);
//! #
//! # Ok(())
//! # }
//! ```

use std::fmt;

use crate::{Request, Response, StatusCode};
use rama_core::{Context, Layer, Service};
use rama_utils::macros::define_inner_service_accessors;

/// Layer that applies [`SetStatus`] which overrides the status codes.
#[derive(Debug, Clone)]
pub struct SetStatusLayer {
    status: StatusCode,
}

impl SetStatusLayer {
    /// Create a new [`SetStatusLayer`].
    ///
    /// The response status code will be `status` regardless of what the inner service returns.
    pub const fn new(status: StatusCode) -> Self {
        SetStatusLayer { status }
    }

    /// Create a new [`SetStatusLayer`] layer which will create
    /// a service that will always set the status code at [`StatusCode::OK`].
    #[inline]
    pub const fn ok() -> Self {
        Self::new(StatusCode::OK)
    }
}

impl<S> Layer<S> for SetStatusLayer {
    type Service = SetStatus<S>;

    fn layer(&self, inner: S) -> Self::Service {
        SetStatus::new(inner, self.status)
    }
}

/// Middleware to override status codes.
///
/// See the [module docs](self) for more details.
pub struct SetStatus<S> {
    inner: S,
    status: StatusCode,
}

impl<S> SetStatus<S> {
    /// Create a new [`SetStatus`].
    ///
    /// The response status code will be `status` regardless of what the inner service returns.
    pub const fn new(inner: S, status: StatusCode) -> Self {
        Self { status, inner }
    }

    /// Create a new [`SetStatus`] service which will always set the
    /// status code at [`StatusCode::OK`].
    #[inline]
    pub const fn ok(inner: S) -> Self {
        Self::new(inner, StatusCode::OK)
    }

    define_inner_service_accessors!();
}

impl<S: fmt::Debug> fmt::Debug for SetStatus<S> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SetStatus")
            .field("inner", &self.inner)
            .field("status", &self.status)
            .finish()
    }
}

impl<S: Clone> Clone for SetStatus<S> {
    fn clone(&self) -> Self {
        SetStatus {
            inner: self.inner.clone(),
            status: self.status,
        }
    }
}

impl<S: Copy> Copy for SetStatus<S> {}

impl<State, S, ReqBody, ResBody> Service<State, Request<ReqBody>> for SetStatus<S>
where
    State: Clone + Send + Sync + 'static,
    S: Service<State, Request<ReqBody>, Response = Response<ResBody>>,
    ReqBody: Send + 'static,
    ResBody: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;

    async fn serve(
        &self,
        ctx: Context<State>,
        req: Request<ReqBody>,
    ) -> Result<Self::Response, Self::Error> {
        let mut response = self.inner.serve(ctx, req).await?;
        *response.status_mut() = self.status;
        Ok(response)
    }
}