alice.pds.demo.boltless.dev
This repository has no description
did:plc:3fwecdnvtcscjnrx2p4n7a…
1use std::net::{IpAddr, SocketAddr};
2use std::num::{NonZeroU32, NonZeroU64};
3use std::sync::Arc;
4use std::time::Duration;
5
6use axum::Router;
7use axum::body::Body;
8use axum::error_handling::HandleErrorLayer;
9use axum::extract::{ConnectInfo, State};
10use axum::middleware::{Next, from_fn_with_state};
11use axum::response::{IntoResponse, Response};
12use governor::middleware::NoOpMiddleware;
13use http::{HeaderName, Method, Request, StatusCode};
14use tokio_util::sync::CancellationToken;
15use tower::limit::GlobalConcurrencyLimitLayer;
16use tower::load_shed::LoadShedLayer;
17use tower::{BoxError, ServiceBuilder};
18use tower_governor::GovernorLayer;
19use tower_governor::errors::GovernorError;
20use tower_governor::governor::{GovernorConfig, GovernorConfigBuilder};
21use tower_governor::key_extractor::KeyExtractor;
22use tower_http::map_request_body::MapRequestBodyLayer;
23use tower_http::timeout::{RequestBodyTimeoutLayer, TimeoutBody};
24
25const NANOS_PER_SECOND: u64 = 1_000_000_000;
26const CLEANUP_INTERVAL: Duration = Duration::from_secs(60);
27
28#[derive(Debug, Clone, Copy)]
29pub struct RequestsPerSecond(NonZeroU32);
30
31impl RequestsPerSecond {
32 pub fn new(value: NonZeroU32) -> Self {
33 Self(value)
34 }
35
36 fn period(self) -> Duration {
37 Duration::from_nanos((NANOS_PER_SECOND / u64::from(self.0.get())).max(1))
38 }
39}
40
41knot_types::scalar_newtype! {
42 pub struct BurstSize(NonZeroU32);
43 pub struct MaxInflightRequests(NonZeroU32);
44}
45
46#[derive(Debug, Clone, Copy)]
47pub struct RequestTimeout(Duration);
48
49impl RequestTimeout {
50 pub fn from_millis(millis: NonZeroU64) -> Self {
51 Self(Duration::from_millis(millis.get()))
52 }
53}
54
55#[derive(Debug, Clone, Copy)]
56pub struct BodyInactivityTimeout(Duration);
57
58impl BodyInactivityTimeout {
59 pub fn from_millis(millis: NonZeroU64) -> Self {
60 Self(Duration::from_millis(millis.get()))
61 }
62}
63
64#[derive(Debug, Clone, Copy)]
65pub struct WriteRequestTimeout(Duration);
66
67impl WriteRequestTimeout {
68 pub fn from_millis(millis: NonZeroU64) -> Self {
69 Self(Duration::from_millis(millis.get()))
70 }
71}
72
73pub struct EdgeGuards {
74 rate: RequestsPerSecond,
75 burst: BurstSize,
76 max_inflight: MaxInflightRequests,
77 request_timeout: RequestTimeout,
78 body_timeout: BodyInactivityTimeout,
79 write_request_timeout: WriteRequestTimeout,
80 proxy_header: Option<HeaderName>,
81}
82
83impl EdgeGuards {
84 pub fn new(
85 rate: RequestsPerSecond,
86 burst: BurstSize,
87 max_inflight: MaxInflightRequests,
88 request_timeout: RequestTimeout,
89 body_timeout: BodyInactivityTimeout,
90 write_request_timeout: WriteRequestTimeout,
91 proxy_header: Option<HeaderName>,
92 ) -> Self {
93 Self {
94 rate,
95 burst,
96 max_inflight,
97 request_timeout,
98 body_timeout,
99 write_request_timeout,
100 proxy_header,
101 }
102 }
103
104 pub(crate) fn prepare(self, shutdown: &CancellationToken) -> GuardLayers {
105 let governor = build_governor(self.rate, self.burst, self.proxy_header);
106 spawn_state_cleanup(Arc::clone(&governor), shutdown.clone());
107 GuardLayers {
108 governor,
109 max_inflight: self.max_inflight.0.get() as usize,
110 body_timeout: self.body_timeout.0,
111 timeout: TimeoutBudget {
112 standard: self.request_timeout.0,
113 extended: self.write_request_timeout.0,
114 },
115 }
116 }
117}
118
119type GuardGovernor = GovernorConfig<ProxyAwareIp, NoOpMiddleware>;
120
121pub(crate) struct GuardLayers {
122 governor: Arc<GuardGovernor>,
123 max_inflight: usize,
124 body_timeout: Duration,
125 timeout: TimeoutBudget,
126}
127
128#[derive(Clone, Copy)]
129struct TimeoutBudget {
130 standard: Duration,
131 extended: Duration,
132}
133
134#[derive(Clone)]
135struct ProxyAwareIp {
136 header: Option<HeaderName>,
137}
138
139impl KeyExtractor for ProxyAwareIp {
140 type Key = IpAddr;
141
142 fn extract<T>(&self, request: &Request<T>) -> Result<Self::Key, GovernorError> {
143 let from_header = self
144 .header
145 .as_ref()
146 .and_then(|header| knot_types::forwarded_peer(request.headers(), header));
147 from_header
148 .or_else(|| {
149 request
150 .extensions()
151 .get::<ConnectInfo<SocketAddr>>()
152 .map(|info| info.0.ip())
153 })
154 .ok_or(GovernorError::UnableToExtractKey)
155 }
156}
157
158fn build_governor(
159 rate: RequestsPerSecond,
160 burst: BurstSize,
161 proxy_header: Option<HeaderName>,
162) -> Arc<GuardGovernor> {
163 let mut builder = GovernorConfigBuilder::default();
164 builder.period(rate.period()).burst_size(burst.0.get());
165 let config = builder
166 .key_extractor(ProxyAwareIp {
167 header: proxy_header,
168 })
169 .finish()
170 .expect("a non-zero rate period and burst size always yield a governor config");
171 Arc::new(config)
172}
173
174fn spawn_state_cleanup(governor: Arc<GuardGovernor>, shutdown: CancellationToken) {
175 tokio::spawn(async move {
176 let mut ticker = tokio::time::interval(CLEANUP_INTERVAL);
177 ticker.tick().await;
178 loop {
179 tokio::select! {
180 () = shutdown.cancelled() => break,
181 _ = ticker.tick() => governor.limiter().retain_recent(),
182 }
183 }
184 });
185}
186
187async fn shed_overloaded(_error: BoxError) -> StatusCode {
188 StatusCode::SERVICE_UNAVAILABLE
189}
190
191async fn apply_request_timeout(
192 State(budget): State<TimeoutBudget>,
193 request: Request<Body>,
194 next: Next,
195) -> Response {
196 let limit = match is_streaming_write(&request) {
197 true => budget.extended,
198 false => budget.standard,
199 };
200 match tokio::time::timeout(limit, next.run(request)).await {
201 Ok(response) => response,
202 Err(_) => StatusCode::REQUEST_TIMEOUT.into_response(),
203 }
204}
205
206fn is_streaming_write(request: &Request<Body>) -> bool {
207 let path = request.uri().path();
208 match *request.method() {
209 Method::POST => path.ends_with("/git-receive-pack"),
210 Method::PUT => path.contains("/info/lfs/objects/"),
211 _ => false,
212 }
213}
214
215fn rewrap_body(body: TimeoutBody<axum::body::Body>) -> axum::body::Body {
216 axum::body::Body::new(body)
217}
218
219pub(crate) fn apply(router: Router, layers: GuardLayers) -> Router {
220 let GuardLayers {
221 governor,
222 max_inflight,
223 body_timeout,
224 timeout,
225 } = layers;
226 let rate_limit: GovernorLayer<ProxyAwareIp, NoOpMiddleware, axum::body::Body> =
227 GovernorLayer::new(governor);
228 router
229 .layer(
230 ServiceBuilder::new()
231 .layer(RequestBodyTimeoutLayer::new(body_timeout))
232 .layer(MapRequestBodyLayer::new(rewrap_body)),
233 )
234 .layer(from_fn_with_state(timeout, apply_request_timeout))
235 .layer(
236 ServiceBuilder::new()
237 .layer(HandleErrorLayer::new(shed_overloaded))
238 .layer(LoadShedLayer::new())
239 .layer(GlobalConcurrencyLimitLayer::new(max_inflight)),
240 )
241 .layer(rate_limit)
242}
243
244#[cfg(test)]
245mod tests {
246 use super::*;
247
248 use axum::body::{Body, Bytes};
249 use axum::routing::{get, post};
250 use futures::StreamExt;
251 use http::StatusCode;
252 use tower::ServiceExt;
253
254 fn guards(
255 rate: u32,
256 burst: u32,
257 inflight: u32,
258 request_timeout_ms: u64,
259 body_timeout_ms: u64,
260 proxy_header: Option<&str>,
261 ) -> EdgeGuards {
262 guards_with_write(
263 rate,
264 burst,
265 inflight,
266 request_timeout_ms,
267 body_timeout_ms,
268 request_timeout_ms,
269 proxy_header,
270 )
271 }
272
273 #[allow(clippy::too_many_arguments)]
274 fn guards_with_write(
275 rate: u32,
276 burst: u32,
277 inflight: u32,
278 request_timeout_ms: u64,
279 body_timeout_ms: u64,
280 write_request_timeout_ms: u64,
281 proxy_header: Option<&str>,
282 ) -> EdgeGuards {
283 EdgeGuards::new(
284 RequestsPerSecond::new(NonZeroU32::new(rate).unwrap()),
285 BurstSize::new(NonZeroU32::new(burst).unwrap()),
286 MaxInflightRequests::new(NonZeroU32::new(inflight).unwrap()),
287 RequestTimeout::from_millis(NonZeroU64::new(request_timeout_ms).unwrap()),
288 BodyInactivityTimeout::from_millis(NonZeroU64::new(body_timeout_ms).unwrap()),
289 WriteRequestTimeout::from_millis(NonZeroU64::new(write_request_timeout_ms).unwrap()),
290 proxy_header.map(|header| HeaderName::from_bytes(header.as_bytes()).unwrap()),
291 )
292 }
293
294 fn guarded_router(router: Router, guards: EdgeGuards) -> Router {
295 apply(router, guards.prepare(&CancellationToken::new()))
296 }
297
298 fn from_peer(request: Request<Body>, host: u8) -> Request<Body> {
299 let mut request = request;
300 request
301 .extensions_mut()
302 .insert(ConnectInfo(SocketAddr::from(([127, 0, 0, host], 47000))));
303 request
304 }
305
306 fn get_request() -> Request<Body> {
307 Request::get("/").body(Body::empty()).unwrap()
308 }
309
310 #[test]
311 fn the_extractor_keys_on_the_trusted_proxy_header_when_configured() {
312 let extractor = ProxyAwareIp {
313 header: Some(HeaderName::from_static("x-forwarded-for")),
314 };
315 let request = Request::get("/")
316 .header("x-forwarded-for", "203.0.113.7, 198.51.100.4")
317 .extension(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 5000))))
318 .body(())
319 .unwrap();
320 assert_eq!(
321 extractor.extract(&request).unwrap(),
322 "198.51.100.4".parse::<IpAddr>().unwrap(),
323 "the rightmost forwarded entry is the client the proxy appended"
324 );
325 }
326
327 #[test]
328 fn the_extractor_ignores_a_forgeable_header_when_no_proxy_is_trusted() {
329 let extractor = ProxyAwareIp { header: None };
330 let request = Request::get("/")
331 .header("x-forwarded-for", "203.0.113.7")
332 .extension(ConnectInfo(SocketAddr::from(([10, 0, 0, 9], 5000))))
333 .body(())
334 .unwrap();
335 assert_eq!(
336 extractor.extract(&request).unwrap(),
337 "10.0.0.9".parse::<IpAddr>().unwrap(),
338 "with no trusted proxy the socket peer wins over a client-forgeable header"
339 );
340 }
341
342 #[test]
343 fn the_extractor_falls_back_to_the_peer_when_the_trusted_header_is_absent() {
344 let extractor = ProxyAwareIp {
345 header: Some(HeaderName::from_static("x-forwarded-for")),
346 };
347 let request = Request::get("/")
348 .extension(ConnectInfo(SocketAddr::from(([10, 0, 0, 9], 5000))))
349 .body(())
350 .unwrap();
351 assert_eq!(
352 extractor.extract(&request).unwrap(),
353 "10.0.0.9".parse::<IpAddr>().unwrap()
354 );
355 }
356
357 #[test]
358 fn the_extractor_fails_when_no_peer_can_be_identified() {
359 let extractor = ProxyAwareIp { header: None };
360 let request = Request::get("/").body(()).unwrap();
361 assert!(matches!(
362 extractor.extract(&request),
363 Err(GovernorError::UnableToExtractKey)
364 ));
365 }
366
367 #[tokio::test]
368 async fn a_well_behaved_request_passes_every_guard() {
369 let app = guarded_router(
370 Router::new().route("/", get(|| async { "ok" })),
371 guards(50, 200, 1_024, 60_000, 30_000, None),
372 );
373 let status = app
374 .oneshot(from_peer(get_request(), 1))
375 .await
376 .unwrap()
377 .status();
378 assert_eq!(status, StatusCode::OK);
379 }
380
381 #[tokio::test]
382 async fn a_burst_beyond_the_per_ip_limit_is_rejected_with_429() {
383 let app = guarded_router(
384 Router::new().route("/", get(|| async { "ok" })),
385 guards(1, 2, 1_024, 60_000, 30_000, None),
386 );
387 let first = app
388 .clone()
389 .oneshot(from_peer(get_request(), 7))
390 .await
391 .unwrap()
392 .status();
393 let second = app
394 .clone()
395 .oneshot(from_peer(get_request(), 7))
396 .await
397 .unwrap()
398 .status();
399 let third = app
400 .clone()
401 .oneshot(from_peer(get_request(), 7))
402 .await
403 .unwrap()
404 .status();
405 let other_ip = app
406 .clone()
407 .oneshot(from_peer(get_request(), 8))
408 .await
409 .unwrap()
410 .status();
411 assert_eq!(first, StatusCode::OK);
412 assert_eq!(second, StatusCode::OK);
413 assert_eq!(
414 third,
415 StatusCode::TOO_MANY_REQUESTS,
416 "a third request inside the window exhausts the burst for this IP"
417 );
418 assert_eq!(
419 other_ip,
420 StatusCode::OK,
421 "a different IP keeps its own independent quota"
422 );
423 }
424
425 #[tokio::test]
426 async fn requests_beyond_the_inflight_limit_are_shed_with_503() {
427 let app = guarded_router(
428 Router::new().route(
429 "/slow",
430 get(|| async {
431 tokio::time::sleep(Duration::from_millis(300)).await;
432 "ok"
433 }),
434 ),
435 guards(10_000, 10_000, 1, 60_000, 30_000, None),
436 );
437 let holder = {
438 let app = app.clone();
439 tokio::spawn(async move {
440 let request = from_peer(Request::get("/slow").body(Body::empty()).unwrap(), 1);
441 app.oneshot(request).await.unwrap().status()
442 })
443 };
444 tokio::time::sleep(Duration::from_millis(50)).await;
445 let shed = app
446 .clone()
447 .oneshot(from_peer(
448 Request::get("/slow").body(Body::empty()).unwrap(),
449 2,
450 ))
451 .await
452 .unwrap()
453 .status();
454 assert_eq!(
455 shed,
456 StatusCode::SERVICE_UNAVAILABLE,
457 "with the single inflight slot held, the next request sheds rather than queues"
458 );
459 assert_eq!(holder.await.unwrap(), StatusCode::OK);
460 }
461
462 #[tokio::test]
463 async fn a_request_slower_than_the_timeout_is_cut_with_408() {
464 let app = guarded_router(
465 Router::new().route(
466 "/slow",
467 get(|| async {
468 tokio::time::sleep(Duration::from_millis(500)).await;
469 "ok"
470 }),
471 ),
472 guards(10_000, 10_000, 1_024, 80, 30_000, None),
473 );
474 let status = app
475 .oneshot(from_peer(
476 Request::get("/slow").body(Body::empty()).unwrap(),
477 1,
478 ))
479 .await
480 .unwrap()
481 .status();
482 assert_eq!(status, StatusCode::REQUEST_TIMEOUT);
483 }
484
485 #[tokio::test]
486 async fn a_streaming_write_runs_under_the_extended_budget_while_reads_keep_the_standard_one() {
487 let app = guarded_router(
488 Router::new()
489 .route(
490 "/did/name/git-receive-pack",
491 post(|| async {
492 tokio::time::sleep(Duration::from_millis(200)).await;
493 "ok"
494 }),
495 )
496 .route(
497 "/did/name/git-upload-pack",
498 post(|| async {
499 tokio::time::sleep(Duration::from_millis(200)).await;
500 "ok"
501 }),
502 ),
503 guards_with_write(10_000, 10_000, 1_024, 80, 30_000, 5_000, None),
504 );
505 let push = app
506 .clone()
507 .oneshot(from_peer(
508 Request::post("/did/name/git-receive-pack")
509 .body(Body::empty())
510 .unwrap(),
511 1,
512 ))
513 .await
514 .unwrap()
515 .status();
516 assert_eq!(
517 push,
518 StatusCode::OK,
519 "a push slower than the standard timeout survives on the extended write budget"
520 );
521 let fetch = app
522 .oneshot(from_peer(
523 Request::post("/did/name/git-upload-pack")
524 .body(Body::empty())
525 .unwrap(),
526 1,
527 ))
528 .await
529 .unwrap()
530 .status();
531 assert_eq!(
532 fetch,
533 StatusCode::REQUEST_TIMEOUT,
534 "a non-write request past the standard timeout is still cut"
535 );
536 }
537
538 #[tokio::test]
539 async fn a_stalled_request_body_is_cut_and_never_hangs() {
540 let app = guarded_router(
541 Router::new().route("/upload", post(|_body: Bytes| async { "ok" })),
542 guards(10_000, 10_000, 1_024, 60_000, 80, None),
543 );
544 let body = Body::from_stream(
545 futures::stream::once(async {
546 Ok::<_, std::io::Error>(Bytes::from_static(b"partial"))
547 })
548 .chain(futures::stream::pending::<Result<Bytes, std::io::Error>>()),
549 );
550 let request = from_peer(Request::post("/upload").body(body).unwrap(), 1);
551 let status = tokio::time::timeout(Duration::from_secs(5), app.oneshot(request))
552 .await
553 .expect("the body inactivity timeout must cut a stalled upload instead of hanging")
554 .unwrap()
555 .status();
556 assert_ne!(
557 status,
558 StatusCode::OK,
559 "a body that stalls past the inactivity timeout is never accepted"
560 );
561 }
562}