#[cfg(tracing_unstable)]
use crate::json::JsonValuable;
use crate::response::IsProxy;
use crate::tcp::Sni;
use crate::{
GMT_FORMAT, HeadersDebug, LatencyDisplay, SERVER, WrappedRedactedHashingAlg, X_CORRELATION_ID,
X_REQUEST_ID, get_http_version_str, response_for_panic,
};
use ahash::AHasher;
use axum::Router;
use axum::body::HttpBody;
use axum::extract::Request;
use axum::http::{HeaderValue, StatusCode, header};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use base64::{Engine as B64Engine, engine::general_purpose::URL_SAFE_NO_PAD as b64};
use hyper::HeaderMap;
use hyper::header::LAST_MODIFIED;
use ordinary_config::{HttpCache, HttpEtagAlgorithm, XXH3Variation};
use std::hash::{BuildHasher, Hasher};
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use time::UtcDateTime;
use tower::ServiceBuilder;
use tower_http::catch_panic::CatchPanicLayer;
use tower_http::classify::ServerErrorsFailureClass;
use tower_http::compression::CompressionLayer;
use tower_http::decompression::RequestDecompressionLayer;
use tower_http::request_id::{
MakeRequestId, PropagateRequestIdLayer, RequestId, SetRequestIdLayer,
};
use tower_http::set_header::SetResponseHeaderLayer;
use tower_http::trace::TraceLayer;
use tracing::Span;
use uuid::Uuid;
#[derive(Clone, Debug)]
pub struct ViaHeader(pub String);
pub enum QueryFmt {
String(String),
#[cfg(tracing_unstable)]
Value(JsonValuable),
}
impl QueryFmt {
#[cfg(tracing_unstable)]
#[must_use]
pub fn trace_display(&self) -> valuable::Value<'_> {
match self {
Self::String(s) => tracing::field::valuable(s),
#[cfg(tracing_unstable)]
Self::Value(json) => tracing::field::valuable(json),
}
}
#[cfg(not(tracing_unstable))]
pub fn trace_display(&self) -> tracing::field::DebugValue<&String> {
match self {
Self::String(s) => tracing::field::debug(s),
}
}
}
#[derive(Clone, Default)]
struct MakeRequestUuidV7;
impl MakeRequestId for MakeRequestUuidV7 {
fn make_request_id<B>(&mut self, _: &Request<B>) -> Option<RequestId> {
let uuid = Uuid::now_v7().to_string();
uuid.parse().ok().map(RequestId::new)
}
}
#[derive(Clone)]
pub enum ServiceKind {
App,
Api,
Redirect,
Proxy,
}
#[allow(clippy::too_many_arguments, clippy::too_many_lines)]
pub fn apply_common_middleware<T>(
router: Router<T>,
state: &T,
server_span: Option<&Span>,
domain: String,
log_headers: bool,
log_ips: bool,
redacted_hash: Arc<Option<WrappedRedactedHashingAlg>>,
kind: ServiceKind,
via_domain: Option<String>,
) -> Router
where
T: Clone + Send + Sync + 'static,
{
let server_span = server_span.map(ToOwned::to_owned);
let redacted_hash_clone = redacted_hash.clone();
router
.with_state(state.clone())
.layer(
ServiceBuilder::new()
.layer(CatchPanicLayer::custom(response_for_panic))
.layer(RequestDecompressionLayer::new())
.layer(CompressionLayer::new()),
)
.layer(
ServiceBuilder::new()
.layer(SetRequestIdLayer::new(X_REQUEST_ID, MakeRequestUuidV7))
.layer(
TraceLayer::new_for_http()
.make_span_with(move |req: &axum::http::Request<_>| {
let request_id = req.extensions().get::<RequestId>().and_then(|rid| {
rid.header_value()
.to_str()
.ok()
.map(tracing::field::display)
});
let correlation_id =
req.headers().get(X_CORRELATION_ID).and_then(|cid| {
cid.to_str()
.ok()
.and_then(|cid| Uuid::parse_str(cid).ok())
.map(tracing::field::display)
});
let uri = req.uri();
let host = {
if let Some(sni) = req.extensions().get::<Sni>() {
Some(tracing::field::display(sni.0.as_str()))
} else {
req.headers()
.get(header::HOST)
.and_then(|host| host.to_str().ok())
.map(tracing::field::display)
}
};
let ip = crate::get_display_ip(log_ips, req);
let query = uri.query().map(tracing::field::display);
match kind {
ServiceKind::App => {
if let Some(server_span) = &server_span {
server_span.in_scope(|| {
tracing::info_span!(
"app",
%domain,
host,
ip,
rid = request_id,
cid = correlation_id,
path = %uri.path(),
query,
)
})
} else {
tracing::info_span!(
"app",
%domain,
host,
ip,
rid = request_id,
cid = correlation_id,
path = %uri.path(),
query,
)
}
}
ServiceKind::Proxy => {
if let Some(server_span) = &server_span {
server_span.in_scope(|| {
tracing::info_span!(
"proxy",
%domain,
host,
ip,
rid = request_id,
cid = correlation_id,
path = %uri.path(),
query,
)
})
} else {
tracing::info_span!(
"proxy",
%domain,
host,
ip,
rid = request_id,
cid = correlation_id,
path = %uri.path(),
query,
)
}
}
ServiceKind::Api => {
if let Some(server_span) = &server_span {
server_span.in_scope(|| {
tracing::info_span!(
"api",
%domain,
host,
ip,
rid = request_id,
cid = correlation_id,
path = %uri.path(),
query,
)
})
} else {
tracing::info_span!(
"api",
%domain,
host,
ip,
rid = request_id,
cid = correlation_id,
path = %uri.path(),
query,
)
}
}
ServiceKind::Redirect => {
if let Some(server_span) = &server_span {
server_span.in_scope(|| {
tracing::info_span!(
"redirect",
host,
ip,
rid = request_id,
cid = correlation_id,
path = %uri.path(),
query,
)
})
} else {
tracing::info_span!(
"redirect",
host,
ip,
rid = request_id,
cid = correlation_id,
path = %uri.path(),
query,
)
}
}
}
})
.on_request(move |req: &axum::http::Request<_>, _: &Span| {
let hd = log_headers
.then_some(HeadersDebug(req.headers(), redacted_hash.clone()));
#[cfg(tracing_unstable)]
let headers = log_headers.then_some(tracing::field::valuable(&hd));
#[cfg(not(tracing_unstable))]
let headers = log_headers.then_some(tracing::field::debug(&hd));
tracing::info!(
version = ?req.version(),
method = %req.method(),
headers,
"req"
);
})
.on_response(move |res: &Response<_>, latency: Duration, _: &Span| {
let hd = log_headers.then_some(HeadersDebug(
res.headers(),
redacted_hash_clone.clone(),
));
#[cfg(tracing_unstable)]
let headers = log_headers.then_some(tracing::field::valuable(&hd));
#[cfg(not(tracing_unstable))]
let headers = log_headers.then_some(tracing::field::debug(&hd));
let status = res.status().as_u16();
let latency = LatencyDisplay(latency.as_nanos() as f64);
if status >= 500 {
tracing::error!(status, headers, %latency, "res");
} else if status >= 400 {
tracing::warn!(status, headers, %latency, "res");
} else {
tracing::info!(status, headers, %latency, "res");
}
})
.on_failure(|error: ServerErrorsFailureClass, _: Duration, _: &Span| {
tracing::error!(
err = %error,
"fail"
);
}),
)
.layer(PropagateRequestIdLayer::new(X_REQUEST_ID))
.layer(SetResponseHeaderLayer::if_not_present(
header::SERVER,
HeaderValue::from_static(SERVER),
))
.layer(SetResponseHeaderLayer::overriding(
header::ETAG,
modify_etag_for_encoding,
))
.option_layer(via_domain.map(|domain| {
ServiceBuilder::new()
.layer(axum::middleware::from_fn(
move |mut req: axum::http::Request<_>, next: Next| {
let req_version = get_http_version_str(req.version());
req.extensions_mut().insert(ViaHeader(format!(
"{req_version} {domain} (ordinaryd)",
)));
next.run(req)
},
))
.layer(axum::middleware::from_fn(propagate_via))
})),
)
}
#[allow(clippy::similar_names)]
pub async fn http_cache_middleware(
last_modified: UtcDateTime,
last_modified_header: HeaderValue,
req_headers: HeaderMap,
request: Request,
next: Next,
) -> Response {
let response = next.run(request).await;
let (mut parts, body) = response.into_parts();
let body_bytes = if let Some(limit) = body.size_hint().upper()
&& let Ok(limit) = usize::try_from(limit)
&& let Ok(body_bytes) = axum::body::to_bytes(body, limit).await
{
body_bytes
} else {
return StatusCode::INTERNAL_SERVER_ERROR.into_response();
};
let mut res_headers = HeaderMap::new();
res_headers.insert(LAST_MODIFIED, last_modified_header);
let etag_string = get_etag_hash(body_bytes.as_ref(), None);
let etag_str = etag_string.as_str();
if let Some(if_none_match) = req_headers.get(header::IF_NONE_MATCH)
&& let Ok(if_none_match_str) = if_none_match.to_str()
&& if_none_match_str == etag_str
{
res_headers.insert(header::ETAG, if_none_match.to_owned());
return (StatusCode::NOT_MODIFIED, res_headers).into_response();
} else if let Ok(etag_header) = HeaderValue::from_str(etag_str) {
if let Some(if_modified_since) = req_headers.get(header::IF_MODIFIED_SINCE)
&& let Ok(if_modified_since_str) = if_modified_since.to_str()
&& let Ok(if_modified_since) = UtcDateTime::parse(if_modified_since_str, &GMT_FORMAT)
&& if_modified_since >= last_modified
{
res_headers.insert(header::ETAG, etag_header);
return (StatusCode::NOT_MODIFIED, res_headers).into_response();
}
parts.headers.insert(header::ETAG, etag_header);
}
(parts, body_bytes).into_response()
}
static FOLDHASH: OnceLock<foldhash::fast::FixedState> = OnceLock::new();
#[must_use]
pub fn get_etag_hash(content: &[u8], http_cache: Option<&HttpCache>) -> String {
if let Some(http_cache) = http_cache
&& let Some(etag_config) = &http_cache.etag
&& let Some(etag_alg) = &etag_config.alg
{
return match etag_alg {
HttpEtagAlgorithm::Foldhash => {
let hasher = FOLDHASH.get_or_init(foldhash::fast::FixedState::default);
let result = hasher.hash_one(content);
b64.encode(result.to_be_bytes())
}
HttpEtagAlgorithm::AHash => {
let mut hasher = AHasher::default();
hasher.write(content);
b64.encode(hasher.finish().to_be_bytes())
}
HttpEtagAlgorithm::XXH3(variation) => match variation {
XXH3Variation::Bit64 => {
b64.encode(xxhash_rust::xxh3::xxh3_64(content).to_be_bytes())
}
XXH3Variation::Bit128 => {
b64.encode(xxhash_rust::xxh3::xxh3_128(content).to_be_bytes())
}
},
HttpEtagAlgorithm::Rustc => {
let mut hasher = rustc_hash::FxHasher::default();
hasher.write(content);
b64.encode(hasher.finish().to_be_bytes())
}
HttpEtagAlgorithm::Blake3 => b64.encode(&blake3::hash(content).as_bytes()[0..16]),
};
}
let mut hasher = AHasher::default();
hasher.write(content);
b64.encode(hasher.finish().to_be_bytes())
}
#[must_use]
pub async fn propagate_via(request: Request, next: Next) -> Response {
let via_ext = request
.extensions()
.get::<ViaHeader>()
.map(ToOwned::to_owned);
let mut response = next.run(request).await;
if let Some(is_proxy) = response.extensions().get::<IsProxy>()
&& is_proxy.0
{
return response;
}
if let Some(via_ext) = via_ext
&& let Ok(header_value) = HeaderValue::from_str(&via_ext.0)
{
response.headers_mut().insert(header::VIA, header_value);
}
response
}
pub fn modify_etag_for_encoding(res: &Response) -> Option<HeaderValue> {
let headers = res.headers();
if let Some(is_proxy) = res.extensions().get::<IsProxy>()
&& is_proxy.0
{
return headers.get(header::ETAG).map(ToOwned::to_owned);
}
if let Some(curr_etag) = headers.get(header::ETAG)
&& let Ok(curr_etag_str) = curr_etag.to_str()
{
let etag_len = curr_etag_str.len();
if (etag_len == 22 || etag_len == 11)
&& let Some(compression) = headers.get(header::CONTENT_ENCODING)
&& let Ok(compression_str) = compression.to_str()
{
let mut etag_string = curr_etag_str.to_owned();
match compression_str {
"gzip" => etag_string.push('1'),
"zstd" => etag_string.push('2'),
"br" => etag_string.push('3'),
"deflate" => etag_string.push('4'),
_ => (),
}
match HeaderValue::from_str(etag_string.as_str()) {
Ok(v) => return Some(v),
Err(err) => tracing::error!(%err),
}
} else {
return Some(curr_etag.to_owned());
}
}
None
}
pub fn check_if_none_match<'a>(headers: &'a HeaderMap, etag: &'a str) -> Option<&'a str> {
if let Some(if_none_match) = headers.get(header::IF_NONE_MATCH)
&& let Ok(if_none_match_str) = if_none_match.to_str()
{
if if_none_match_str.len() < 11 {
return None;
}
if (etag.len() == 23
|| etag.len() == 12
|| if_none_match_str.len() == 22
|| if_none_match_str.len() == 11)
&& if_none_match_str == etag
{
return Some(etag);
}
if &if_none_match_str[..if_none_match_str.len() - 1] == etag {
Some(if_none_match_str)
} else {
None
}
} else {
None
}
}