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
//! Middleware to override status codes.
//!
//! # Example
//!
//! ```
//! use std::{iter::once, convert::Infallible};
//! use rama_core::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::{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),
//! ).into_layer(service_fn(handle));
//!
//! // Call the service.
//! let request = Request::builder().body(Body::empty())?;
//!
//! let response = service.serve(request).await?;
//!
//! assert_eq!(response.status(), StatusCode::NOT_FOUND);
//! #
//! # Ok(())
//! # }
//! ```
use crate::{Request, Response, StatusCode};
use rama_core::{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.
#[must_use]
pub const fn new(status: StatusCode) -> Self {
Self { status }
}
/// Create a new [`SetStatusLayer`] layer which will create
/// a service that will always set the status code at [`StatusCode::OK`].
#[inline]
#[must_use]
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.
#[derive(Debug, Clone)]
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: Copy> Copy for SetStatus<S> {}
impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for SetStatus<S>
where
S: Service<Request<ReqBody>, Output = Response<ResBody>>,
ReqBody: Send + 'static,
ResBody: Send + 'static,
{
type Output = S::Output;
type Error = S::Error;
async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
let mut response = self.inner.serve(req).await?;
*response.status_mut() = self.status;
Ok(response)
}
}