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, OnceLock};
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::{Method, Request, StatusCode};
14use knot_types::ProxyTrust;
15use tokio_util::sync::CancellationToken;
16use tower::limit::GlobalConcurrencyLimitLayer;
17use tower::load_shed::LoadShedLayer;
18use tower::{BoxError, ServiceBuilder};
19use tower_governor::GovernorLayer;
20use tower_governor::errors::GovernorError;
21use tower_governor::governor::{GovernorConfig, GovernorConfigBuilder};
22use tower_governor::key_extractor::KeyExtractor;
23use tower_http::map_request_body::MapRequestBodyLayer;
24use tower_http::timeout::{RequestBodyTimeoutLayer, TimeoutBody};
25
26const NANOS_PER_SECOND: u64 = 1_000_000_000;
27const CLEANUP_INTERVAL: Duration = Duration::from_secs(60);
28
29#[derive(Debug, Clone, Copy)]
30pub struct RequestsPerSecond(NonZeroU32);
31
32impl RequestsPerSecond {
33 pub fn new(value: NonZeroU32) -> Self {
34 Self(value)
35 }
36
37 fn period(self) -> Duration {
38 Duration::from_nanos((NANOS_PER_SECOND / u64::from(self.0.get())).max(1))
39 }
40}
41
42knot_types::scalar_newtype! {
43 pub struct BurstSize(NonZeroU32);
44 pub struct MaxInflightRequests(NonZeroU32);
45}
46
47#[derive(Debug, Clone, Copy)]
48pub struct RequestTimeout(Duration);
49
50impl RequestTimeout {
51 pub fn from_millis(millis: NonZeroU64) -> Self {
52 Self(Duration::from_millis(millis.get()))
53 }
54}
55
56#[derive(Debug, Clone, Copy)]
57pub struct BodyInactivityTimeout(Duration);
58
59impl BodyInactivityTimeout {
60 pub fn from_millis(millis: NonZeroU64) -> Self {
61 Self(Duration::from_millis(millis.get()))
62 }
63}
64
65#[derive(Debug, Clone, Copy)]
66pub struct WriteRequestTimeout(Duration);
67
68impl WriteRequestTimeout {
69 pub fn from_millis(millis: NonZeroU64) -> Self {
70 Self(Duration::from_millis(millis.get()))
71 }
72}
73
74pub struct EdgeGuards {
75 rate: RequestsPerSecond,
76 burst: BurstSize,
77 max_inflight: MaxInflightRequests,
78 request_timeout: RequestTimeout,
79 body_timeout: BodyInactivityTimeout,
80 write_request_timeout: WriteRequestTimeout,
81 proxy_trust: ProxyTrust,
82}
83
84impl EdgeGuards {
85 pub fn new(
86 rate: RequestsPerSecond,
87 burst: BurstSize,
88 max_inflight: MaxInflightRequests,
89 request_timeout: RequestTimeout,
90 body_timeout: BodyInactivityTimeout,
91 write_request_timeout: WriteRequestTimeout,
92 proxy_trust: ProxyTrust,
93 ) -> Self {
94 Self {
95 rate,
96 burst,
97 max_inflight,
98 request_timeout,
99 body_timeout,
100 write_request_timeout,
101 proxy_trust,
102 }
103 }
104
105 pub(crate) fn prepare(self, shutdown: &CancellationToken) -> GuardLayers {
106 let governor = build_governor(self.rate, self.burst, self.proxy_trust);
107 spawn_state_cleanup(Arc::clone(&governor), shutdown.clone());
108 GuardLayers {
109 governor,
110 max_inflight: self.max_inflight.0.get() as usize,
111 body_timeout: self.body_timeout.0,
112 timeout: TimeoutBudget {
113 standard: self.request_timeout.0,
114 extended: self.write_request_timeout.0,
115 },
116 }
117 }
118}
119
120type GuardGovernor = GovernorConfig<ProxyAwareIp, NoOpMiddleware>;
121
122pub(crate) struct GuardLayers {
123 governor: Arc<GuardGovernor>,
124 max_inflight: usize,
125 body_timeout: Duration,
126 timeout: TimeoutBudget,
127}
128
129#[derive(Clone, Copy)]
130struct TimeoutBudget {
131 standard: Duration,
132 extended: Duration,
133}
134
135#[derive(Clone, Default)]
136struct IgnoredHeaderNotice(Arc<OnceLock<IpAddr>>);
137
138impl IgnoredHeaderNotice {
139 fn report(&self, peer: IpAddr) {
140 if self.0.set(peer).is_ok() {
141 tracing::warn!(
142 %peer,
143 "a peer outside xrpc.trusted_proxies sent xrpc.trusted_proxy_header, so the knot ignored the header and rate-limits that peer by the address it connected from. Add this address to xrpc.trusted_proxies if it is the reverse proxy, since a proxy reaches the knot over one address family and listing the other one silently loses the header. This warning reports the first such peer only."
144 );
145 }
146 }
147}
148
149#[derive(Clone)]
150struct ProxyAwareIp {
151 trust: ProxyTrust,
152 ignored_header: IgnoredHeaderNotice,
153}
154
155impl KeyExtractor for ProxyAwareIp {
156 type Key = IpAddr;
157
158 fn extract<T>(&self, request: &Request<T>) -> Result<Self::Key, GovernorError> {
159 let socket = request
160 .extensions()
161 .get::<ConnectInfo<SocketAddr>>()
162 .map(|info| info.0.ip());
163 let key = self.trust.peer_key(request.headers(), socket);
164 if let Some(peer) = key.ignored_header() {
165 self.ignored_header.report(peer);
166 }
167 key.address().ok_or(GovernorError::UnableToExtractKey)
168 }
169}
170
171fn build_governor(
172 rate: RequestsPerSecond,
173 burst: BurstSize,
174 proxy_trust: ProxyTrust,
175) -> Arc<GuardGovernor> {
176 let mut builder = GovernorConfigBuilder::default();
177 builder.period(rate.period()).burst_size(burst.0.get());
178 let config = builder
179 .key_extractor(ProxyAwareIp {
180 trust: proxy_trust,
181 ignored_header: IgnoredHeaderNotice::default(),
182 })
183 .finish()
184 .expect("a non-zero rate period and burst size always yield a governor config");
185 Arc::new(config)
186}
187
188fn spawn_state_cleanup(governor: Arc<GuardGovernor>, shutdown: CancellationToken) {
189 tokio::spawn(async move {
190 let mut ticker = tokio::time::interval(CLEANUP_INTERVAL);
191 ticker.tick().await;
192 loop {
193 tokio::select! {
194 () = shutdown.cancelled() => break,
195 _ = ticker.tick() => governor.limiter().retain_recent(),
196 }
197 }
198 });
199}
200
201async fn shed_overloaded(_error: BoxError) -> StatusCode {
202 StatusCode::SERVICE_UNAVAILABLE
203}
204
205async fn apply_request_timeout(
206 State(budget): State<TimeoutBudget>,
207 request: Request<Body>,
208 next: Next,
209) -> Response {
210 let limit = match is_streaming_write(&request) {
211 true => budget.extended,
212 false => budget.standard,
213 };
214 match tokio::time::timeout(limit, next.run(request)).await {
215 Ok(response) => response,
216 Err(_) => StatusCode::REQUEST_TIMEOUT.into_response(),
217 }
218}
219
220fn is_streaming_write(request: &Request<Body>) -> bool {
221 let path = request.uri().path();
222 match *request.method() {
223 Method::POST => path.ends_with("/git-receive-pack"),
224 Method::PUT => path.contains("/info/lfs/objects/"),
225 _ => false,
226 }
227}
228
229fn rewrap_body(body: TimeoutBody<axum::body::Body>) -> axum::body::Body {
230 axum::body::Body::new(body)
231}
232
233pub(crate) fn apply(router: Router, layers: GuardLayers) -> Router {
234 let GuardLayers {
235 governor,
236 max_inflight,
237 body_timeout,
238 timeout,
239 } = layers;
240 let rate_limit: GovernorLayer<ProxyAwareIp, NoOpMiddleware, axum::body::Body> =
241 GovernorLayer::new(governor);
242 router
243 .layer(
244 ServiceBuilder::new()
245 .layer(RequestBodyTimeoutLayer::new(body_timeout))
246 .layer(MapRequestBodyLayer::new(rewrap_body)),
247 )
248 .layer(from_fn_with_state(timeout, apply_request_timeout))
249 .layer(
250 ServiceBuilder::new()
251 .layer(HandleErrorLayer::new(shed_overloaded))
252 .layer(LoadShedLayer::new())
253 .layer(GlobalConcurrencyLimitLayer::new(max_inflight)),
254 )
255 .layer(rate_limit)
256}
257
258#[cfg(test)]
259mod tests {
260 use super::*;
261
262 use axum::body::{Body, Bytes};
263 use axum::routing::{get, post};
264 use futures::StreamExt;
265 use http::StatusCode;
266 use tower::ServiceExt;
267
268 fn guards(
269 rate: u32,
270 burst: u32,
271 inflight: u32,
272 request_timeout_ms: u64,
273 body_timeout_ms: u64,
274 proxy_trust: ProxyTrust,
275 ) -> EdgeGuards {
276 guards_with_write(
277 rate,
278 burst,
279 inflight,
280 request_timeout_ms,
281 body_timeout_ms,
282 request_timeout_ms,
283 proxy_trust,
284 )
285 }
286
287 #[allow(clippy::too_many_arguments)]
288 fn guards_with_write(
289 rate: u32,
290 burst: u32,
291 inflight: u32,
292 request_timeout_ms: u64,
293 body_timeout_ms: u64,
294 write_request_timeout_ms: u64,
295 proxy_trust: ProxyTrust,
296 ) -> EdgeGuards {
297 EdgeGuards::new(
298 RequestsPerSecond::new(NonZeroU32::new(rate).unwrap()),
299 BurstSize::new(NonZeroU32::new(burst).unwrap()),
300 MaxInflightRequests::new(NonZeroU32::new(inflight).unwrap()),
301 RequestTimeout::from_millis(NonZeroU64::new(request_timeout_ms).unwrap()),
302 BodyInactivityTimeout::from_millis(NonZeroU64::new(body_timeout_ms).unwrap()),
303 WriteRequestTimeout::from_millis(NonZeroU64::new(write_request_timeout_ms).unwrap()),
304 proxy_trust,
305 )
306 }
307
308 fn forwarded_for() -> http::HeaderName {
309 http::HeaderName::from_static("x-forwarded-for")
310 }
311
312 fn trusting_any_peer() -> ProxyTrust {
313 ProxyTrust::new(Some(forwarded_for()), knot_types::TrustedProxies::default())
314 }
315
316 fn trusting_loopback() -> ProxyTrust {
317 ProxyTrust::new(
318 Some(forwarded_for()),
319 knot_types::TrustedProxies::parse(["127.0.0.1"]).unwrap(),
320 )
321 }
322
323 fn extractor(trust: ProxyTrust) -> ProxyAwareIp {
324 ProxyAwareIp {
325 trust,
326 ignored_header: IgnoredHeaderNotice::default(),
327 }
328 }
329
330 fn guarded_router(router: Router, guards: EdgeGuards) -> Router {
331 apply(router, guards.prepare(&CancellationToken::new()))
332 }
333
334 fn from_peer(request: Request<Body>, host: u8) -> Request<Body> {
335 let mut request = request;
336 request
337 .extensions_mut()
338 .insert(ConnectInfo(SocketAddr::from(([127, 0, 0, host], 47000))));
339 request
340 }
341
342 fn get_request() -> Request<Body> {
343 Request::get("/").body(Body::empty()).unwrap()
344 }
345
346 #[test]
347 fn the_extractor_keys_on_the_trusted_proxy_header_when_configured() {
348 let extractor = extractor(trusting_any_peer());
349 let request = Request::get("/")
350 .header("x-forwarded-for", "203.0.113.7, 198.51.100.4")
351 .extension(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 5000))))
352 .body(())
353 .unwrap();
354 assert_eq!(
355 extractor.extract(&request).unwrap(),
356 "198.51.100.4".parse::<IpAddr>().unwrap(),
357 "the rightmost forwarded entry is the client the proxy appended"
358 );
359 }
360
361 #[test]
362 fn the_extractor_ignores_a_forgeable_header_when_no_proxy_is_trusted() {
363 let extractor = extractor(ProxyTrust::default());
364 let request = Request::get("/")
365 .header("x-forwarded-for", "203.0.113.7")
366 .extension(ConnectInfo(SocketAddr::from(([10, 0, 0, 9], 5000))))
367 .body(())
368 .unwrap();
369 assert_eq!(
370 extractor.extract(&request).unwrap(),
371 "10.0.0.9".parse::<IpAddr>().unwrap(),
372 "with no trusted proxy the socket peer wins over a client-forgeable header"
373 );
374 }
375
376 #[test]
377 fn the_extractor_keys_an_unlisted_peer_on_its_socket_however_it_fills_the_header() {
378 let extractor = extractor(trusting_loopback());
379 let forged = |host| {
380 Request::get("/")
381 .header("x-forwarded-for", "198.51.100.4")
382 .extension(ConnectInfo(SocketAddr::from(([127, 0, 0, host], 5000))))
383 .body(())
384 .unwrap()
385 };
386 assert_eq!(
387 extractor.extract(&forged(1)).unwrap(),
388 "198.51.100.4".parse::<IpAddr>().unwrap(),
389 "the listed proxy relayed this one, so the extractor keys on the header address"
390 );
391 assert_eq!(
392 extractor.extract(&forged(9)).unwrap(),
393 "127.0.0.9".parse::<IpAddr>().unwrap(),
394 "an unlisted peer picked its own token bucket by forging the header"
395 );
396 }
397
398 #[test]
399 fn the_extractor_falls_back_to_the_peer_when_the_trusted_header_is_absent() {
400 let extractor = extractor(trusting_any_peer());
401 let request = Request::get("/")
402 .extension(ConnectInfo(SocketAddr::from(([10, 0, 0, 9], 5000))))
403 .body(())
404 .unwrap();
405 assert_eq!(
406 extractor.extract(&request).unwrap(),
407 "10.0.0.9".parse::<IpAddr>().unwrap()
408 );
409 }
410
411 #[test]
412 fn the_extractor_fails_when_no_peer_can_be_identified() {
413 let extractor = extractor(ProxyTrust::default());
414 let request = Request::get("/").body(()).unwrap();
415 assert!(matches!(
416 extractor.extract(&request),
417 Err(GovernorError::UnableToExtractKey)
418 ));
419 }
420
421 #[test]
422 fn a_missing_connect_info_fails_closed_rather_than_taking_the_header_on_trust() {
423 let listed = extractor(trusting_loopback());
424 let headed = || {
425 Request::get("/")
426 .header("x-forwarded-for", "198.51.100.4")
427 .body(())
428 .unwrap()
429 };
430 assert!(
431 matches!(
432 listed.extract(&headed()),
433 Err(GovernorError::UnableToExtractKey)
434 ),
435 "with no socket to match against the allowlist the extractor identifies no client"
436 );
437 assert_eq!(
438 extractor(trusting_any_peer()).extract(&headed()).unwrap(),
439 "198.51.100.4".parse::<IpAddr>().unwrap(),
440 "an operator who lists no proxy already told the knot to take the header from anyone"
441 );
442 }
443
444 #[test]
445 fn the_first_unlisted_peer_sending_the_header_is_reported_once() {
446 let extractor = extractor(trusting_loopback());
447 let forged = |host| {
448 Request::get("/")
449 .header("x-forwarded-for", "198.51.100.4")
450 .extension(ConnectInfo(SocketAddr::from(([203, 0, 113, host], 5000))))
451 .body(())
452 .unwrap()
453 };
454 extractor.extract(&forged(7)).unwrap();
455 extractor.extract(&forged(9)).unwrap();
456 assert_eq!(
457 extractor.ignored_header.0.get(),
458 Some(&"203.0.113.7".parse::<IpAddr>().unwrap()),
459 "a wrong-family allowlist sends every request down this path, so only the first peer is reported"
460 );
461 }
462
463 #[test]
464 fn a_listed_proxy_and_a_headerless_request_report_nothing() {
465 let listed = extractor(trusting_loopback());
466 listed
467 .extract(
468 &Request::get("/")
469 .header("x-forwarded-for", "198.51.100.4")
470 .extension(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 5000))))
471 .body(())
472 .unwrap(),
473 )
474 .unwrap();
475 listed
476 .extract(
477 &Request::get("/")
478 .extension(ConnectInfo(SocketAddr::from(([203, 0, 113, 7], 5000))))
479 .body(())
480 .unwrap(),
481 )
482 .unwrap();
483 assert_eq!(
484 listed.ignored_header.0.get(),
485 None,
486 "neither a relayed request or a request without the header says anything about the allowlist"
487 );
488 }
489
490 #[tokio::test]
491 async fn a_well_behaved_request_passes_every_guard() {
492 let app = guarded_router(
493 Router::new().route("/", get(|| async { "ok" })),
494 guards(50, 200, 1_024, 60_000, 30_000, ProxyTrust::default()),
495 );
496 let status = app
497 .oneshot(from_peer(get_request(), 1))
498 .await
499 .unwrap()
500 .status();
501 assert_eq!(status, StatusCode::OK);
502 }
503
504 #[tokio::test]
505 async fn a_burst_beyond_the_per_ip_limit_is_rejected_with_429() {
506 let app = guarded_router(
507 Router::new().route("/", get(|| async { "ok" })),
508 guards(1, 2, 1_024, 60_000, 30_000, ProxyTrust::default()),
509 );
510 let first = app
511 .clone()
512 .oneshot(from_peer(get_request(), 7))
513 .await
514 .unwrap()
515 .status();
516 let second = app
517 .clone()
518 .oneshot(from_peer(get_request(), 7))
519 .await
520 .unwrap()
521 .status();
522 let third = app
523 .clone()
524 .oneshot(from_peer(get_request(), 7))
525 .await
526 .unwrap()
527 .status();
528 let other_ip = app
529 .clone()
530 .oneshot(from_peer(get_request(), 8))
531 .await
532 .unwrap()
533 .status();
534 assert_eq!(first, StatusCode::OK);
535 assert_eq!(second, StatusCode::OK);
536 assert_eq!(
537 third,
538 StatusCode::TOO_MANY_REQUESTS,
539 "a third request inside the window exhausts the burst for this IP"
540 );
541 assert_eq!(
542 other_ip,
543 StatusCode::OK,
544 "a different IP keeps its own independent quota"
545 );
546 }
547
548 #[tokio::test]
549 async fn requests_beyond_the_inflight_limit_are_shed_with_503() {
550 let app = guarded_router(
551 Router::new().route(
552 "/slow",
553 get(|| async {
554 tokio::time::sleep(Duration::from_millis(300)).await;
555 "ok"
556 }),
557 ),
558 guards(10_000, 10_000, 1, 60_000, 30_000, ProxyTrust::default()),
559 );
560 let holder = {
561 let app = app.clone();
562 tokio::spawn(async move {
563 let request = from_peer(Request::get("/slow").body(Body::empty()).unwrap(), 1);
564 app.oneshot(request).await.unwrap().status()
565 })
566 };
567 tokio::time::sleep(Duration::from_millis(50)).await;
568 let shed = app
569 .clone()
570 .oneshot(from_peer(
571 Request::get("/slow").body(Body::empty()).unwrap(),
572 2,
573 ))
574 .await
575 .unwrap()
576 .status();
577 assert_eq!(
578 shed,
579 StatusCode::SERVICE_UNAVAILABLE,
580 "with the single inflight slot held, the next request sheds rather than queues"
581 );
582 assert_eq!(holder.await.unwrap(), StatusCode::OK);
583 }
584
585 #[tokio::test]
586 async fn a_request_slower_than_the_timeout_is_cut_with_408() {
587 let app = guarded_router(
588 Router::new().route(
589 "/slow",
590 get(|| async {
591 tokio::time::sleep(Duration::from_millis(500)).await;
592 "ok"
593 }),
594 ),
595 guards(10_000, 10_000, 1_024, 80, 30_000, ProxyTrust::default()),
596 );
597 let status = app
598 .oneshot(from_peer(
599 Request::get("/slow").body(Body::empty()).unwrap(),
600 1,
601 ))
602 .await
603 .unwrap()
604 .status();
605 assert_eq!(status, StatusCode::REQUEST_TIMEOUT);
606 }
607
608 #[tokio::test]
609 async fn a_streaming_write_runs_under_the_extended_budget_while_reads_keep_the_standard_one() {
610 let app = guarded_router(
611 Router::new()
612 .route(
613 "/did/name/git-receive-pack",
614 post(|| async {
615 tokio::time::sleep(Duration::from_millis(200)).await;
616 "ok"
617 }),
618 )
619 .route(
620 "/did/name/git-upload-pack",
621 post(|| async {
622 tokio::time::sleep(Duration::from_millis(200)).await;
623 "ok"
624 }),
625 ),
626 guards_with_write(
627 10_000,
628 10_000,
629 1_024,
630 80,
631 30_000,
632 5_000,
633 ProxyTrust::default(),
634 ),
635 );
636 let push = app
637 .clone()
638 .oneshot(from_peer(
639 Request::post("/did/name/git-receive-pack")
640 .body(Body::empty())
641 .unwrap(),
642 1,
643 ))
644 .await
645 .unwrap()
646 .status();
647 assert_eq!(
648 push,
649 StatusCode::OK,
650 "a push slower than the standard timeout survives on the extended write budget"
651 );
652 let fetch = app
653 .oneshot(from_peer(
654 Request::post("/did/name/git-upload-pack")
655 .body(Body::empty())
656 .unwrap(),
657 1,
658 ))
659 .await
660 .unwrap()
661 .status();
662 assert_eq!(
663 fetch,
664 StatusCode::REQUEST_TIMEOUT,
665 "a non-write request past the standard timeout is still cut"
666 );
667 }
668
669 #[tokio::test]
670 async fn a_stalled_request_body_is_cut_and_never_hangs() {
671 let app = guarded_router(
672 Router::new().route("/upload", post(|_body: Bytes| async { "ok" })),
673 guards(10_000, 10_000, 1_024, 60_000, 80, ProxyTrust::default()),
674 );
675 let body = Body::from_stream(
676 futures::stream::once(async {
677 Ok::<_, std::io::Error>(Bytes::from_static(b"partial"))
678 })
679 .chain(futures::stream::pending::<Result<Bytes, std::io::Error>>()),
680 );
681 let request = from_peer(Request::post("/upload").body(body).unwrap(), 1);
682 let status = tokio::time::timeout(Duration::from_secs(5), app.oneshot(request))
683 .await
684 .expect("the body inactivity timeout must cut a stalled upload instead of hanging")
685 .unwrap()
686 .status();
687 assert_ne!(
688 status,
689 StatusCode::OK,
690 "a body that stalls past the inactivity timeout is never accepted"
691 );
692 }
693}