Skip to main content

ordinary_utils/
middleware.rs

1// Copyright (C) 2026 Ordinary Labs, LLC.
2//
3// SPDX-License-Identifier: AGPL-3.0-only
4
5use crate::{
6    GMT_FORMAT, HeadersDebug, LatencyDisplay, SERVER, WrappedRedactedHashingAlg, X_CORRELATION_ID,
7    X_REQUEST_ID, X_VIA, get_host, get_http_version_str, response_for_panic,
8};
9use ahash::AHasher;
10use axum::Router;
11use axum::body::HttpBody;
12use axum::extract::Request;
13use axum::http::{HeaderValue, StatusCode, header};
14use axum::middleware::Next;
15use axum::response::{IntoResponse, Response};
16use base64::{Engine as B64Engine, engine::general_purpose::URL_SAFE_NO_PAD as b64};
17use hyper::HeaderMap;
18use hyper::header::LAST_MODIFIED;
19use ordinary_config::{HttpCache, HttpEtagAlgorithm, XXH3Variation};
20use std::hash::Hasher;
21use std::sync::Arc;
22use std::time::Duration;
23use time::UtcDateTime;
24use tower::ServiceBuilder;
25use tower_http::catch_panic::CatchPanicLayer;
26use tower_http::classify::ServerErrorsFailureClass;
27use tower_http::compression::CompressionLayer;
28use tower_http::decompression::RequestDecompressionLayer;
29use tower_http::request_id::{MakeRequestUuid, PropagateRequestIdLayer, SetRequestIdLayer};
30use tower_http::set_header::{SetRequestHeaderLayer, SetResponseHeaderLayer};
31use tower_http::trace::TraceLayer;
32use tracing::Span;
33use uuid::Uuid;
34
35#[derive(Clone)]
36pub enum ServiceKind {
37    App,
38    Api,
39    Redirect,
40    Proxy,
41}
42
43#[allow(clippy::too_many_arguments, clippy::too_many_lines)]
44pub fn apply_common_middleware<T>(
45    router: Router<T>,
46    state: &T,
47    server_span: Option<Span>,
48    domain: String,
49    log_headers: bool,
50    log_ips: bool,
51    redacted_hash: Arc<Option<WrappedRedactedHashingAlg>>,
52    kind: ServiceKind,
53    via_domain: Option<String>,
54) -> Router
55where
56    T: Clone + Send + Sync + 'static,
57{
58    let redacted_hash_clone = redacted_hash.clone();
59
60    router
61        .with_state(state.clone())
62        .layer(
63            ServiceBuilder::new()
64                .layer(CatchPanicLayer::custom(response_for_panic))
65                .layer(RequestDecompressionLayer::new())
66                .layer(CompressionLayer::new()),
67        )
68        .layer(
69            ServiceBuilder::new()
70                .layer(SetRequestIdLayer::new(X_REQUEST_ID, MakeRequestUuid))
71                .layer(
72                    TraceLayer::new_for_http()
73                        .make_span_with(move |req: &axum::http::Request<_>| {
74                            let request_id = req.headers().get(X_REQUEST_ID).and_then(|rid| {
75                                rid.to_str()
76                                    .ok()
77                                    .and_then(|rid| Uuid::parse_str(rid).ok())
78                                    .map(tracing::field::display)
79                            });
80
81                            let correlation_id =
82                                req.headers().get(X_CORRELATION_ID).and_then(|cid| {
83                                    cid.to_str()
84                                        .ok()
85                                        .and_then(|cid| Uuid::parse_str(cid).ok())
86                                        .map(tracing::field::display)
87                                });
88
89                            let host =
90                                get_host(req.headers(), req.uri()).map(tracing::field::display);
91
92                            let ip = crate::get_display_ip(log_ips, req);
93
94                            let query = req.uri().query().map(tracing::field::display);
95
96                            match kind {
97                                ServiceKind::App => {
98                                    if let Some(server_span) = &server_span {
99                                        server_span.in_scope(|| {
100                                            tracing::info_span!(
101                                                "app",
102                                                %domain,
103                                                host,
104                                                ip,
105                                                rid = request_id,
106                                                cid = correlation_id,
107                                                path = %req.uri().path(),
108                                                query,
109                                            )
110                                        })
111                                    } else {
112                                        tracing::info_span!(
113                                            "app",
114                                            %domain,
115                                            host,
116                                            ip,
117                                            rid = request_id,
118                                            cid = correlation_id,
119                                            path = %req.uri().path(),
120                                            query,
121                                        )
122                                    }
123                                }
124                                ServiceKind::Proxy => {
125                                    if let Some(server_span) = &server_span {
126                                        server_span.in_scope(|| {
127                                            tracing::info_span!(
128                                                "proxy",
129                                                %domain,
130                                                host,
131                                                ip,
132                                                rid = request_id,
133                                                cid = correlation_id,
134                                                path = %req.uri().path(),
135                                                query,
136                                            )
137                                        })
138                                    } else {
139                                        tracing::info_span!(
140                                            "proxy",
141                                            %domain,
142                                            host,
143                                            ip,
144                                            rid = request_id,
145                                            cid = correlation_id,
146                                            path = %req.uri().path(),
147                                            query,
148                                        )
149                                    }
150                                }
151                                ServiceKind::Api => {
152                                    if let Some(server_span) = &server_span {
153                                        server_span.in_scope(|| {
154                                            tracing::info_span!(
155                                                "api",
156                                                %domain,
157                                                host,
158                                                ip,
159                                                rid = request_id,
160                                                cid = correlation_id,
161                                                path = %req.uri().path(),
162                                                query,
163                                            )
164                                        })
165                                    } else {
166                                        tracing::info_span!(
167                                            "api",
168                                            %domain,
169                                            host,
170                                            ip,
171                                            rid = request_id,
172                                            cid = correlation_id,
173                                            path = %req.uri().path(),
174                                            query,
175                                        )
176                                    }
177                                }
178                                ServiceKind::Redirect => {
179                                    if let Some(server_span) = &server_span {
180                                        server_span.in_scope(|| {
181                                            tracing::info_span!(
182                                                "redirect",
183                                                host,
184                                                ip,
185                                                rid = request_id,
186                                                cid = correlation_id,
187                                                path = %req.uri().path(),
188                                                query,
189                                            )
190                                        })
191                                    } else {
192                                        tracing::info_span!(
193                                            "redirect",
194                                            host,
195                                            ip,
196                                            rid = request_id,
197                                            cid = correlation_id,
198                                            path = %req.uri().path(),
199                                            query,
200                                        )
201                                    }
202                                }
203                            }
204                        })
205                        .on_request(move |req: &axum::http::Request<_>, _: &Span| {
206                            let hd = log_headers
207                                .then_some(HeadersDebug(req.headers(), redacted_hash.clone()));
208
209                            #[cfg(tracing_unstable)]
210                            let headers = log_headers.then_some(tracing::field::valuable(&hd));
211                            #[cfg(not(tracing_unstable))]
212                            let headers = log_headers.then_some(tracing::field::debug(&hd));
213
214                            tracing::info!(
215                                version = ?req.version(),
216                                method = %req.method(),
217                                headers,
218                                "req"
219                            );
220                        })
221                        .on_response(move |res: &Response<_>, latency: Duration, _: &Span| {
222                            let hd = log_headers.then_some(HeadersDebug(
223                                res.headers(),
224                                redacted_hash_clone.clone(),
225                            ));
226
227                            #[cfg(tracing_unstable)]
228                            let headers = log_headers.then_some(tracing::field::valuable(&hd));
229
230                            #[cfg(not(tracing_unstable))]
231                            let headers = log_headers.then_some(tracing::field::debug(&hd));
232
233                            let status = res.status().as_u16();
234                            let latency = LatencyDisplay(latency.as_nanos() as f64);
235
236                            if status >= 500 {
237                                tracing::error!(status, headers, %latency, "res");
238                            } else if status >= 400 {
239                                tracing::warn!(status, headers, %latency, "res");
240                            } else {
241                                tracing::info!(status, headers, %latency, "res");
242                            }
243                        })
244                        .on_failure(|error: ServerErrorsFailureClass, _: Duration, _: &Span| {
245                            tracing::error!(
246                                err = %error,
247                                "fail"
248                            );
249                        }),
250                )
251                .layer(PropagateRequestIdLayer::new(X_REQUEST_ID))
252                .layer(SetResponseHeaderLayer::if_not_present(
253                    header::SERVER,
254                    HeaderValue::from_static(SERVER),
255                ))
256                .option_layer(via_domain.map(|domain| {
257                    SetRequestHeaderLayer::overriding(X_VIA, move |req: &axum::http::Request<_>| {
258                        let req_version = get_http_version_str(req.version());
259                        HeaderValue::from_str(&format!("{req_version} {domain} (ordinaryd)")).ok()
260                    })
261                })),
262        )
263}
264
265#[allow(clippy::similar_names)]
266pub async fn http_cache_middleware(
267    last_modified: UtcDateTime,
268    last_modified_header: HeaderValue,
269    req_headers: HeaderMap,
270    request: Request,
271    next: Next,
272) -> Response {
273    let response = next.run(request).await;
274    let (mut parts, body) = response.into_parts();
275
276    let body_bytes = if let Some(limit) = body.size_hint().upper()
277        && let Ok(limit) = usize::try_from(limit)
278        && let Ok(body_bytes) = axum::body::to_bytes(body, limit).await
279    {
280        body_bytes
281    } else {
282        return StatusCode::INTERNAL_SERVER_ERROR.into_response();
283    };
284
285    let mut res_headers = HeaderMap::new();
286    res_headers.insert(LAST_MODIFIED, last_modified_header);
287
288    let etag_string = get_etag_hash(body_bytes.as_ref(), None);
289    let etag_str = etag_string.as_str();
290
291    if let Some(if_none_match) = req_headers.get(header::IF_NONE_MATCH)
292        && let Ok(if_none_match_str) = if_none_match.to_str()
293        && if_none_match_str == etag_str
294    {
295        res_headers.insert(header::ETAG, if_none_match.to_owned());
296
297        return (StatusCode::NOT_MODIFIED, res_headers).into_response();
298    } else if let Ok(etag_header) = HeaderValue::from_str(etag_str) {
299        if let Some(if_modified_since) = req_headers.get(header::IF_MODIFIED_SINCE)
300            && let Ok(if_modified_since_str) = if_modified_since.to_str()
301            && let Ok(if_modified_since) = UtcDateTime::parse(if_modified_since_str, &GMT_FORMAT)
302            && if_modified_since >= last_modified
303        {
304            res_headers.insert(header::ETAG, etag_header);
305            return (StatusCode::NOT_MODIFIED, res_headers).into_response();
306        }
307
308        parts.headers.insert(header::ETAG, etag_header);
309    }
310
311    (parts, body_bytes).into_response()
312}
313
314#[must_use]
315pub fn get_etag_hash(content: &[u8], http_cache: Option<&HttpCache>) -> String {
316    if let Some(http_cache) = http_cache
317        && let Some(etag_config) = &http_cache.etag
318        && let Some(etag_alg) = &etag_config.alg
319    {
320        return match etag_alg {
321            HttpEtagAlgorithm::AHash => {
322                let mut hasher = AHasher::default();
323                hasher.write(content);
324                b64.encode(hasher.finish().to_be_bytes())
325            }
326            HttpEtagAlgorithm::XXH3(variation) => match variation {
327                XXH3Variation::Bit64 => {
328                    b64.encode(xxhash_rust::xxh3::xxh3_64(content).to_be_bytes())
329                }
330                XXH3Variation::Bit128 => {
331                    b64.encode(xxhash_rust::xxh3::xxh3_128(content).to_be_bytes())
332                }
333            },
334            HttpEtagAlgorithm::Rustc => {
335                let mut hasher = rustc_hash::FxHasher::default();
336                hasher.write(content);
337
338                b64.encode(hasher.finish().to_be_bytes())
339            }
340            HttpEtagAlgorithm::Blake3 => b64.encode(&blake3::hash(content).as_bytes()[0..16]),
341        };
342    }
343
344    let mut hasher = AHasher::default();
345    hasher.write(content);
346    b64.encode(hasher.finish().to_be_bytes())
347}
348
349// todo: switch to using request extensions
350#[must_use]
351pub async fn x_via(headers: HeaderMap, request: Request, next: Next) -> Response {
352    let mut response = next.run(request).await;
353
354    if let Some(x_via) = headers.get(X_VIA) {
355        response.headers_mut().insert(header::VIA, x_via.to_owned());
356    }
357
358    response
359}
360
361pub fn modify_etag_for_encoding(res: &Response) -> Option<HeaderValue> {
362    let headers = res.headers();
363
364    if let Some(curr_etag) = headers.get(header::ETAG)
365        && let Ok(curr_etag_str) = curr_etag.to_str()
366    {
367        let etag_len = curr_etag_str.len();
368
369        if (etag_len == 22 || etag_len == 11)
370            && let Some(compression) = headers.get(header::CONTENT_ENCODING)
371            && let Ok(compression_str) = compression.to_str()
372        {
373            let mut etag_string = curr_etag_str.to_owned();
374
375            match compression_str {
376                "gzip" => etag_string.push('1'),
377                "zstd" => etag_string.push('2'),
378                "br" => etag_string.push('3'),
379                "deflate" => etag_string.push('4'),
380                _ => (),
381            }
382
383            match HeaderValue::from_str(etag_string.as_str()) {
384                Ok(v) => return Some(v),
385                Err(err) => tracing::error!(%err),
386            }
387        } else {
388            return Some(curr_etag.clone());
389        }
390    }
391
392    None
393}
394
395pub fn check_if_none_match<'a>(headers: &'a HeaderMap, etag: &'a str) -> Option<&'a str> {
396    if let Some(if_none_match) = headers.get(header::IF_NONE_MATCH)
397        && let Ok(if_none_match_str) = if_none_match.to_str()
398    {
399        if if_none_match_str.len() < 11 {
400            return None;
401        }
402
403        if (etag.len() == 23
404            || etag.len() == 12
405            || if_none_match_str.len() == 22
406            || if_none_match_str.len() == 11)
407            && if_none_match_str == etag
408        {
409            return Some(etag);
410        }
411
412        if &if_none_match_str[..if_none_match_str.len() - 1] == etag {
413            Some(if_none_match_str)
414        } else {
415            None
416        }
417    } else {
418        None
419    }
420}