fix: revoke mcp oauth grants on password change and reset (#11908)

* fix: revoke mcp oauth grants on password change and reset Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * fix: revoke pending mcp codes and serialize grants against revocation Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * fix: run mcp refresh theft check on the grant transaction Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * fix: log commit failures of rejected mcp oauth exchanges Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-22 16:03:43 +00:00 · 2026-10-21 13:58:03 +02:00
parent d35c86896f
commit c04477723c
6 changed files with 287 additions and 36 deletions
@@ -0,0 +1,22 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "SELECT user_email FROM mcp_oauth_refresh_token WHERE refresh_token = $1",
+  "describe": {
+    "columns": [
+      {
+        "ordinal": 0,
+        "name": "user_email",
+        "type_info": "Varchar"
+      }
+    ],
+    "parameters": {
+      "Left": [
+        "Text"
+      ]
+    },
+    "nullable": [
+      false
+    ]
+  },
+  "hash": "257b3bc19f738fe36e3c995c5ba239e1f2855cf39f9fe7a92d722732653fd98b"
+}
@@ -0,0 +1,22 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "SELECT user_email FROM mcp_oauth_server_code WHERE code = $1",
+  "describe": {
+    "columns": [
+      {
+        "ordinal": 0,
+        "name": "user_email",
+        "type_info": "Varchar"
+      }
+    ],
+    "parameters": {
+      "Left": [
+        "Text"
+      ]
+    },
+    "nullable": [
+      false
+    ]
+  },
+  "hash": "706872f05247a7d26fcd1b40b129f2defef86f8139dc4e38ae3bf63db039c48b"
+}
@@ -0,0 +1,14 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "WITH revoked AS (\n             DELETE FROM token WHERE email = $1 AND label LIKE 'mcp-oauth-%' RETURNING token_prefix\n         ), refresh AS (\n             DELETE FROM mcp_oauth_refresh_token WHERE user_email = $1\n         ), code AS (\n             DELETE FROM mcp_oauth_server_code WHERE user_email = $1\n         )\n         INSERT INTO notify_event (channel, payload)\n         SELECT 'notify_token_invalidation', token_prefix FROM revoked",
+  "describe": {
+    "columns": [],
+    "parameters": {
+      "Left": [
+        "Text"
+      ]
+    },
+    "nullable": []
+  },
+  "hash": "b7c7384b335a51f7d1abd2c627ea62ec651c5c6eee947540e7ee5dcea4780b37"
+}
@@ -1046,3 +1046,117 @@ async fn test_impersonation_session_is_never_refreshed(db: Pool<Postgres>) -> an

    Ok(())
 }
