nidus_http/middleware/
security.rs1use std::{
2 task::{Context, Poll},
3 time::Duration,
4};
5
6use axum::{body::Body, extract::Request};
7use futures_util::{
8 FutureExt, TryFutureExt,
9 future::{Either, Map, MapOk, Ready, ready},
10};
11use http::{HeaderValue, Response, StatusCode, header};
12use tower::{Layer, Service};
13use tower_http::limit::RequestBodyLimitLayer;
14
15pub fn security_headers_layer() -> SecurityHeadersLayer {
24 SecurityHeadersLayer
25}
26
27#[derive(Clone, Copy, Debug, Default)]
33pub struct SecurityHeadersLayer;
34
35impl<S> Layer<S> for SecurityHeadersLayer {
36 type Service = SecurityHeadersService<S>;
37
38 fn layer(&self, inner: S) -> Self::Service {
39 SecurityHeadersService { inner }
40 }
41}
42
43#[derive(Clone, Debug)]
45pub struct SecurityHeadersService<S> {
46 inner: S,
47}
48
49impl<S> Service<Request> for SecurityHeadersService<S>
50where
51 S: Service<Request, Response = Response<Body>> + Send + 'static,
52 S::Future: Send + 'static,
53 S::Error: Send + 'static,
54{
55 type Response = Response<Body>;
56 type Error = S::Error;
57 type Future = MapOk<S::Future, fn(Response<Body>) -> Response<Body>>;
58
59 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
60 self.inner.poll_ready(cx)
61 }
62
63 fn call(&mut self, request: Request) -> Self::Future {
64 self.inner
65 .call(request)
66 .map_ok(apply_security_headers as fn(Response<Body>) -> Response<Body>)
67 }
68}
69
70fn apply_security_headers(mut response: Response<Body>) -> Response<Body> {
71 response.headers_mut().insert(
72 "x-content-type-options",
73 HeaderValue::from_static("nosniff"),
74 );
75 response
76 .headers_mut()
77 .insert("x-frame-options", HeaderValue::from_static("DENY"));
78 response
79 .headers_mut()
80 .insert("referrer-policy", HeaderValue::from_static("no-referrer"));
81 response
82}
83
84pub fn body_limit_layer(max_bytes: u64) -> BodyLimitLayer {
95 BodyLimitLayer {
96 max_bytes,
97 webhook_boundary: false,
98 }
99}
100
101pub fn streaming_body_limit_layer(max_bytes: usize) -> RequestBodyLimitLayer {
113 RequestBodyLimitLayer::new(max_bytes)
114}
115
116pub fn webhook_body_limit_layer(max_bytes: u64) -> BodyLimitLayer {
124 BodyLimitLayer {
125 max_bytes,
126 webhook_boundary: true,
127 }
128}
129
130#[derive(Clone, Copy, Debug)]
136pub struct BodyLimitLayer {
137 max_bytes: u64,
138 webhook_boundary: bool,
139}
140
141impl<S> Layer<S> for BodyLimitLayer {
142 type Service = BodyLimitService<S>;
143
144 fn layer(&self, inner: S) -> Self::Service {
145 BodyLimitService {
146 inner,
147 max_bytes: self.max_bytes,
148 webhook_boundary: self.webhook_boundary,
149 }
150 }
151}
152
153#[derive(Clone, Debug)]
155pub struct BodyLimitService<S> {
156 inner: S,
157 max_bytes: u64,
158 webhook_boundary: bool,
159}
160
161impl<S> Service<Request> for BodyLimitService<S>
162where
163 S: Service<Request, Response = Response<Body>> + Send + 'static,
164 S::Future: Send + 'static,
165 S::Error: Send + 'static,
166{
167 type Response = Response<Body>;
168 type Error = S::Error;
169 type Future = Either<Ready<Result<Self::Response, Self::Error>>, S::Future>;
170
171 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
172 self.inner.poll_ready(cx)
173 }
174
175 fn call(&mut self, request: Request) -> Self::Future {
176 let too_large = request
177 .headers()
178 .get(header::CONTENT_LENGTH)
179 .and_then(|value| value.to_str().ok())
180 .and_then(|value| value.parse::<u64>().ok())
181 .is_some_and(|length| length > self.max_bytes);
182 if too_large {
183 return Either::Left(ready(Ok(body_too_large_response(self.webhook_boundary))));
184 }
185
186 Either::Right(self.inner.call(request))
187 }
188}
189
190fn body_too_large_response(webhook_boundary: bool) -> Response<Body> {
191 let mut response = Response::new(Body::from("payload too large"));
192 *response.status_mut() = StatusCode::PAYLOAD_TOO_LARGE;
193 if webhook_boundary {
194 response.headers_mut().insert(
195 "x-nidus-body-limit",
196 HeaderValue::from_static("webhook-raw-body"),
197 );
198 }
199 response
200}
201
202pub fn timeout_response_layer(timeout: Duration) -> TimeoutResponseLayer {
208 TimeoutResponseLayer { timeout }
209}
210
211#[derive(Clone, Copy, Debug)]
217pub struct TimeoutResponseLayer {
218 timeout: Duration,
219}
220
221impl<S> Layer<S> for TimeoutResponseLayer {
222 type Service = TimeoutResponseService<S>;
223
224 fn layer(&self, inner: S) -> Self::Service {
225 TimeoutResponseService {
226 inner,
227 timeout: self.timeout,
228 }
229 }
230}
231
232#[derive(Clone, Debug)]
234pub struct TimeoutResponseService<S> {
235 inner: S,
236 timeout: Duration,
237}
238
239type TimeoutResult<E> = std::result::Result<Result<Response<Body>, E>, tokio::time::error::Elapsed>;
240type TimeoutResultMapper<E> = fn(TimeoutResult<E>) -> Result<Response<Body>, E>;
241
242impl<S> Service<Request> for TimeoutResponseService<S>
243where
244 S: Service<Request, Response = Response<Body>> + Send + 'static,
245 S::Future: Send + 'static,
246 S::Error: Send + 'static,
247{
248 type Response = Response<Body>;
249 type Error = S::Error;
250 type Future = Map<tokio::time::Timeout<S::Future>, TimeoutResultMapper<S::Error>>;
251
252 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
253 self.inner.poll_ready(cx)
254 }
255
256 fn call(&mut self, request: Request) -> Self::Future {
257 tokio::time::timeout(self.timeout, self.inner.call(request))
258 .map(map_timeout_result::<S::Error> as TimeoutResultMapper<S::Error>)
259 }
260}
261
262fn map_timeout_result<E>(result: TimeoutResult<E>) -> Result<Response<Body>, E> {
263 match result {
264 Ok(response) => response,
265 Err(_) => Ok(timeout_response()),
266 }
267}
268
269fn timeout_response() -> Response<Body> {
270 let mut response = Response::new(Body::from("request timed out"));
271 *response.status_mut() = StatusCode::REQUEST_TIMEOUT;
272 response
273}