This repository has no description
1use axum::Router;
2use axum::extract::Request;
3use axum::handler::Handler;
4use axum::middleware::{Next, from_fn};
5use axum::response::{IntoResponse, Response};
6use axum::routing::get;
7use http::StatusCode;
8
9use crate::protocol::NegotiatedProtocol;
10
11pub struct ZeroRttSafe<H> {
12 handler: H,
13}
14
15impl<H> ZeroRttSafe<H> {
16 pub fn new(handler: H) -> Self {
17 Self { handler }
18 }
19}
20
21pub struct RequiresFullHandshake {
22 router: Router,
23}
24
25impl RequiresFullHandshake {
26 pub fn new(router: Router) -> Self {
27 Self {
28 router: router.layer(from_fn(reject_early_writes)),
29 }
30 }
31
32 pub(crate) fn into_router(self) -> Router {
33 self.router
34 }
35}
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum EarlyData {
39 Yes,
40 No,
41}
42
43impl EarlyData {
44 pub fn is_early(self) -> bool {
45 matches!(self, EarlyData::Yes)
46 }
47}
48
49const EARLY_DATA_HEADER: &str = "early-data";
50
51pub struct ZeroRttRoutes {
52 router: Router,
53 count: usize,
54}
55
56impl ZeroRttRoutes {
57 pub fn new() -> Self {
58 Self {
59 router: Router::new(),
60 count: 0,
61 }
62 }
63
64 pub fn get<H, T>(mut self, path: &str, handler: ZeroRttSafe<H>) -> Self
65 where
66 H: Handler<T, ()>,
67 T: 'static,
68 {
69 self.router = self.router.route(path, get(handler.handler));
70 self.count += 1;
71 self
72 }
73
74 pub fn into_router(self) -> Router {
75 self.router
76 }
77
78 pub(crate) fn early_data_policy(&self) -> EarlyDataPolicy {
79 match self.count {
80 0 => EarlyDataPolicy::Disabled,
81 _ => EarlyDataPolicy::Enabled,
82 }
83 }
84}
85
86impl Default for ZeroRttRoutes {
87 fn default() -> Self {
88 Self::new()
89 }
90}
91
92#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub(crate) enum EarlyDataPolicy {
94 Disabled,
95 Enabled,
96}
97
98impl EarlyDataPolicy {
99 pub(crate) fn max_early_data_size(self) -> u32 {
100 match self {
101 EarlyDataPolicy::Disabled => 0,
102 // rustls takes 0 or u32::MAX,
103 // and quinn unwraps errors, quite unfortunate.
104 //
105 // The 0 above is hard-off marker, for when no route opted in.
106 EarlyDataPolicy::Enabled => u32::MAX,
107 }
108 }
109}
110
111pub(crate) async fn tag_from_header(mut request: Request, next: Next) -> Response {
112 if request.extensions().get::<EarlyData>().is_none() {
113 let early = request
114 .headers()
115 .get(EARLY_DATA_HEADER)
116 .and_then(|value| value.to_str().ok())
117 .is_some_and(|value| value.trim() == "1");
118 request
119 .extensions_mut()
120 .insert(if early { EarlyData::Yes } else { EarlyData::No });
121 }
122 next.run(request).await
123}
124
125async fn reject_early_writes(request: Request, next: Next) -> Response {
126 let early = request
127 .extensions()
128 .get::<EarlyData>()
129 .copied()
130 .unwrap_or(EarlyData::Yes);
131 if early.is_early() {
132 let protocol = request
133 .extensions()
134 .get::<NegotiatedProtocol>()
135 .map(|protocol| protocol.as_str())
136 .unwrap_or("unknown");
137 tracing::debug!(
138 protocol,
139 path = %request.uri().path(),
140 "refused an early-data request to a full-handshake route with 425"
141 );
142 return too_early();
143 }
144 next.run(request).await
145}
146
147fn too_early() -> Response {
148 (
149 StatusCode::TOO_EARLY,
150 "this route refuses early data and needs a completed handshake",
151 )
152 .into_response()
153}
154
155#[cfg(test)]
156mod tests {
157 use super::*;
158
159 #[test]
160 fn an_empty_safe_set_keeps_early_data_disabled() {
161 assert_eq!(
162 ZeroRttRoutes::new().early_data_policy(),
163 EarlyDataPolicy::Disabled
164 );
165 assert_eq!(EarlyDataPolicy::Disabled.max_early_data_size(), 0);
166 }
167
168 #[test]
169 fn a_proven_safe_set_unlocks_a_nonzero_early_data_size() {
170 let routes = ZeroRttRoutes::new().get("/info/refs", ZeroRttSafe::new(|| async { "ok" }));
171 assert_eq!(routes.early_data_policy(), EarlyDataPolicy::Enabled);
172 assert_eq!(EarlyDataPolicy::Enabled.max_early_data_size(), u32::MAX);
173 }
174
175 #[tokio::test]
176 async fn a_full_handshake_route_fails_closed_when_the_request_is_unclassified() {
177 use axum::body::Body;
178 use axum::routing::post;
179 use http::{Request, StatusCode};
180 use tower::ServiceExt;
181
182 let router = Router::new()
183 .route("/git-upload-pack", post(|| async { "pack" }))
184 .layer(from_fn(reject_early_writes));
185 let request = Request::post("/git-upload-pack")
186 .body(Body::empty())
187 .unwrap();
188 assert_eq!(
189 router.oneshot(request).await.unwrap().status(),
190 StatusCode::TOO_EARLY,
191 "a write whose early-data status was never tagged must fail closed with 425"
192 );
193 }
194}