crank/crates/crank-registry/tests/integration/credential_touch.rs

use super::common::*;

use std::time::{Duration, Instant};

use crank_core::{
    PlatformApiKey, PlatformApiKeyId, PlatformApiKeyKind, PlatformApiKeyScope,
    PlatformApiKeyStatus, Secret, SecretId, SecretKind, SecretStatus, Workspace, WorkspaceId,
    WorkspaceStatus,
};
use serde_json::json;
use sqlx::{PgPool, Row};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};

use crank_registry::{
    CreatePlatformApiKeyRequest, CreateSecretRequest, CreateWorkspaceRequest,
    MASTER_KEY_CIPHER_CONTRACT, MasterKeyIdentityCandidate, RegistryError,
};

fn timestamp(value: &str) -> OffsetDateTime {
    OffsetDateTime::parse(value, &Rfc3339).unwrap()
}

fn workspace() -> Workspace {
    Workspace {
        id: WorkspaceId::new("ws_credential_touch"),
        slug: "credential-touch".to_owned(),
        display_name: "Credential Touch".to_owned(),
        status: WorkspaceStatus::Active,
        settings: json!({}),
        created_at: timestamp("2026-08-24T12:00:00Z"),
        updated_at: timestamp("2026-08-24T12:00:00Z"),
    }
}

async fn wait_for_row_lock(inspector: &PgPool, holder_pid: i32, query_marker: &str) {
    let deadline = Instant::now() + Duration::from_secs(5);

    loop {
        let blocked: i64 = sqlx::query_scalar(
            "select count(*)::bigint
             from pg_stat_activity activity
             where $1 = any(pg_blocking_pids(activity.pid))
               and activity.wait_event_type = 'Lock'
               and activity.query like '%' || $2 || '%'",
        )
        .bind(holder_pid)
        .bind(query_marker)
        .fetch_one(inspector)
        .await
        .unwrap();

        if blocked > 0 {
            return;
        }

        if Instant::now() >= deadline {
            let diagnostics = sqlx::query(
                "select pid, state, wait_event_type, wait_event, query
                 from pg_stat_activity
                 where $1 = any(pg_blocking_pids(pid))
                    or pid = $1",
            )
            .bind(holder_pid)
            .fetch_all(inspector)
            .await
            .unwrap()
            .into_iter()
            .map(|row| {
                format!(
                    "pid={}, state={}, wait_event_type={:?}, wait_event={:?}, query={}",
                    row.get::<i32, _>("pid"),
                    row.get::<String, _>("state"),
                    row.get::<Option<String>, _>("wait_event_type"),
                    row.get::<Option<String>, _>("wait_event"),
                    row.get::<String, _>("query"),
                )
            })
            .collect::<Vec<_>>();
            panic!(
                "touch query did not block on backend {holder_pid} within 5 seconds; \
                 pg_stat_activity: {diagnostics:#?}"
            );
        }

        tokio::task::yield_now().await;
    }
}

#[tokio::test]
async fn touch_platform_api_key_rejects_revoke_that_won_row_lock() {
    let database = TestDatabase::new().await;
    let registry = database.registry().await;
    let workspace = workspace();
    let key = PlatformApiKey {
        id: PlatformApiKeyId::new("key_touch_race"),
        workspace_id: workspace.id.clone(),
        agent_id: None,
        key_kind: PlatformApiKeyKind::McpClient,
        name: "Race key".to_owned(),
        prefix: "crk_live".to_owned(),
        scopes: vec![PlatformApiKeyScope::Read],
        status: PlatformApiKeyStatus::Active,
        created_at: timestamp("2026-08-24T12:00:00Z"),
        last_used_at: None,
        expires_at: None,
        allowed_origins: Vec::new(),
    };

    registry
        .create_workspace(CreateWorkspaceRequest {
            workspace: &workspace,
        })
        .await
        .unwrap();
    registry
        .create_platform_api_key(CreatePlatformApiKeyRequest {
            api_key: &key,
            secret_hash: "touch-race-secret-hash",
        })
        .await
        .unwrap();

    let lock_pool = database.raw_pool().await;
    let inspector = database.raw_pool().await;
    let mut lock_tx = lock_pool.begin().await.unwrap();
    let holder_pid: i32 = sqlx::query_scalar("select pg_backend_pid()")
        .fetch_one(&mut *lock_tx)
        .await
        .unwrap();
    sqlx::query(
        "select id
         from platform_api_keys
         where workspace_id = $1 and id = $2
         for update",
    )
    .bind(workspace.id.as_str())
    .bind(key.id.as_str())
    .execute(&mut *lock_tx)
    .await
    .unwrap();

    let touch_registry = registry.clone();
    let touch_workspace_id = workspace.id.clone();
    let touch_key_id = key.id.clone();
    let touch = tokio::spawn(async move {
        touch_registry
            .touch_platform_api_key(
                &touch_workspace_id,
                &touch_key_id,
                &timestamp("2026-08-24T12:05:00Z"),
            )
            .await
    });

    wait_for_row_lock(&inspector, holder_pid, "from platform_api_keys").await;

    sqlx::query(
        "update platform_api_keys
         set status = 'revoked'
         where workspace_id = $1 and id = $2",
    )
    .bind(workspace.id.as_str())
    .bind(key.id.as_str())
    .execute(&mut *lock_tx)
    .await
    .unwrap();
    lock_tx.commit().await.unwrap();

    let touch_result = touch.await.unwrap();
    assert!(matches!(
        touch_result,
        Err(RegistryError::PlatformApiKeyInactive { key_id }) if key_id == key.id.as_str()
    ));
    let last_used_at: Option<OffsetDateTime> = sqlx::query_scalar(
        "select last_used_at
         from platform_api_keys
         where workspace_id = $1 and id = $2",
    )
    .bind(workspace.id.as_str())
    .bind(key.id.as_str())
    .fetch_one(registry.pool())
    .await
    .unwrap();
    assert_eq!(last_used_at, None);

    database.cleanup().await;
}

