This repository has no description
0

Configure Feed

Select the types of activity you want to include in your feed.

core / knot2 / crates / knot-edge / src / altsvc.rs
3.0 kB 103 lines
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}