This repository has no description
1use axum::Router;
2use axum::body::Body;
3use http::header::ALT_SVC;
4use http::{HeaderValue, Request, Response, StatusCode};
5
6const ALT_SVC_MAX_AGE_SECS: u32 = 86_400;
7
8knot_types::scalar_newtype! {
9 pub struct Port(u16);
10}
11
12pub fn alt_svc_header(port: Port) -> HeaderValue {
13 let port = port.get();
14 HeaderValue::from_str(&format!("h3=\":{port}\"; ma={ALT_SVC_MAX_AGE_SECS}"))
15 .expect("alt-svc header value is valid ascii")
16}
17
18pub fn with_alt_svc(app: Router, port: Port) -> Router {
19 let value = alt_svc_header(port);
20 app.layer(axum::middleware::map_response(
21 move |mut response: Response<Body>| {
22 let value = value.clone();
23 async move {
24 if response.status() != StatusCode::SWITCHING_PROTOCOLS {
25 response.headers_mut().insert(ALT_SVC, value);
26 }
27 response
28 }
29 },
30 ))
31}
32
33pub fn with_host_from_authority(app: Router) -> Router {
34 app.layer(axum::middleware::map_request(
35 |mut request: Request<Body>| async move {
36 let authority = request
37 .uri()
38 .authority()
39 .map(|authority| HeaderValue::from_str(authority.as_str()));
40 match (
41 request.headers().contains_key(http::header::HOST),
42 authority,
43 ) {
44 (false, Some(Ok(value))) => {
45 request.headers_mut().insert(http::header::HOST, value);
46 request
47 }
48 _ => request,
49 }
50 },
51 ))
52}
53
54#[cfg(test)]
55mod tests {
56 use super::*;
57
58 use axum::routing::get;
59 use tower::ServiceExt;
60
61 #[test]
62 fn alt_svc_header_advertises_h3() {
63 assert_eq!(
64 alt_svc_header(Port::new(443)).to_str().unwrap(),
65 "h3=\":443\"; ma=86400"
66 );
67 }
68
69 #[tokio::test]
70 async fn alt_svc_added_to_responses_except_switching_protocols() {
71 let app = with_alt_svc(
72 Router::new().route("/ok", get(|| async { "ok" })).route(
73 "/upgrade",
74 get(|| async {
75 Response::builder()
76 .status(StatusCode::SWITCHING_PROTOCOLS)
77 .body(Body::empty())
78 .unwrap()
79 }),
80 ),
81 Port::new(443),
82 );
83
84 let normal = app
85 .clone()
86 .oneshot(Request::get("/ok").body(Body::empty()).unwrap())
87 .await
88 .unwrap();
89 assert_eq!(
90 normal.headers().get(ALT_SVC).and_then(|v| v.to_str().ok()),
91 Some("h3=\":443\"; ma=86400")
92 );
93
94 let upgrade = app
95 .oneshot(Request::get("/upgrade").body(Body::empty()).unwrap())
96 .await
97 .unwrap();
98 assert!(
99 upgrade.headers().get(ALT_SVC).is_none(),
100 "101 responses mustn't include Alt-Svc"
101 );
102 }
103}