Device login: concurrent polls of a code no longer hold the pool

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<Option<PolledRequest>> {
+) -> sqlx::Result<PollLock> {
     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,
     }))
 }
 

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<Option<Repository>> {
+) -> sqlx::Result<Option<Repository>>
+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
 }
 

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<TokenScope> {
+    /// 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<TokenScope> {
         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),
         },

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

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<F: Future>(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() {

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` (`<base>/device`), `verification_uri_complete` (`<base>/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 \