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