From 7586e695d9fb0829b00232116de93e2755479bda Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 3 Nov 2026 09:10:16 +0000 Subject: [PATCH] Device login: concurrent polls of a code no longer hold the pool A poll locked its request FOR UPDATE and, for an approved one, asked the pool for a second connection to read the user and the scope's repositories, while every other poll of the same code held a connection waiting for that lock. A handful of concurrent polls took the whole pool, and every request of the instance waited for the acquire timeout and failed, until the request expired. The poll now locks the request FOR UPDATE NOWAIT and answers slow_down at once when another poll holds it, and collect reads the user and the repositories on the transaction's own connection, so a poll never waits for a connection or a lock while it holds one. find_repository_by_name takes any executor for that. The test polls one approved code six times at once on a pool of two connections: one poll gets the token, every poll answers within seconds, and a request beside them is served (before: 10 s, the acquire timeout). Fixes #794 Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_012e6DJackg4JuAo1ir7BJjS --- ...01d608f6220dca563034db86142c75332425.json} | 4 +- crates/db/src/device_auth_dao.rs | 69 ++++++++++----- crates/db/src/repository_dao.rs | 14 ++-- crates/server/src/auth/device.rs | 49 +++++++---- crates/server/tests/common/mod.rs | 17 ++++ crates/server/tests/device_auth.rs | 83 ++++++++++++++++++- docs/deployment.md | 2 +- 7 files changed, 192 insertions(+), 46 deletions(-) rename .sqlx/{query-d308dc16d28a0e5a6629c141569ac745dc25768f6de2bc8c0b424993e3397dc8.json => query-561cc809dfe2e381d20b3cd0949601d608f6220dca563034db86142c75332425.json} (95%) diff --git a/.sqlx/query-d308dc16d28a0e5a6629c141569ac745dc25768f6de2bc8c0b424993e3397dc8.json b/.sqlx/query-561cc809dfe2e381d20b3cd0949601d608f6220dca563034db86142c75332425.json similarity index 95% rename from .sqlx/query-d308dc16d28a0e5a6629c141569ac745dc25768f6de2bc8c0b424993e3397dc8.json rename to .sqlx/query-561cc809dfe2e381d20b3cd0949601d608f6220dca563034db86142c75332425.json index f6c995d9..50ec6327 100644 --- a/.sqlx/query-d308dc16d28a0e5a6629c141569ac745dc25768f6de2bc8c0b424993e3397dc8.json +++ b/.sqlx/query-561cc809dfe2e381d20b3cd0949601d608f6220dca563034db86142c75332425.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n SELECT id, client_id, scope, status, user_id, interval_seconds,\n expires_at <= now() AS \"expired!\",\n COALESCE(\n last_polled_at > now() - make_interval(secs => interval_seconds),\n false\n ) AS \"too_fast!\"\n FROM device_authorization\n WHERE device_code_hash = $1\n FOR UPDATE\n ", + "query": "\n SELECT id, client_id, scope, status, user_id, interval_seconds,\n expires_at <= now() AS \"expired!\",\n COALESCE(\n last_polled_at > now() - make_interval(secs => interval_seconds),\n false\n ) AS \"too_fast!\"\n FROM device_authorization\n WHERE device_code_hash = $1\n FOR UPDATE NOWAIT\n ", "describe": { "columns": [ { @@ -98,5 +98,5 @@ null ] }, - "hash": "d308dc16d28a0e5a6629c141569ac745dc25768f6de2bc8c0b424993e3397dc8" + "hash": "561cc809dfe2e381d20b3cd0949601d608f6220dca563034db86142c75332425" } diff --git a/crates/db/src/device_auth_dao.rs b/crates/db/src/device_auth_dao.rs index b0f6b90a..071490c8 100644 --- a/crates/db/src/device_auth_dao.rs +++ b/crates/db/src/device_auth_dao.rs @@ -125,13 +125,34 @@ pub struct PolledRequest { pub too_fast: bool, } +/// What a poll finds when it locks its request ([`lock_for_poll`]). +#[derive(Debug, Clone)] +pub enum PollLock { + /// The request, locked for the poll's transaction. + Locked(PolledRequest), + /// Another transaction holds the request: a poll of the same code is + /// being answered (or the user is deciding it right now). + Busy, + /// No request has the device code. + Unknown, +} + +/// SQLSTATE `lock_not_available`: `FOR UPDATE NOWAIT` found the row locked. +const LOCK_NOT_AVAILABLE: &str = "55P03"; + /// Find and lock the request whose device code hashes to `device_code_hash` -/// in the transaction on `conn` (`SELECT … FOR UPDATE`, so two polls of one -/// code are answered one after the other). The clock is the database's. +/// in the transaction on `conn` (`SELECT … FOR UPDATE NOWAIT`). A request +/// another transaction holds is [`PollLock::Busy`] at once, never waited +/// for: a poll waiting on the lock would hold its connection meanwhile, and +/// enough of them would take the whole pool (#794). The clock is the +/// database's. +/// +/// A `Busy` answer leaves the transaction aborted (Postgres refused the +/// statement): the caller drops it. pub async fn lock_for_poll( conn: &mut PgConnection, device_code_hash: &str, -) -> sqlx::Result> { +) -> sqlx::Result { let row = sqlx::query!( r#" SELECT id, client_id, scope, status, user_id, interval_seconds, @@ -142,29 +163,35 @@ pub async fn lock_for_poll( ) AS "too_fast!" FROM device_authorization WHERE device_code_hash = $1 - FOR UPDATE + FOR UPDATE NOWAIT "#, device_code_hash, ) .fetch_optional(&mut *conn) - .await?; - Ok(row.map(|row| { - let state = match (row.status.as_str(), row.user_id) { - _ if row.expired => PollState::Expired, - ("pending", _) => PollState::Pending, - ("approved", Some(user_id)) => PollState::Approved { user_id }, - ("denied", _) => PollState::Denied, - // `collected`, or an approval without its user (deleted). - _ => PollState::Expired, - }; - PolledRequest { - id: row.id, - client_id: row.client_id, - scope: row.scope, - state, - interval_seconds: row.interval_seconds, - too_fast: row.too_fast, + .await; + let row = match row { + Ok(Some(row)) => row, + Ok(None) => return Ok(PollLock::Unknown), + Err(sqlx::Error::Database(db)) if db.code().as_deref() == Some(LOCK_NOT_AVAILABLE) => { + return Ok(PollLock::Busy) } + Err(e) => return Err(e), + }; + let state = match (row.status.as_str(), row.user_id) { + _ if row.expired => PollState::Expired, + ("pending", _) => PollState::Pending, + ("approved", Some(user_id)) => PollState::Approved { user_id }, + ("denied", _) => PollState::Denied, + // `collected`, or an approval without its user (deleted). + _ => PollState::Expired, + }; + Ok(PollLock::Locked(PolledRequest { + id: row.id, + client_id: row.client_id, + scope: row.scope, + state, + interval_seconds: row.interval_seconds, + too_fast: row.too_fast, })) } diff --git a/crates/db/src/repository_dao.rs b/crates/db/src/repository_dao.rs index e92af728..1a64b18b 100644 --- a/crates/db/src/repository_dao.rs +++ b/crates/db/src/repository_dao.rs @@ -101,11 +101,15 @@ where .await } -/// Find a repository by its unique name. -pub async fn find_repository_by_name( - pool: &PgPool, +/// Find a repository by its unique name, on a pool or on a connection (in a +/// transaction). +pub async fn find_repository_by_name<'e, E>( + executor: E, name: &str, -) -> sqlx::Result> { +) -> sqlx::Result> +where + E: sqlx::PgExecutor<'e>, +{ sqlx::query_as!( Repository, r#" @@ -119,7 +123,7 @@ pub async fn find_repository_by_name( "#, name, ) - .fetch_optional(pool) + .fetch_optional(executor) .await } diff --git a/crates/server/src/auth/device.rs b/crates/server/src/auth/device.rs index bb571022..31a9a940 100644 --- a/crates/server/src/auth/device.rs +++ b/crates/server/src/auth/device.rs @@ -11,9 +11,9 @@ //! * `POST /auth/device/token` (§3.4–3.5), with `grant_type //! urn:ietf:params:oauth:grant-type:device_code`, answers //! `authorization_pending` until the user decides, `slow_down` (and five -//! more seconds of interval) when polled sooner than the interval, -//! `access_denied`, `expired_token`, or, once, the token. Errors are RFC -//! 6749 §5.2 JSON. +//! more seconds of interval) when polled sooner than the interval or +//! while another poll of the same code is answered, `access_denied`, +//! `expired_token`, or, once, the token. Errors are RFC 6749 §5.2 JSON. //! * `GET /.well-known/oauth-authorization-server` (RFC 8414) names both //! endpoints, so a client finds everything from the server URL alone. //! * `POST /auth/device/lookup` and `/auth/device/decide` serve the `/device` @@ -58,7 +58,7 @@ use forge_core::auth::token::{generate_token_secret, hash_token}; use forge_core::auth::{AttemptRefused, LoginAttempt, LoginThrottle}; use forge_core::authz::{Level, TokenScope}; use forge_db::device_auth_dao::{ - self, DeviceAuthorization, NewDeviceAuthorization, PollState, PolledRequest, + self, DeviceAuthorization, NewDeviceAuthorization, PollLock, PollState, PolledRequest, }; use forge_db::{repository_dao, token_dao, user_dao}; use serde::{Deserialize, Serialize}; @@ -235,15 +235,15 @@ impl DeviceScope { })) } - /// The token scope: the repositories by id. A name no repository has - /// reaches nothing (a scope only ever narrows). - async fn token_scope(&self, state: &AppState) -> sqlx::Result { + /// The token scope: the repositories by id, read on `conn`. A name no + /// repository has reaches nothing (a scope only ever narrows). + async fn token_scope(&self, conn: &mut sqlx::PgConnection) -> sqlx::Result { let repositories = if self.repositories.is_empty() { None } else { let mut ids = Vec::with_capacity(self.repositories.len()); for name in &self.repositories { - match repository_dao::find_repository_by_name(&state.pool, name).await? { + match repository_dao::find_repository_by_name(&mut *conn, name).await? { Some(repo) => ids.push(repo.id), None => tracing::info!( repository = %name, @@ -554,18 +554,32 @@ async fn device_token( } /// Answer a poll of `device_code` by `client_id`. The request's row is -/// locked for the answer, so two polls of one code are answered one after -/// the other, and the token is minted in the transaction that marks the -/// request collected: it is answered once, and a failure leaves the -/// approval to collect on the next poll. +/// locked for the answer, and the token is minted in the transaction that +/// marks the request collected: it is answered once, and a failure leaves +/// the approval to collect on the next poll. +/// +/// A poll never waits for a database connection or a lock while it holds +/// one (#794): everything it reads runs on its transaction's connection, +/// and a poll that finds the request locked by another poll of the same +/// code answers `slow_down` at once, rather than wait for it with a +/// connection taken from the pool. async fn poll(state: &AppState, client_id: &str, device_code: &str) -> Response { let mut tx = match forge_db::begin(&state.pool).await { Ok(tx) => tx, Err(e) => return oauth_unavailable(&e), }; let request = match device_auth_dao::lock_for_poll(&mut tx, &hash_token(device_code)).await { - Ok(Some(request)) => request, - Ok(None) => return bad_request("invalid_grant", "unknown device_code"), + Ok(PollLock::Locked(request)) => request, + Ok(PollLock::Unknown) => return bad_request("invalid_grant", "unknown device_code"), + // A client polls once at a time: a second poll while the first is + // answered came sooner than the interval. + Ok(PollLock::Busy) => { + return bad_request( + "slow_down", + "polled again while a poll of this device_code is answered; wait five \ + seconds longer from now on", + ) + } Err(e) => return oauth_unavailable(&e), }; if request.client_id != client_id { @@ -611,13 +625,16 @@ async fn poll(state: &AppState, client_id: &str, device_code: &str) -> Response } /// Mint the token of an approved request, in its poll's transaction `tx`. +/// Everything is read on the transaction's connection: the poll holds the +/// request's lock, and waiting for a second connection meanwhile could wait +/// on the polls that hold the rest of the pool (#794). async fn collect( state: &AppState, mut tx: sqlx::Transaction<'static, sqlx::Postgres>, request: &PolledRequest, user_id: i64, ) -> Response { - let user = match user_dao::find_user_by_id(&state.pool, user_id).await { + let user = match user_dao::find_user_by_id(&mut *tx, user_id).await { Ok(user) => user.filter(|u| u.active), Err(e) => return oauth_unavailable(&e), }; @@ -658,7 +675,7 @@ async fn collect( }; let token_scope = match &scope { None => None, - Some(scope) => match scope.token_scope(state).await { + Some(scope) => match scope.token_scope(&mut tx).await { Ok(token_scope) => Some(token_scope), Err(e) => return oauth_unavailable(&e), }, diff --git a/crates/server/tests/common/mod.rs b/crates/server/tests/common/mod.rs index fe0209a2..2682976e 100644 --- a/crates/server/tests/common/mod.rs +++ b/crates/server/tests/common/mod.rs @@ -126,6 +126,23 @@ pub async fn pool() -> PgPool { pool } +/// A pool on the shared database the way `main` builds the server's +/// ([`forge_db::connect_with`]): at most `max_connections` +/// (`database.max_connections`), giving up on a free one after +/// `acquire_timeout` (`database.acquire_timeout_seconds`). For a test of what +/// a server does with a pool that small. +pub async fn sized_pool(max_connections: u32, acquire_timeout: Duration) -> PgPool { + // Migrates the shared database first. + pool().await.close().await; + let url = match std::env::var("FORGE_TEST_DATABASE_URL") { + Ok(url) => url, + Err(_) => format!("postgres://forge:forge@{}/forge", pg_endpoint().await), + }; + forge_db::connect_with(&url, max_connections, acquire_timeout) + .await + .expect("connect") +} + /// Every test database starts with its local realm accepting passwords from /// clients (tests sign in over Basic), and the auth settings at their /// defaults (anonymous access off). A test that needs others changes them on diff --git a/crates/server/tests/device_auth.rs b/crates/server/tests/device_auth.rs index 03b634fc..66c68edf 100644 --- a/crates/server/tests/device_auth.rs +++ b/crates/server/tests/device_auth.rs @@ -4,7 +4,8 @@ //! web UI's `/device` page (`/auth/device/lookup`, `/auth/device/decide`). //! //! Covers a pending request, `slow_down`, expiry, denial and approval, a -//! token collected once (and what it is), the scope, unknown clients and +//! token collected once (and what it is), concurrent polls of one code on a +//! small connection pool, the scope, unknown clients and //! malformed requests (always RFC 6749 JSON), the rate limit on guessing //! codes per network and per session, and the settings that turn the //! device login off and set its interval. Requires Docker (Postgres + @@ -13,6 +14,9 @@ mod common; +use std::future::Future; +use std::time::{Duration, Instant}; + use common::{fresh_pool, start_with, unique, BrowserSession, ServerOpts, TestServer}; use reqwest::StatusCode; use serde_json::{json, Value}; @@ -360,6 +364,83 @@ async fn an_approved_login_collects_its_token_once() { assert_eq!(resp.status(), 401); } +/// `fut`'s output and how long it took. +async fn timed(fut: F) -> (F::Output, Duration) { + let started = Instant::now(); + let out = fut.await; + (out, started.elapsed()) +} + +/// Concurrent polls of an approved code neither wait for one another nor +/// take the pool (#794). On a server of two database connections +/// (`database.max_connections = 2`, the default acquire timeout of 10 s), +/// six polls at once: one gets the token, every other answers at once +/// (`slow_down` while it is minted, `expired_token` after), and a request +/// of the signed-in user is served meanwhile. Before, the poll minting the +/// token held the request's lock and waited for a second connection, which +/// the polls waiting for the lock held: each answered 503 after the acquire +/// timeout, and so did every other request. +#[tokio::test] +async fn concurrent_polls_of_an_approved_code_leave_the_pool_alone() { + let srv = start_on(common::sized_pool(2, Duration::from_secs(10)).await).await; + // A scope naming a repository: minting the token looks it up. + let scope = format!("read repository:dev-raw-{}", unique()); + let started = start_login(&srv, "forge-cli", Some(&scope)).await; + let device_code = started["device_code"].as_str().unwrap(); + let user_code = started["user_code"].as_str().unwrap(); + let (user_id, username, session) = signed_in(&srv, "writer").await; + assert_eq!( + decide(&srv, &session, user_code, true, "198.51.100.20") + .await + .status(), + 204 + ); + + // Well below the acquire timeout, with slack for a loaded machine. + let bound = Duration::from_secs(5); + let polls = + futures_util::future::join_all((0..6).map(|_| timed(poll(&srv, "forge-cli", device_code)))); + let unrelated = timed(async { + session + .apply( + srv.client() + .get(format!("{}/admin/me/tokens", srv.base_url)), + ) + .send() + .await + .unwrap() + .status() + }); + let (polls, (status, took)) = tokio::join!(polls, unrelated); + assert_eq!(status, 200, "a request beside the polls"); + assert!(took < bound, "a request beside the polls took {took:?}"); + let mut collected = 0; + for ((status, body), took) in polls { + assert!(took < bound, "a poll took {took:?}: {status} {body}"); + if status == 200 { + collected += 1; + assert_eq!(body["username"], username.as_str(), "{body}"); + assert_eq!(body["scope"], scope.as_str(), "{body}"); + } else { + assert_eq!(status, 400, "{body}"); + assert!( + matches!(body["error"].as_str(), Some("slow_down" | "expired_token")), + "{body}" + ); + } + } + assert_eq!(collected, 1, "exactly one poll gets the token"); + assert_error(poll(&srv, "forge-cli", device_code).await, "expired_token"); + let minted: i64 = sqlx::query_scalar( + "SELECT count(*) FROM token WHERE user_id = $1 AND name = 'forge-cli login (test.local)'", + ) + .bind(user_id) + .fetch_one(&srv.pool) + .await + .unwrap(); + assert_eq!(minted, 1); +} + /// A denied login answers `access_denied`, and mints nothing. #[tokio::test] async fn a_denied_login_is_access_denied() { diff --git a/docs/deployment.md b/docs/deployment.md index 0b93978e..17f2c413 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -698,7 +698,7 @@ authorization server metadata ([RFC 8414](https://www.rfc-editor.org/rfc/rfc8414 | --- | --- | --- | | `GET` | `/.well-known/oauth-authorization-server` | `issuer` (`server.public_base_url`), `device_authorization_endpoint`, `token_endpoint`, `grant_types_supported` (`urn:ietf:params:oauth:grant-type:device_code`). | | `POST` | `/auth/device/code` | Start a login (RFC 8628 §3.1): `client_id` and an optional `scope`, form-encoded. Answers `device_code`, `user_code`, `verification_uri` (`/device`), `verification_uri_complete` (`/device?code=BCDF-GHJK`), `expires_in` and `interval`. | -| `POST` | `/auth/device/token` | Poll it (§3.4): `grant_type=urn:ietf:params:oauth:grant-type:device_code`, `client_id` and `device_code`. Answers `authorization_pending` until the user decides, `slow_down` when polled sooner than the interval (which then grows by five seconds), `access_denied`, `expired_token`, or the token: `access_token`, `token_type` `Bearer`, `expires_in`, `scope` when one was asked for, and `username`, whose token it is. | +| `POST` | `/auth/device/token` | Poll it (§3.4): `grant_type=urn:ietf:params:oauth:grant-type:device_code`, `client_id` and `device_code`. Answers `authorization_pending` until the user decides, `slow_down` when polled sooner than the interval (which then grows by five seconds) or while another poll of the same code is answered, `access_denied`, `expired_token`, or the token: `access_token`, `token_type` `Bearer`, `expires_in`, `scope` when one was asked for, and `username`, whose token it is. | ```bash curl -s https://forge.example.com/auth/device/code \