From 0a4059d0da1efe9c9e9df07e372fc195284b5985 Mon Sep 17 00:00:00 2001 From: Alex Bates Date: Mon, 12 Oct 2026 03:37:19 +0100 Subject: [PATCH] renew a session once however many requests arrive, and only when it's about to expire --- src/oauth/refresh.rs | 134 ++++++++++++++++++++++++++++++++++++++----- src/state.rs | 4 ++ 2 files changed, 123 insertions(+), 15 deletions(-) diff --git a/src/oauth/refresh.rs b/src/oauth/refresh.rs index 767ed9f..e98adcc 100644 --- a/src/oauth/refresh.rs +++ b/src/oauth/refresh.rs @@ -11,6 +11,9 @@ use axum::response::{IntoResponse, Response}; use axum::routing::post; use axum::Router; use chrono::{Duration, Utc}; +use std::collections::HashMap; +use std::sync::{Arc, Mutex, PoisonError}; +use std::time::{Duration as StdDuration, Instant}; use crate::error::AppError; use crate::oauth::session::{ @@ -19,6 +22,75 @@ use crate::oauth::session::{ }; use crate::state::AppState; +/// How close to expiry an access token is when it's renewed. The PDS gives +/// them short lives, so renewing early would be renewing on nearly every +/// request. +const RENEW_WITHIN: Duration = Duration::seconds(30); + +/// How long a renewal is remembered, for requests that were already on their +/// way with the cookie it replaced. +const REMEMBER: StdDuration = StdDuration::from_secs(10 * 60); + +/// The session a renewal made, and the cookie that carries it. +type Renewed = (SessionCookie, String); + +/// What came of renewing, and when: `None` if it failed, which is also +/// remembered, so a refresh token that won't work isn't tried on every request. +type Outcome = Option<(Instant, Option)>; + +/// Sessions being renewed, or just renewed, by the refresh token each began +/// with. +/// +/// A page loads with several requests at once, all with the same cookie, and +/// every one of them would renew it. The PDS takes a refresh token used twice +/// for a stolen one, and ends the session, so they share one renewal instead. +#[derive(Default)] +pub struct Renewals(Mutex>>>); + +impl Renewals { + fn slot(&self, refresh_token: &str) -> Arc> { + let mut slots = self.0.lock().unwrap_or_else(PoisonError::into_inner); + // Those long done with, and not being waited on, are forgotten. + if slots.len() > 1000 { + slots.retain(|_, slot| { + Arc::strong_count(slot) > 1 + || slot.try_lock().map_or(true, |done| { + done.as_ref().is_some_and(|(at, _)| at.elapsed() < REMEMBER) + }) + }); + } + slots.entry(refresh_token.to_string()).or_default().clone() + } +} + +/// Renews `session` if its access token is about to expire, once however +/// many requests ask, returning the session and cookie it was renewed to. +async fn renew(state: &AppState, session: SessionCookie) -> Option { + if !session.expires_within(RENEW_WITHIN) { + return None; + } + let refresh_token = session.refresh_token.clone()?; + let slot = state.renewals.slot(&refresh_token); + let mut done = slot.lock().await; + if let Some((at, outcome)) = done.as_ref() { + if at.elapsed() < REMEMBER { + return outcome.clone(); + } + } + + let did = session.did.clone(); + let outcome = match try_refresh_session(state, session, RENEW_WITHIN).await { + Ok((session, Some(set_cookie))) => Some((session, set_cookie)), + Ok((_, None)) => None, + Err(err) => { + eprintln!("failed to renew {did}'s session: {err:?}"); + None + } + }; + *done = Some((Instant::now(), outcome.clone())); + outcome +} + /// Renews a session whose access token is about to expire, or has, before the /// request is handled, so a handler never goes to the PDS with a token it will /// refuse. The renewed cookie goes back with the response, and this request's @@ -32,22 +104,17 @@ pub async fn renew_expiring_session( next: Next, ) -> Response { let mut renewed = None; - // Nothing under /assets, or in a style or script's request, needs one. - if !request.uri().path().starts_with("/assets") { + // Nothing a browser fetches under these is for someone signed in. + let path = request.uri().path(); + if !path.starts_with("/assets") && !path.starts_with("/blob/") { if let Some(session) = get_session_from_headers(&state.secrets.cookie_secret, request.headers()) { - let did = session.did.clone(); - match try_refresh_session(&state, session).await { - Ok((session, Some(set_cookie))) => { - if let Ok(value) = encode_session_cookie(&state.secrets.cookie_secret, &session) - { - replace_cookie(request.headers_mut(), SESSION_COOKIE_NAME, &value); - } - renewed = HeaderValue::from_str(&set_cookie).ok(); + if let Some((session, set_cookie)) = renew(&state, session).await { + if let Ok(value) = encode_session_cookie(&state.secrets.cookie_secret, &session) { + replace_cookie(request.headers_mut(), SESSION_COOKIE_NAME, &value); } - Ok((_, None)) => {} - Err(err) => eprintln!("failed to renew {did}'s session: {err:?}"), + renewed = HeaderValue::from_str(&set_cookie).ok(); } } } @@ -87,7 +154,7 @@ pub async fn refresh( let session = get_session_from_headers(&state.secrets.cookie_secret, &headers) .ok_or(AppError::Unauthorized)?; - let (_session, set_cookie) = try_refresh_session(&state, session).await?; + let (_session, set_cookie) = try_refresh_session(&state, session, RENEW_WITHIN).await?; match set_cookie { Some(value) => { @@ -101,14 +168,15 @@ pub async fn refresh( } } -/// Refresh `session` in place if it's within 5 minutes of expiry, returning +/// Refresh `session` in place if it's within `within` of expiry, returning /// the (possibly updated) session and an optional `Set-Cookie` header value /// for the caller to attach. pub async fn try_refresh_session( state: &AppState, session: SessionCookie, + within: Duration, ) -> Result<(SessionCookie, Option), AppError> { - if !session.expires_within(Duration::minutes(5)) { + if !session.expires_within(within) { return Ok((session, None)); } @@ -225,6 +293,42 @@ mod tests { ); } + #[tokio::test] + async fn requests_with_one_refresh_token_share_one_renewal() { + let renewals = Renewals::default(); + + let first = renewals.slot("token-a"); + let same = renewals.slot("token-a"); + let other = renewals.slot("token-b"); + + assert!(Arc::ptr_eq(&first, &same)); + assert!(!Arc::ptr_eq(&first, &other)); + } + + #[tokio::test] + async fn a_request_waits_for_the_renewal_in_progress() { + let renewals = Arc::new(Renewals::default()); + let slot = renewals.slot("token"); + let mut renewing = slot.lock().await; + + // Another request for the same session, which has to wait its turn. + let waiting = { + let renewals = renewals.clone(); + tokio::spawn(async move { + let slot = renewals.slot("token"); + let done = slot.lock().await; + done.as_ref().map(|(_, outcome)| outcome.is_none()) + }) + }; + tokio::task::yield_now().await; + assert!(!waiting.is_finished()); + + // It finds what the first came to, rather than renewing again. + *renewing = Some((Instant::now(), None)); + drop(renewing); + assert_eq!(waiting.await.unwrap(), Some(true)); + } + #[test] fn leaves_a_request_without_cookies_alone() { let mut headers = HeaderMap::new(); diff --git a/src/state.rs b/src/state.rs index 5bb13f9..101cf69 100644 --- a/src/state.rs +++ b/src/state.rs @@ -53,6 +53,9 @@ pub struct Inner { pub downloads_seen: stats::Seen, /// Who viewed what today, likewise. pub views_seen: stats::Seen, + /// Sessions being renewed, or just renewed, so requests that arrive + /// together share one renewal. + pub renewals: crate::oauth::refresh::Renewals, } impl AppState { @@ -107,6 +110,7 @@ impl AppState { signup, downloads_seen: stats::Seen::default(), views_seen: stats::Seen::default(), + renewals: Default::default(), }))) } } -- knotmirror 0.0.1