#[tokio::test]
async fn touch_secret_rejects_disable_that_won_row_lock() {
    let database = TestDatabase::new().await;
    let registry = database.registry().await;
    let workspace = workspace();
    let secret = Secret {
        id: SecretId::new("secret_touch_race"),
        workspace_id: workspace.id.clone(),
        name: "Race secret".to_owned(),
        kind: SecretKind::Token,
        status: SecretStatus::Active,
        current_version: 1,
        created_at: timestamp("2026-08-24T12:00:00Z"),
        updated_at: timestamp("2026-08-24T12:00:00Z"),
        last_used_at: None,
    };

    registry
        .create_workspace(CreateWorkspaceRequest {
            workspace: &workspace,
        })
        .await
        .unwrap();
    registry
        .verify_or_register_master_key_identity(MasterKeyIdentityCandidate {
            epoch: 1,
            fingerprint: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
            cipher_contract: MASTER_KEY_CIPHER_CONTRACT,
            observed_at: &timestamp("2026-08-24T12:00:00Z"),
        })
        .await
        .unwrap();
    registry
        .create_secret(CreateSecretRequest {
            secret: &secret,
            ciphertext: "touch-race-ciphertext",
            key_version: "test-key-v1",
            master_key_epoch: 1,
            created_by: None,
        })
        .await
        .unwrap();

    let lock_pool = database.raw_pool().await;
    let inspector = database.raw_pool().await;
    let mut lock_tx = lock_pool.begin().await.unwrap();
    let holder_pid: i32 = sqlx::query_scalar("select pg_backend_pid()")
        .fetch_one(&mut *lock_tx)
        .await
        .unwrap();
    sqlx::query(
        "select id
         from secrets
         where workspace_id = $1 and id = $2
         for update",
    )
    .bind(workspace.id.as_str())
    .bind(secret.id.as_str())
    .execute(&mut *lock_tx)
    .await
    .unwrap();

    let touch_registry = registry.clone();
    let touch_workspace_id = workspace.id.clone();
    let touch_secret_id = secret.id.clone();
    let touch = tokio::spawn(async move {
        touch_registry
            .touch_secret(
                &touch_workspace_id,
                &touch_secret_id,
                &timestamp("2026-08-24T12:05:00Z"),
            )
            .await
    });

    wait_for_row_lock(&inspector, holder_pid, "from secrets").await;

    sqlx::query(
        "update secrets
         set status = 'disabled'
         where workspace_id = $1 and id = $2",
    )
    .bind(workspace.id.as_str())
    .bind(secret.id.as_str())
    .execute(&mut *lock_tx)
    .await
    .unwrap();
    lock_tx.commit().await.unwrap();

    let touch_result = touch.await.unwrap();
    assert!(matches!(
        touch_result,
        Err(RegistryError::SecretInactive { secret_id }) if secret_id == secret.id.as_str()
    ));
    let last_used_at: Option<OffsetDateTime> = sqlx::query_scalar(
        "select last_used_at
         from secrets
         where workspace_id = $1 and id = $2",
    )
    .bind(workspace.id.as_str())
    .bind(secret.id.as_str())
    .fetch_one(registry.pool())
    .await
    .unwrap();
    assert_eq!(last_used_at, None);

    database.cleanup().await;
}