1use std::collections::HashMap;
4use std::net::{IpAddr, SocketAddr};
5use std::sync::{Arc, Mutex};
6
7use axum::Json;
8use axum::extract::{ConnectInfo, State};
9use axum::http::StatusCode;
10use axum::http::header::{COOKIE, SET_COOKIE};
11use axum::response::{AppendHeaders, IntoResponse, Response};
12use axum::routing::post;
13use serde::{Deserialize, Serialize};
14
15use koan_core::auth;
16use koan_core::db::pool::{Handle, Pool};
17use koan_core::db::queries::auth as auth_queries;
18
19pub(crate) const REFRESH_COOKIE: &str = "koan_refresh";
23const REFRESH_COOKIE_PATH: &str = "/auth";
24const STALE_REFRESH_COOKIE_PATH: &str = "/auth/refresh";
28
29const LOGIN_WINDOW_SECS: u64 = 60;
36const LOGIN_MAX_PER_WINDOW: u32 = 10;
37const TRACKED_IPS_MAX: usize = 4096;
39
40pub struct RateLimiter {
41 windows: Mutex<HashMap<IpAddr, (u64, u32)>>,
42 window_secs: u64,
43 max: u32,
44}
45
46impl Default for RateLimiter {
47 fn default() -> Self {
48 Self::new(LOGIN_WINDOW_SECS, LOGIN_MAX_PER_WINDOW)
49 }
50}
51
52impl RateLimiter {
53 pub fn new(window_secs: u64, max: u32) -> Self {
54 Self {
55 windows: Mutex::default(),
56 window_secs,
57 max,
58 }
59 }
60
61 pub(crate) fn allow(&self, ip: IpAddr) -> bool {
64 let now = koan_core::clock::unix_secs() as u64;
65 let mut windows = self.windows.lock().unwrap_or_else(|e| e.into_inner());
66
67 if windows.len() > TRACKED_IPS_MAX {
68 windows.retain(|_, (start, _)| now.saturating_sub(*start) < self.window_secs);
69 }
70
71 let entry = windows.entry(network(ip)).or_insert((now, 0));
72 if now.saturating_sub(entry.0) >= self.window_secs {
73 *entry = (now, 0);
74 }
75 entry.1 += 1;
76 entry.1 <= self.max
77 }
78}
79
80pub(crate) fn client_ip(request: &axum::extract::Request) -> IpAddr {
94 address(request.extensions(), request.headers())
95}
96
97pub(crate) struct ClientIp(pub IpAddr);
99
100impl<S: Send + Sync> axum::extract::FromRequestParts<S> for ClientIp {
101 type Rejection = std::convert::Infallible;
102
103 async fn from_request_parts(
104 parts: &mut axum::http::request::Parts,
105 _: &S,
106 ) -> Result<Self, Self::Rejection> {
107 Ok(Self(address(&parts.extensions, &parts.headers)))
108 }
109}
110
111fn address(extensions: &axum::http::Extensions, headers: &axum::http::HeaderMap) -> IpAddr {
112 let Some(ConnectInfo(peer)) = extensions.get::<ConnectInfo<SocketAddr>>() else {
113 return IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED);
114 };
115 let peer = peer.ip().to_canonical();
116 if !is_internal(peer) {
117 return peer;
118 }
119 headers
120 .get_all("x-forwarded-for")
121 .iter()
122 .filter_map(|v| v.to_str().ok())
123 .flat_map(|v| v.split(','))
124 .filter_map(|ip| ip.trim().parse::<IpAddr>().ok())
125 .next_back()
126 .map_or(peer, |ip| ip.to_canonical())
127}
128
129pub(crate) fn network(ip: IpAddr) -> IpAddr {
133 match ip.to_canonical() {
134 IpAddr::V4(v4) => IpAddr::V4(v4),
135 IpAddr::V6(v6) => IpAddr::V6((u128::from(v6) & !((1u128 << 64) - 1)).into()),
136 }
137}
138
139fn is_internal(ip: IpAddr) -> bool {
140 match ip.to_canonical() {
141 IpAddr::V4(v4) => v4.is_private() || v4.is_loopback(),
142 IpAddr::V6(v6) => v6.is_loopback() || (v6.segments()[0] & 0xfe00) == 0xfc00,
143 }
144}
145
146#[derive(Clone)]
151pub struct AuthRouteState {
152 pub pool: Arc<Pool>,
153 pub private_pem: Arc<Vec<u8>>,
154 pub public_pem: Arc<Vec<u8>>,
155 pub access_ttl_secs: u64,
156 pub refresh_ttl_secs: u64,
157 pub cookie_secure: bool,
161 pub login_limiter: Arc<RateLimiter>,
162 pub users: Arc<super::password::PasswordVerifier>,
164}
165
166impl AuthRouteState {
167 fn cookie(&self, name: &str, value: &str, path: &str, max_age: u64) -> String {
170 let secure = if self.cookie_secure { "; Secure" } else { "" };
171 format!("{name}={value}; HttpOnly; SameSite=Lax; Path={path}; Max-Age={max_age}{secure}")
172 }
173
174 fn access_cookie(&self, token: &str) -> String {
175 self.cookie("koan_access", token, "/", self.access_ttl_secs)
176 }
177
178 fn refresh_cookie(&self, token: &str) -> String {
179 self.cookie(
180 REFRESH_COOKIE,
181 token,
182 REFRESH_COOKIE_PATH,
183 self.refresh_ttl_secs,
184 )
185 }
186
187 fn stale_refresh_cookie(&self) -> String {
188 self.cookie(REFRESH_COOKIE, "", STALE_REFRESH_COOKIE_PATH, 0)
189 }
190
191 pub(crate) fn session_cookies(
194 &self,
195 access_token: &str,
196 refresh_token: &str,
197 ) -> AppendHeaders<[(axum::http::HeaderName, String); 3]> {
198 AppendHeaders([
199 (SET_COOKIE, self.access_cookie(access_token)),
200 (SET_COOKIE, self.refresh_cookie(refresh_token)),
201 (SET_COOKIE, self.stale_refresh_cookie()),
202 ])
203 }
204
205 pub(crate) fn proxied_cookies(
210 &self,
211 access_token: &str,
212 ) -> AppendHeaders<[(axum::http::HeaderName, String); 3]> {
213 AppendHeaders([
214 (SET_COOKIE, self.access_cookie(access_token)),
215 (
216 SET_COOKIE,
217 self.cookie(REFRESH_COOKIE, "", REFRESH_COOKIE_PATH, 0),
218 ),
219 (SET_COOKIE, self.stale_refresh_cookie()),
220 ])
221 }
222
223 pub(crate) fn cleared_cookies(&self) -> AppendHeaders<[(axum::http::HeaderName, String); 3]> {
225 AppendHeaders([
226 (SET_COOKIE, self.cookie("koan_access", "", "/", 0)),
227 (
228 SET_COOKIE,
229 self.cookie(REFRESH_COOKIE, "", REFRESH_COOKIE_PATH, 0),
230 ),
231 (SET_COOKIE, self.stale_refresh_cookie()),
232 ])
233 }
234}
235
236pub(crate) fn dummy_password_hash() -> &'static str {
239 static HASH: std::sync::OnceLock<String> = std::sync::OnceLock::new();
240 HASH.get_or_init(|| auth::hash_password("koan-dummy-password").unwrap_or_default())
241}
242
243pub(crate) fn refresh_token_from(
246 body: Option<&str>,
247 headers: &axum::http::HeaderMap,
248) -> Option<String> {
249 if let Some(t) = body.filter(|t| !t.is_empty()) {
250 return Some(t.to_owned());
251 }
252 headers
253 .get(COOKIE)
254 .and_then(|v| v.to_str().ok())
255 .and_then(|cookies| {
256 cookies.split(';').find_map(|c| {
257 c.trim()
258 .strip_prefix(&format!("{REFRESH_COOKIE}="))
259 .map(str::to_owned)
260 })
261 })
262}
263
264pub(crate) async fn login_rate_limit(
269 State(state): State<AuthRouteState>,
270 request: axum::extract::Request,
271 next: axum::middleware::Next,
272) -> Response {
273 let ip = client_ip(&request);
274
275 if !state.login_limiter.allow(ip) {
276 return (
277 StatusCode::TOO_MANY_REQUESTS,
278 Json(MessageResponse {
279 message: "too many login attempts".into(),
280 }),
281 )
282 .into_response();
283 }
284 next.run(request).await
285}
286
287pub(crate) async fn rate_limit(
289 State(limiter): State<Arc<RateLimiter>>,
290 request: axum::extract::Request,
291 next: axum::middleware::Next,
292) -> Response {
293 if !limiter.allow(client_ip(&request)) {
294 return (StatusCode::TOO_MANY_REQUESTS, "too many requests").into_response();
295 }
296 next.run(request).await
297}
298
299impl AuthRouteState {
300 fn open_db(&self) -> Result<Handle<'_>, (StatusCode, String)> {
301 self.pool.get().map_err(|e| {
302 log::error!("auth db open error: {}", e);
303 (
304 StatusCode::INTERNAL_SERVER_ERROR,
305 "internal error".to_string(),
306 )
307 })
308 }
309}
310
311#[derive(Deserialize)]
316pub struct LoginRequest {
317 pub username: String,
318 pub password: String,
319}
320
321#[derive(Serialize)]
322pub struct LoginResponse {
323 pub access_token: String,
324 pub refresh_token: String,
325 pub token_type: String,
326 pub expires_in: u64,
327 pub user: UserInfo,
328}
329
330#[derive(Serialize)]
331pub struct UserInfo {
332 pub id: i64,
333 pub username: String,
334 pub role: String,
335}
336
337#[derive(Deserialize, Default)]
338#[serde(default)]
339pub struct RefreshRequest {
340 pub refresh_token: Option<String>,
341}
342
343#[derive(Serialize)]
344pub struct RefreshResponse {
345 pub access_token: String,
346 pub refresh_token: String,
347 pub token_type: String,
348 pub expires_in: u64,
349}
350
351#[derive(Deserialize, Default)]
352#[serde(default)]
353pub struct LogoutRequest {
354 pub refresh_token: Option<String>,
355}
356
357#[derive(Serialize)]
358pub struct MessageResponse {
359 pub message: String,
360}
361
362pub fn auth_router(state: AuthRouteState) -> axum::Router {
367 let router = axum::Router::new()
368 .route(
369 "/auth/login",
370 post(login).layer(axum::middleware::from_fn_with_state(
371 state.clone(),
372 login_rate_limit,
373 )),
374 )
375 .route("/auth/refresh", post(refresh))
376 .route("/auth/logout", post(logout));
377 auth_perimeter(router, AUTH_TIMEOUT, AUTH_BODY_TIMEOUT).with_state(state)
378}
379
380const AUTH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
382
383const AUTH_BODY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
385
386const AUTH_BODY_LIMIT: usize = 16 * 1024;
388
389fn auth_perimeter<S>(
396 router: axum::Router<S>,
397 timeout: std::time::Duration,
398 body_timeout: std::time::Duration,
399) -> axum::Router<S>
400where
401 S: Clone + Send + Sync + 'static,
402{
403 router
404 .layer(tower_http::timeout::TimeoutLayer::with_status_code(
405 StatusCode::REQUEST_TIMEOUT,
406 timeout,
407 ))
408 .layer(
409 tower::ServiceBuilder::new()
410 .layer(axum::error_handling::HandleErrorLayer::new(
411 |_: tower::BoxError| async { (StatusCode::SERVICE_UNAVAILABLE, "busy") },
412 ))
413 .load_shed()
414 .concurrency_limit(2),
415 )
416 .layer(axum::middleware::from_fn(
417 move |request: axum::extract::Request, next: axum::middleware::Next| {
418 buffer_body(request, next, body_timeout)
419 },
420 ))
421}
422
423async fn buffer_body(
425 request: axum::extract::Request,
426 next: axum::middleware::Next,
427 deadline: std::time::Duration,
428) -> Response {
429 use futures_util::StreamExt as _;
430 let (parts, body) = request.into_parts();
431 let read = async {
432 let mut stream = body.into_data_stream();
433 let mut bytes = Vec::new();
434 while let Some(chunk) = stream.next().await {
435 let chunk = chunk.map_err(|_| StatusCode::BAD_REQUEST)?;
436 if bytes.len() + chunk.len() > AUTH_BODY_LIMIT {
437 return Err(StatusCode::PAYLOAD_TOO_LARGE);
438 }
439 bytes.extend_from_slice(&chunk);
440 }
441 Ok(bytes)
442 };
443 match tokio::time::timeout(deadline, read).await {
444 Ok(Ok(bytes)) => {
445 next.run(axum::extract::Request::from_parts(parts, bytes.into()))
446 .await
447 }
448 Ok(Err(status)) => status.into_response(),
449 Err(_) => StatusCode::REQUEST_TIMEOUT.into_response(),
450 }
451}
452
453pub(crate) async fn authenticate(
466 state: &AuthRouteState,
467 username: &str,
468 password: &str,
469 from: IpAddr,
470) -> Result<(auth_queries::UserRow, String, String), Box<Response>> {
471 let state = state.clone();
472 let (username, password) = (username.to_owned(), password.to_owned());
473 tokio::task::spawn_blocking(move || authenticate_blocking(&state, &username, &password, from))
474 .await
475 .unwrap_or_else(|_| Err(Box::new(StatusCode::INTERNAL_SERVER_ERROR.into_response())))
476}
477
478fn authenticate_blocking(
479 state: &AuthRouteState,
480 username: &str,
481 password: &str,
482 from: IpAddr,
483) -> Result<(auth_queries::UserRow, String, String), Box<Response>> {
484 use super::password::Refused;
485 let refused = |status, message: &str| {
486 Box::new(
487 (
488 status,
489 Json(MessageResponse {
490 message: message.into(),
491 }),
492 )
493 .into_response(),
494 )
495 };
496 if state.users.spent(username, from) {
497 return Err(refused(
498 StatusCode::TOO_MANY_REQUESTS,
499 "too many failed sign-ins for this account; try again in a minute",
500 ));
501 }
502 let user = match state.users.verify(username, password) {
503 Ok(user) => {
504 state.users.signed_in(username, from);
505 user
506 }
507 Err(Refused::Wrong) => {
508 state.users.failed(username);
509 return Err(refused(
510 StatusCode::UNAUTHORIZED,
511 "invalid username or password",
512 ));
513 }
514 Err(Refused::Busy) => {
515 return Err(refused(
516 StatusCode::SERVICE_UNAVAILABLE,
517 "busy; try again in a moment",
518 ));
519 }
520 };
521
522 let db = match state.open_db() {
523 Ok(db) => db,
524 Err((status, msg)) => return Err(Box::new((status, msg).into_response())),
525 };
526
527 let access_token = match auth::mint_access_token(
528 &state.private_pem,
529 user.id,
530 &user.username,
531 user.role,
532 state.access_ttl_secs,
533 ) {
534 Ok(t) => t,
535 Err(e) => {
536 log::error!("auth mint token error: {}", e);
537 return Err(Box::new(
538 (StatusCode::INTERNAL_SERVER_ERROR, "token error").into_response(),
539 ));
540 }
541 };
542
543 let refresh_token_id = match auth::random_token() {
544 Ok(t) => t,
545 Err(e) => {
546 log::error!("auth refresh token generation error: {}", e);
547 return Err(Box::new(
548 (StatusCode::INTERNAL_SERVER_ERROR, "token error").into_response(),
549 ));
550 }
551 };
552 let refresh_expires = koan_core::clock::unix_secs() + state.refresh_ttl_secs as i64;
553 if let Err(e) =
554 auth_queries::store_refresh_token(&db.conn, &refresh_token_id, user.id, refresh_expires)
555 {
556 log::error!("auth store refresh token error: {}", e);
557 return Err(Box::new(
558 (StatusCode::INTERNAL_SERVER_ERROR, "token error").into_response(),
559 ));
560 }
561
562 let _ = auth_queries::cleanup_expired_tokens(&db.conn);
564
565 Ok((user, access_token, refresh_token_id))
566}
567
568pub(crate) async fn proxied_access(state: &AuthRouteState, username: &str) -> Option<String> {
573 let state = state.clone();
574 let username = username.to_owned();
575 tokio::task::spawn_blocking(move || {
576 let db = state.open_db().ok()?;
577 let user = auth_queries::get_user_by_username(&db.conn, &username).ok()??;
578 auth::mint_access_token(
579 &state.private_pem,
580 user.id,
581 &user.username,
582 user.role,
583 state.access_ttl_secs,
584 )
585 .map_err(|e| log::error!("auth mint token error: {e}"))
586 .ok()
587 })
588 .await
589 .ok()
590 .flatten()
591}
592
593async fn login(
594 State(state): State<AuthRouteState>,
595 ClientIp(from): ClientIp,
596 Json(req): Json<LoginRequest>,
597) -> Response {
598 let (user, access_token, refresh_token_id) =
599 match authenticate(&state, &req.username, &req.password, from).await {
600 Ok(session) => session,
601 Err(resp) => return *resp,
602 };
603
604 let cookies = state.session_cookies(&access_token, &refresh_token_id);
605
606 let resp = LoginResponse {
607 access_token,
608 refresh_token: refresh_token_id,
611 token_type: "Bearer".into(),
612 expires_in: state.access_ttl_secs,
613 user: UserInfo {
614 id: user.id,
615 username: user.username,
616 role: user.role.as_str().into(),
617 },
618 };
619
620 (StatusCode::OK, cookies, Json(resp)).into_response()
621}
622
623const REPLAY_GRACE_SECS: i64 = 30;
626
627pub(crate) fn rotate(
637 state: &AuthRouteState,
638 supplied: &str,
639 from: Option<IpAddr>,
640) -> Result<(String, String), Box<Response>> {
641 let db = match state.open_db() {
642 Ok(db) => db,
643 Err((status, msg)) => return Err(Box::new((status, msg).into_response())),
644 };
645
646 let token = match auth_queries::consume_refresh_token(&db.conn, supplied) {
649 Ok(Some(t)) => t,
650 Ok(None) => {
651 match auth_queries::revoke_replayed_grant(&db.conn, supplied, REPLAY_GRACE_SECS) {
652 Ok(0) => {}
653 Ok(n) => log::warn!("a spent OAuth refresh token came back: revoked {n} tokens"),
654 Err(e) => log::error!("auth replay check error: {e}"),
655 }
656 return Err(Box::new(
657 (
658 StatusCode::UNAUTHORIZED,
659 Json(MessageResponse {
660 message: "invalid or expired refresh token".into(),
661 }),
662 )
663 .into_response(),
664 ));
665 }
666 Err(e) => {
667 log::error!("auth refresh db error: {}", e);
668 return Err(Box::new(
669 (StatusCode::INTERNAL_SERVER_ERROR, "internal error").into_response(),
670 ));
671 }
672 };
673
674 let user = match auth_queries::get_user_by_id(&db.conn, token.user_id) {
675 Ok(Some(u)) => u,
676 Ok(None) => {
677 return Err(Box::new(
678 (
679 StatusCode::UNAUTHORIZED,
680 Json(MessageResponse {
681 message: "user not found".into(),
682 }),
683 )
684 .into_response(),
685 ));
686 }
687 Err(e) => {
688 log::error!("auth refresh user lookup error: {}", e);
689 return Err(Box::new(
690 (StatusCode::INTERNAL_SERVER_ERROR, "internal error").into_response(),
691 ));
692 }
693 };
694
695 let access_token = match auth::mint_scoped_token(
697 &state.private_pem,
698 user.id,
699 &user.username,
700 user.role,
701 state.access_ttl_secs,
702 token.grant.as_ref().map(|_| auth::MCP_SCOPE),
703 ) {
704 Ok(t) => t,
705 Err(e) => {
706 log::error!("auth mint token error: {}", e);
707 return Err(Box::new(
708 (StatusCode::INTERNAL_SERVER_ERROR, "token error").into_response(),
709 ));
710 }
711 };
712
713 let new_refresh_id = match auth::random_token() {
714 Ok(t) => t,
715 Err(e) => {
716 log::error!("auth refresh token generation error: {}", e);
717 return Err(Box::new(
718 (StatusCode::INTERNAL_SERVER_ERROR, "token error").into_response(),
719 ));
720 }
721 };
722 let refresh_expires = koan_core::clock::unix_secs() + state.refresh_ttl_secs as i64;
723 if let Err(e) = auth_queries::store_grant_token(
724 &db.conn,
725 &new_refresh_id,
726 user.id,
727 refresh_expires,
728 token.grant.as_ref(),
729 ) {
730 log::error!("auth store refresh token error: {}", e);
731 return Err(Box::new(
732 (StatusCode::INTERNAL_SERVER_ERROR, "token error").into_response(),
733 ));
734 }
735
736 if let Some(from) = from {
737 state.users.signed_in(&user.username, from);
738 }
739 Ok((access_token, new_refresh_id))
740}
741
742async fn refresh(
743 State(state): State<AuthRouteState>,
744 ClientIp(from): ClientIp,
745 headers: axum::http::HeaderMap,
746 body: Option<Json<RefreshRequest>>,
747) -> Response {
748 let supplied = body.and_then(|Json(req)| req.refresh_token);
749 let Some(supplied) = refresh_token_from(supplied.as_deref(), &headers) else {
750 return (
751 StatusCode::UNAUTHORIZED,
752 Json(MessageResponse {
753 message: "missing refresh token".into(),
754 }),
755 )
756 .into_response();
757 };
758
759 let rotating = state.clone();
760 let rotated = tokio::task::spawn_blocking(move || rotate(&rotating, &supplied, Some(from)))
761 .await
762 .unwrap_or_else(|_| Err(Box::new(StatusCode::INTERNAL_SERVER_ERROR.into_response())));
763 let (access_token, new_refresh_id) = match rotated {
764 Ok(pair) => pair,
765 Err(resp) => return *resp,
766 };
767
768 let cookies = state.session_cookies(&access_token, &new_refresh_id);
769
770 let resp = RefreshResponse {
771 access_token,
772 refresh_token: new_refresh_id,
773 token_type: "Bearer".into(),
774 expires_in: state.access_ttl_secs,
775 };
776
777 (StatusCode::OK, cookies, Json(resp)).into_response()
778}
779
780async fn logout(
781 State(state): State<AuthRouteState>,
782 headers: axum::http::HeaderMap,
783 body: Option<Json<LogoutRequest>>,
784) -> Response {
785 let supplied = body.and_then(|Json(req)| req.refresh_token);
786 let token = refresh_token_from(supplied.as_deref(), &headers);
787 let revoking = state.clone();
788 let revoked = tokio::task::spawn_blocking(move || {
789 let db = revoking.open_db()?;
790 if let Some(token) = token {
791 let _ = auth_queries::revoke_refresh_token(&db.conn, &token);
792 }
793 Ok(())
794 })
795 .await
796 .unwrap_or_else(|_| {
797 Err((
798 StatusCode::INTERNAL_SERVER_ERROR,
799 "internal error".to_string(),
800 ))
801 });
802 if let Err((status, msg)) = revoked {
803 return (status, msg).into_response();
804 }
805
806 let cookies = state.cleared_cookies();
807
808 (
809 StatusCode::OK,
810 cookies,
811 Json(MessageResponse {
812 message: "logged out".into(),
813 }),
814 )
815 .into_response()
816}
817
818#[cfg(test)]
823mod tests {
824 use super::*;
825
826 #[test]
827 fn login_limiter_caps_a_single_ip() {
828 let limiter = RateLimiter::default();
829 let ip: IpAddr = "10.0.0.5".parse().unwrap();
830 for _ in 0..LOGIN_MAX_PER_WINDOW {
831 assert!(limiter.allow(ip));
832 }
833 assert!(!limiter.allow(ip));
834
835 assert!(limiter.allow("10.0.0.6".parse().unwrap()));
837 }
838
839 #[test]
840 fn an_ipv6_client_is_limited_by_its_64() {
841 let limiter = RateLimiter::default();
842 for n in 0..LOGIN_MAX_PER_WINDOW {
843 assert!(limiter.allow(format!("2001:db8:1:2::{n:x}").parse().unwrap()));
844 }
845 assert!(!limiter.allow("2001:db8:1:2:ffff::1".parse().unwrap()));
846 assert!(limiter.allow("2001:db8:1:3::1".parse().unwrap()));
847 }
848
849 #[test]
850 fn a_mapped_ipv4_address_is_the_ipv4_address() {
851 let r = request_from("[::ffff:198.51.100.4]", None);
852 assert_eq!(client_ip(&r), "198.51.100.4".parse::<IpAddr>().unwrap());
853 let r = request_from("[::ffff:10.42.0.7]", Some("::ffff:203.0.113.9"));
855 assert_eq!(client_ip(&r), "203.0.113.9".parse::<IpAddr>().unwrap());
856 }
857
858 fn request_from(peer: &str, forwarded: Option<&str>) -> axum::extract::Request {
859 let mut request = axum::http::Request::new(axum::body::Body::empty());
860 request.extensions_mut().insert(ConnectInfo(
861 format!("{peer}:1234").parse::<SocketAddr>().unwrap(),
862 ));
863 if let Some(f) = forwarded {
864 request
865 .headers_mut()
866 .insert("x-forwarded-for", f.parse().unwrap());
867 }
868 request
869 }
870
871 #[test]
872 fn client_ip_believes_only_an_internal_proxy() {
873 let r = request_from("10.42.0.7", Some("6.6.6.6, 203.0.113.9"));
875 assert_eq!(client_ip(&r), "203.0.113.9".parse::<IpAddr>().unwrap());
876 let r = request_from("198.51.100.4", Some("10.0.0.1"));
878 assert_eq!(client_ip(&r), "198.51.100.4".parse::<IpAddr>().unwrap());
879 let r = request_from("10.42.0.7", None);
881 assert_eq!(client_ip(&r), "10.42.0.7".parse::<IpAddr>().unwrap());
882 let mut r = axum::http::Request::new(axum::body::Body::empty());
884 r.headers_mut()
885 .insert("x-forwarded-for", "203.0.113.9".parse().unwrap());
886 assert_eq!(client_ip(&r), IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED));
887 }
888
889 #[tokio::test]
890 async fn stalled_bodies_time_out_without_holding_a_permit() {
891 use tower::ServiceExt as _;
892 async fn read(_: axum::body::Bytes) -> StatusCode {
893 StatusCode::OK
894 }
895 let app = auth_perimeter(
898 axum::Router::new().route("/auth/login", post(read)),
899 std::time::Duration::from_secs(10),
900 std::time::Duration::from_millis(200),
901 )
902 .with_state(());
903 let stalled = || {
904 axum::http::Request::post("/auth/login")
905 .body(axum::body::Body::from_stream(tokio_stream::pending::<
906 Result<axum::body::Bytes, std::io::Error>,
907 >()))
908 .unwrap()
909 };
910 let held: Vec<_> = (0..4)
911 .map(|_| tokio::spawn(app.clone().oneshot(stalled())))
912 .collect();
913 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
914 let ok = axum::http::Request::post("/auth/login")
916 .body(axum::body::Body::from("{}"))
917 .unwrap();
918 assert_eq!(
919 app.clone().oneshot(ok).await.unwrap().status(),
920 StatusCode::OK
921 );
922 for h in held {
923 assert_eq!(
924 h.await.unwrap().unwrap().status(),
925 StatusCode::REQUEST_TIMEOUT
926 );
927 }
928 let big = axum::http::Request::post("/auth/login")
929 .body(axum::body::Body::from(vec![b' '; AUTH_BODY_LIMIT + 1]))
930 .unwrap();
931 assert_eq!(
932 app.clone().oneshot(big).await.unwrap().status(),
933 StatusCode::PAYLOAD_TOO_LARGE
934 );
935 let broken = axum::http::Request::post("/auth/login")
936 .body(axum::body::Body::from_stream(tokio_stream::once(Err::<
937 axum::body::Bytes,
938 _,
939 >(
940 std::io::Error::other("reset"),
941 ))))
942 .unwrap();
943 assert_eq!(
944 app.oneshot(broken).await.unwrap().status(),
945 StatusCode::BAD_REQUEST
946 );
947 }
948
949 #[tokio::test]
950 async fn slow_work_is_shed_beyond_the_permits() {
951 use tower::ServiceExt as _;
952 async fn slow(_: axum::body::Bytes) -> StatusCode {
953 tokio::time::sleep(std::time::Duration::from_millis(300)).await;
954 StatusCode::OK
955 }
956 let app = auth_perimeter(
957 axum::Router::new().route("/auth/login", post(slow)),
958 std::time::Duration::from_millis(100),
959 std::time::Duration::from_secs(5),
960 )
961 .with_state(());
962 let req = || {
963 axum::http::Request::post("/auth/login")
964 .body(axum::body::Body::from("{}"))
965 .unwrap()
966 };
967 let held: Vec<_> = (0..2)
968 .map(|_| tokio::spawn(app.clone().oneshot(req())))
969 .collect();
970 tokio::time::sleep(std::time::Duration::from_millis(30)).await;
971 assert_eq!(
972 app.clone().oneshot(req()).await.unwrap().status(),
973 StatusCode::SERVICE_UNAVAILABLE
974 );
975 for h in held {
976 assert_eq!(
977 h.await.unwrap().unwrap().status(),
978 StatusCode::REQUEST_TIMEOUT
979 );
980 }
981 }
982
983 #[test]
984 fn refresh_token_falls_back_to_the_cookie() {
985 let mut headers = axum::http::HeaderMap::new();
986 headers.insert(
987 COOKIE,
988 format!("a=1; {REFRESH_COOKIE}=from-cookie; b=2")
989 .parse()
990 .unwrap(),
991 );
992
993 assert_eq!(
994 refresh_token_from(None, &headers).as_deref(),
995 Some("from-cookie")
996 );
997 assert_eq!(
998 refresh_token_from(Some("from-body"), &headers).as_deref(),
999 Some("from-body")
1000 );
1001 assert_eq!(
1002 refresh_token_from(None, &axum::http::HeaderMap::new()),
1003 None
1004 );
1005 }
1006
1007 #[test]
1008 fn cookies_are_lax_and_only_secure_when_tls_is_in_play() {
1009 let state = |cookie_secure| AuthRouteState {
1010 pool: Arc::new(Pool::new("/nonexistent".into())),
1011 private_pem: Arc::new(Vec::new()),
1012 public_pem: Arc::new(Vec::new()),
1013 access_ttl_secs: 900,
1014 refresh_ttl_secs: 60,
1015 cookie_secure,
1016 login_limiter: Arc::new(RateLimiter::default()),
1017 users: Arc::new(crate::auth::password::PasswordVerifier::new(Arc::new(
1018 Pool::new("/nonexistent/koan.db".into()),
1019 ))),
1020 };
1021
1022 let plain = state(false).access_cookie("tok");
1023 assert!(plain.contains("SameSite=Lax"));
1024 assert!(plain.contains("HttpOnly"));
1025 assert!(!plain.contains("Secure"));
1026
1027 assert!(state(true).access_cookie("tok").contains("; Secure"));
1028
1029 let resp = (StatusCode::OK, state(false).session_cookies("a", "r")).into_response();
1031 let set: Vec<_> = resp.headers().get_all(SET_COOKIE).iter().collect();
1032 assert_eq!(set.len(), 3);
1033 assert_eq!(
1034 (StatusCode::OK, state(false).cleared_cookies())
1035 .into_response()
1036 .headers()
1037 .get_all(SET_COOKIE)
1038 .iter()
1039 .count(),
1040 3
1041 );
1042
1043 let refresh = state(false).refresh_cookie("tok");
1045 assert!(refresh.contains("Path=/auth;"));
1046 assert!(refresh.contains("HttpOnly"));
1047 }
1048}