+
+#[cfg(all(feature = "enterprise", feature = "private"))]
+#[sqlx::test(migrations = "../migrations", fixtures("base"))]
+async fn test_set_password_revokes_mcp_oauth_grants(db: Pool<Postgres>) -> anyhow::Result<()> {
+    initialize_tracing().await;
+    let server = ApiServer::start(db.clone()).await?;
+
+    sqlx::raw_sql(
+        "INSERT INTO mcp_oauth_server_client (client_id, client_name, redirect_uris)
+         VALUES ('mcp-client', 'MCP client', ARRAY['http://localhost/cb']);
+         INSERT INTO token (token_hash, token_prefix, email, label, expiration)
+         VALUES (encode(sha256('MCP_ACCESS_TOKEN'::bytea), 'hex'), 'MCP_ACCESS', 'test2@windmill.dev',
+                 'mcp-oauth-mcp-client', now() + interval '7 days'),
+                (encode(sha256('OTHER_MCP_TOKEN'::bytea), 'hex'), 'OTHER_MCP_', 'test3@windmill.dev',
+                 'mcp-oauth-mcp-client', now() + interval '7 days');
+         INSERT INTO mcp_oauth_refresh_token
+             (refresh_token, access_token_hash, client_id, user_email, workspace_id, scopes, token_family, expires_at)
+         VALUES ('test2-refresh', encode(sha256('MCP_ACCESS_TOKEN'::bytea), 'hex'), 'mcp-client',
+                 'test2@windmill.dev', 'test-workspace', ARRAY['mcp:all'], gen_random_uuid(), now() + interval '30 days'),
+                ('test3-refresh', encode(sha256('OTHER_MCP_TOKEN'::bytea), 'hex'), 'mcp-client',
+                 'test3@windmill.dev', 'test-workspace', ARRAY['mcp:all'], gen_random_uuid(), now() + interval '30 days');
+         INSERT INTO mcp_oauth_server_code (code, client_id, user_email, workspace_id, scopes, redirect_uri)
+         VALUES ('test2-code', 'mcp-client', 'test2@windmill.dev', 'test-workspace', ARRAY['mcp:all'],
+                 'http://localhost/cb');",
+    )
+    .execute(&db)
+    .await?;
+
+    // A refresh in flight: it holds the grant lock and writes its rotated tokens while the
+    // password change is waiting on that lock.
+    let mut grant = db.begin().await?;
+    sqlx::query("SELECT pg_advisory_xact_lock_shared(hashtext('mcp_oauth_grant:' || $1))")
+        .bind("test2@windmill.dev")
+        .execute(&mut *grant)
+        .await?;
+    let request = tokio::spawn(
+        authed(client().post(format!(
+            "http://localhost:{}/api/users/set_password_of/test2@windmill.dev",
+            server.addr.port()
+        )))
+        .json(&json!({"password": "new-test-password-123"}))
+        .send(),
+    );
+    let mut waiting = false;
+    for _ in 0..100 {
+        waiting = sqlx::query_scalar(
+            "SELECT EXISTS(SELECT 1 FROM pg_locks WHERE locktype = 'advisory' AND NOT granted
+             AND database = (SELECT oid FROM pg_database WHERE datname = current_database()))",
+        )
+        .fetch_one(&db)
+        .await?;
+        if waiting {
+            break;
+        }
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+    }
+    assert!(
+        waiting,
+        "the password change never waited on the grant lock"
+    );
+    sqlx::raw_sql(
+        "INSERT INTO token (token_hash, token_prefix, email, label, expiration)
+         VALUES (encode(sha256('MCP_ROTATED_TOKEN'::bytea), 'hex'), 'MCP_ROTATE', 'test2@windmill.dev',
+                 'mcp-oauth-mcp-client', now() + interval '7 days');
+         INSERT INTO mcp_oauth_refresh_token
+             (refresh_token, access_token_hash, client_id, user_email, workspace_id, scopes, token_family, expires_at)
+         VALUES ('test2-rotated', encode(sha256('MCP_ROTATED_TOKEN'::bytea), 'hex'), 'mcp-client',
+                 'test2@windmill.dev', 'test-workspace', ARRAY['mcp:all'], gen_random_uuid(), now() + interval '30 days');",
+    )
+    .execute(&mut *grant)
+    .await?;
+    grant.commit().await?;
+
+    let resp = request.await??;
+    assert_eq!(resp.status(), 200, "{}", resp.text().await?);
+
+    let labels: Vec<(String, String)> = sqlx::query_as(
+        "SELECT email, label FROM token WHERE email IN ('test2@windmill.dev', 'test3@windmill.dev') ORDER BY email, label",
+    )
+    .fetch_all(&db)
+    .await?;
+    assert_eq!(
+        labels,
+        vec![
+            ("test2@windmill.dev".to_string(), "test token 2".to_string()),
+            (
+                "test3@windmill.dev".to_string(),
+                "mcp-oauth-mcp-client".to_string()
+            ),
+            ("test3@windmill.dev".to_string(), "test token 3".to_string()),
+        ]
+    );
+    let refresh: Vec<String> = sqlx::query_scalar(
+        "SELECT refresh_token FROM mcp_oauth_refresh_token ORDER BY refresh_token",
+    )
+    .fetch_all(&db)
+    .await?;
+    assert_eq!(refresh, vec!["test3-refresh".to_string()]);
+    let invalidated: Vec<String> = sqlx::query_scalar(
+        "SELECT payload FROM notify_event WHERE channel = 'notify_token_invalidation' ORDER BY payload",
+    )
+    .fetch_all(&db)
+    .await?;
+    assert_eq!(
+        invalidated,
+        vec!["MCP_ACCESS".to_string(), "MCP_ROTATE".to_string()]
+    );
+    let codes: i64 = sqlx::query_scalar("SELECT count(*) FROM mcp_oauth_server_code")
+        .fetch_one(&db)
+        .await?;
+    assert_eq!(codes, 0);
+
+    Ok(())
+}
@@ -365,18 +365,59 @@ pub async fn oauth_token(
    Extension(db): Extension<DB>,
    Form(req): Form<TokenRequest>,
 ) -> std::result::Result<Json<TokenResponse>, OAuthTokenError> {
-    match req.grant_type.as_str() {
-        "authorization_code" => handle_authorization_code_grant(&db, &req).await,
-        "refresh_token" => handle_refresh_token_grant(&db, &req).await,
-        _ => Err(OAuthTokenError::unsupported_grant_type(
-            "Supported grant types: authorization_code, refresh_token",
-        )),
+    let db_error = |e: sqlx::Error| {
+        tracing::error!("Database error issuing MCP OAuth token: {}", e);
+        OAuthTokenError::server_error("Database error")
+    };
+    let owner = match req.grant_type.as_str() {
+        "authorization_code" => sqlx::query_scalar!(
+            "SELECT user_email FROM mcp_oauth_server_code WHERE code = $1",
+            req.code
+        )
+        .fetch_optional(&db)
+        .await
+        .map_err(db_error)?,
+        "refresh_token" => sqlx::query_scalar!(
+            "SELECT user_email FROM mcp_oauth_refresh_token WHERE refresh_token = $1",
+            req.refresh_token.as_deref().unwrap_or_default()
+        )
+        .fetch_optional(&db)
+        .await
+        .map_err(db_error)?,
+        _ => {
+            return Err(OAuthTokenError::unsupported_grant_type(
+                "Supported grant types: authorization_code, refresh_token",
+            ))
+        }
+    };
+    let mut tx = db.begin().await.map_err(db_error)?;
+    if let Some(email) = &owner {
+        crate::users::lock_mcp_oauth_grants(&mut tx, email, false)
+            .await
+            .map_err(db_error)?;
+    }
+    let res = if req.grant_type == "authorization_code" {
+        handle_authorization_code_grant(&mut tx, &req).await
+    } else {
+        handle_refresh_token_grant(&mut tx, &req).await
+    };
+    // Committed on error too, so a rejected exchange still consumes its code or refresh token.
+    let committed = tx.commit().await;
+    match (res, committed) {
+        (Ok(res), Ok(())) => Ok(res),
+        (Err(e), committed) => {
+            if let Err(ce) = committed {
+                tracing::error!("Failed to commit rejected MCP OAuth exchange: {}", ce);
+            }
+            Err(e)
+        }
+        (Ok(_), Err(e)) => Err(db_error(e)),
    }
 }

 /// Handle authorization_code grant type
 async fn handle_authorization_code_grant(
-    db: &DB,
+    tx: &mut sqlx::PgConnection,
    req: &TokenRequest,
 ) -> std::result::Result<Json<TokenResponse>, OAuthTokenError> {
    let auth_code = match sqlx::query_as!(
@@ -387,7 +428,7 @@ async fn handle_authorization_code_grant(
                   code_challenge, code_challenge_method",
        req.code
    )
-    .fetch_optional(db)
+    .fetch_optional(&mut *tx)
    .await
    {
        Ok(Some(code)) => code,
@@ -458,7 +499,7 @@ async fn handle_authorization_code_grant(
        &scopes,
        auth_code.workspace_id,
    )
-    .execute(db)
+    .execute(&mut *tx)
    .await
    .map_err(|e| {
        tracing::error!("Failed to create access token: {}", e);
@@ -471,7 +512,7 @@ async fn handle_authorization_code_grant(
    }

    // Create refresh token — store the hash of the access token so we can delete it later
-    let refresh_token_result = sqlx::query!(
+    sqlx::query!(
        "INSERT INTO mcp_oauth_refresh_token
         (refresh_token, access_token_hash, client_id, user_email, workspace_id, scopes, token_family, expires_at)
         VALUES ($1, $2, $3, $4, $5, $6, $7, now() + ($8 || ' seconds')::interval)",
@@ -484,29 +525,25 @@ async fn handle_authorization_code_grant(
        token_family,
        MCP_OAUTH_REFRESH_TOKEN_EXPIRATION_SECS.to_string(),
    )
-    .execute(db)
-    .await;
-
-    let refresh_token_value = match refresh_token_result {
-        Ok(_) => Some(refresh_token),
-        Err(e) => {
-            tracing::error!("Failed to create refresh token: {}", e);
-            None // Don't include invalid refresh token
-        }
-    };
+    .execute(&mut *tx)
+    .await
+    .map_err(|e| {
+        tracing::error!("Failed to create refresh token: {}", e);
+        OAuthTokenError::server_error("Failed to create refresh token")
+    })?;

    Ok(Json(TokenResponse {
        access_token,
        token_type: "Bearer".to_string(),
        expires_in: MCP_OAUTH_TOKEN_EXPIRATION_SECS,
        scope: Some(scopes.join(" ")),
-        refresh_token: refresh_token_value,
+        refresh_token: Some(refresh_token),
    }))
 }

 /// Handle refresh_token grant type with token rotation and theft detection
 async fn handle_refresh_token_grant(
-    db: &DB,
+    tx: &mut sqlx::PgConnection,
    req: &TokenRequest,
 ) -> std::result::Result<Json<TokenResponse>, OAuthTokenError> {
    let refresh_token_value = req
@@ -533,7 +570,7 @@ async fn handle_refresh_token_grant(
        refresh_token_value,
        req.client_id
    )
-    .fetch_optional(db)
+    .fetch_optional(&mut *tx)
    .await
    {
        Ok(Some(row)) => row,
@@ -544,7 +581,7 @@ async fn handle_refresh_token_grant(
                 WHERE refresh_token = $1 AND used_at IS NOT NULL",
                refresh_token_value
            )
-            .fetch_optional(db)
+            .fetch_optional(&mut *tx)
            .await
            {
                tracing::warn!("Refresh token reuse detected, revoking family {:?}", family);
@@ -552,7 +589,7 @@ async fn handle_refresh_token_grant(
                    "UPDATE mcp_oauth_refresh_token SET revoked = TRUE WHERE token_family = $1",
                    family
                )
-                .execute(db)
+                .execute(&mut *tx)
                .await;
            }
            return Err(OAuthTokenError::invalid_grant("Invalid refresh token"));
@@ -564,16 +601,16 @@ async fn handle_refresh_token_grant(
    };

    // Delete old access token using the stored hash
-    if let Err(e) = sqlx::query!(
+    sqlx::query!(
        "DELETE FROM token WHERE token_hash = $1",
        token_row.access_token_hash
    )
-    .execute(db)
+    .execute(&mut *tx)
    .await
-    {
+    .map_err(|e| {
        tracing::error!("Failed to delete old access token: {}", e);
-        // Non-fatal, continue with token creation
-    }
+        OAuthTokenError::server_error("Database error")
+    })?;

    // Generate new tokens
    let new_access_token = rd_string(32);
@@ -605,7 +642,7 @@ async fn handle_refresh_token_grant(
        &scopes,
        token_row.workspace_id,
    )
-    .execute(db)
+    .execute(&mut *tx)
    .await
    .map_err(|e| {
        tracing::error!("Failed to create new access token: {}", e);
@@ -618,7 +655,7 @@ async fn handle_refresh_token_grant(
    }

    // Create new refresh token (same token family for tracking) — store hash of access token
-    if let Err(e) = sqlx::query!(
+    sqlx::query!(
        "INSERT INTO mcp_oauth_refresh_token
         (refresh_token, access_token_hash, client_id, user_email, workspace_id, scopes, token_family, expires_at)
         VALUES ($1, $2, $3, $4, $5, $6, $7, now() + ($8 || ' seconds')::interval)",
@@ -631,12 +668,12 @@ async fn handle_refresh_token_grant(
        token_row.token_family,
        MCP_OAUTH_REFRESH_TOKEN_EXPIRATION_SECS.to_string(),
    )
-    .execute(db)
+    .execute(&mut *tx)
    .await
-    {
+    .map_err(|e| {
        tracing::error!("Failed to create new refresh token: {}", e);
-        // Access token was created, return success without refresh token
-    }
+        OAuthTokenError::server_error("Failed to create refresh token")
+    })?;

    Ok(Json(TokenResponse {
        access_token: new_access_token,
@@ -194,6 +194,41 @@ async fn list_guests(
    Ok(Json(GuestList { usage, guests }))
 }

+/// Shared while issuing an MCP OAuth grant, exclusive while revoking: no grant outlives a revoke.
+pub(crate) async fn lock_mcp_oauth_grants(
+    conn: &mut sqlx::PgConnection,
+    email: &str,
+    exclusive: bool,
+) -> std::result::Result<(), sqlx::Error> {
+    let q = if exclusive {
+        "SELECT pg_advisory_xact_lock(hashtext('mcp_oauth_grant:' || $1))"
+    } else {
+        "SELECT pg_advisory_xact_lock_shared(hashtext('mcp_oauth_grant:' || $1))"
+    };
+    sqlx::query(q).bind(email).execute(conn).await?;
+    Ok(())
+}
+
+async fn revoke_mcp_oauth_grants(tx: &mut sqlx::PgConnection, email: &str) -> Result<()> {
+    lock_mcp_oauth_grants(&mut *tx, email, true).await?;
+    // The token trigger only notifies for sessions, so evict deleted MCP access tokens explicitly.
+    sqlx::query!(
+        "WITH revoked AS (
+             DELETE FROM token WHERE email = $1 AND label LIKE 'mcp-oauth-%' RETURNING token_prefix
+         ), refresh AS (
+             DELETE FROM mcp_oauth_refresh_token WHERE user_email = $1
+         ), code AS (
+             DELETE FROM mcp_oauth_server_code WHERE user_email = $1
+         )
+         INSERT INTO notify_event (channel, payload)
+         SELECT 'notify_token_invalidation', token_prefix FROM revoked",
+        email
+    )
+    .execute(&mut *tx)
+    .await?;
+    Ok(())
+}
+
 async fn set_password(
    Extension(db): Extension<DB>,
    Extension(argon2): Extension<Arc<Argon2<'_>>>,
@@ -215,6 +250,9 @@ async fn set_password(
    )
    .execute(&db)
    .await?;
+    let mut tx = db.begin().await?;
+    revoke_mcp_oauth_grants(&mut *tx, &email).await?;
+    tx.commit().await?;
    Ok(res)
 }

@@ -235,6 +273,9 @@ async fn set_password_of_user(
    )
    .execute(&db)
    .await?;
+    let mut tx = db.begin().await?;
+    revoke_mcp_oauth_grants(&mut *tx, &email).await?;
+    tx.commit().await?;
    Ok(res)
 }

@@ -1087,6 +1128,7 @@ async fn reset_password(
    )
    .execute(&mut *tx)
    .await?;
+    revoke_mcp_oauth_grants(&mut *tx, &email).await?;

    // Audit log
    let audit_author = AuditAuthor {