remora/crates/remora-admission/tests/substrate.rs

use std::time::Duration;

use remora_admission::{
    AdmissionConfig, AdmissionDecision, AdmissionError, AdmissionService, SubmitOutcome, migrate,
};
use remora_contracts::{
    HttpSourceSubmissionDto, JobSubmissionDto, SourceSubmissionDto, TargetOverrideDto,
};
use remora_credentials::{
    CredentialError, CredentialGrantVersion, CredentialRegistry, GrantBoundResolver, VaultClient,
    VaultConfig, migrate as credentials_migrate,
};
use remora_domain::{
    BucketName, OpaqueId, StorageScope, TargetPrefix, TargetStoreId, ports::ActiveCredentialGrant,
};
use remora_durable::{MutationDisposition, PoolConfig, connect, migrate as durable_migrate};
use remora_identity::{DefaultTarget, Principal};
use sqlx::PgPool;

fn database_url() -> Option<String> {
    std::env::var("DATABASE_URL")
        .ok()
        .filter(|value| !value.is_empty())
}

async fn pool() -> Option<PgPool> {
    let Some(url) = database_url() else {
        if std::env::var("REMORA_ADMISSION_REQUIRE").is_ok() {
            panic!("DATABASE_URL is required when REMORA_ADMISSION_REQUIRE is set");
        }
        return None;
    };
    let pool = connect(
        &url,
        PoolConfig {
            max_connections: 8,
            acquire_timeout: Duration::from_secs(10),
            idle_timeout: Some(Duration::from_secs(30)),
        },
    )
    .await
    .expect("connect postgres");
    durable_migrate(&pool).await.expect("durable migrate");
    credentials_migrate(&pool)
        .await
        .expect("credentials migrate");
    migrate(&pool).await.expect("admission migrate");
    Some(pool)
}

fn unique_id(prefix: &str) -> OpaqueId {
    let nonce = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_nanos();
    OpaqueId::parse(format!("{prefix}{nonce:032x}")).unwrap()
}

fn target_parts() -> (TargetStoreId, BucketName, TargetPrefix) {
    (
        TargetStoreId::parse("primary").unwrap(),
        BucketName::parse("remora-results").unwrap(),
        TargetPrefix::parse("users", ["_remora"]).unwrap(),
    )
}

fn principal(id: OpaqueId, authorized: bool, administrator: bool) -> Principal {
    let (store, bucket, prefix) = target_parts();
    Principal {
        id,
        issuer: "https://issuer.example".into(),
        subject: "subject".into(),
        administrator,
        scopes: authorized
            .then_some(StorageScope {
                store,
                bucket,
                prefix,
            })
            .into_iter()
            .collect(),
    }
}

fn service(pool: PgPool) -> AdmissionService {
    let (store, bucket, prefix) = target_parts();
    AdmissionService::new(
        pool,
        AdmissionConfig {
            default_target: DefaultTarget {
                store,
                bucket,
                prefix,
            },
            reserved_target_segments: vec!["_remora".into()],
            policy_version: "source-policy-v1".into(),
        },
    )
    .unwrap()
}

fn request(url: String) -> JobSubmissionDto {
    JobSubmissionDto {
        source: SourceSubmissionDto::Http(HttpSourceSubmissionDto {
            url,
            result_path: "artifact.bin".into(),
            credential_id: None,
        }),
        target: Some(TargetOverrideDto {
            bucket: None,
            prefix: Some("users/alice".into()),
        }),
        label: Some("artifact".into()),
    }
}

fn accepted(outcome: SubmitOutcome) -> remora_admission::AcceptedJob {
    match outcome.decision {
        AdmissionDecision::Accepted(job) => *job,
        AdmissionDecision::Rejected { .. } => panic!("submission rejected"),
    }
}

fn registry(pool: PgPool) -> CredentialRegistry {
    registry_and_resolver(pool).0
}

fn registry_and_resolver(pool: PgPool) -> (CredentialRegistry, GrantBoundResolver) {
    let vault = VaultClient::new(VaultConfig {
        address: "http://127.0.0.1:1".into(),
        token: "synthetic-test-token".into(),
        mount: "secret".into(),
        path_prefix: "admission-tests".into(),
        timeout: Duration::from_millis(100),
    })
    .unwrap();
    let registry = CredentialRegistry::new(pool, vault.clone());
    let resolver = GrantBoundResolver::new(registry.clone(), vault);
    (registry, resolver)
}

async fn insert_credential(
    pool: &PgPool,
    owner_id: &OpaqueId,
    kind: &str,
) -> CredentialGrantVersion {
    let credential_id = unique_id("cred");
    let version_id = unique_id("version");
    let now = time::OffsetDateTime::now_utc();
    let mut tx = pool.begin().await.unwrap();
    sqlx::query(
        "INSERT INTO source_credentials (
            credential_id, owner_id, label, kind, status, active_version_id,
            secret_mount, secret_path, created_at, updated_at
         ) VALUES ($1, $2, 'fixture', $3, 'active', $4, 'secret', 'fixture', $5, $5)",
    )
    .bind(credential_id.as_str())
    .bind(owner_id.as_str())
    .bind(kind)
    .bind(version_id.as_str())
    .bind(now)
    .execute(&mut *tx)
    .await
    .unwrap();
    sqlx::query(
        "INSERT INTO source_credential_versions (
            version_id, credential_id, vault_kv_version, created_at, destroyed_at, secret_mount, secret_path
         ) VALUES ($1, $2, 1, $3, NULL, 'secret', 'fixture/version')",
    )
    .bind(version_id.as_str())
    .bind(credential_id.as_str())
    .bind(now)
    .execute(&mut *tx)
    .await
    .unwrap();
    tx.commit().await.unwrap();
    CredentialGrantVersion {
        credential_id,
        version_id,
    }
}

#[tokio::test]
async fn exact_replay_and_lost_response_recovery_create_one_durable_acceptance() {
    let Some(pool) = pool().await else { return };
    let service = service(pool.clone());
    let owner = principal(unique_id("owner"), true, false);
    let key = unique_id("key").to_string();
    let submission = request("https://downloads.example/archive?token=private".into());
    let first = service.submit(&owner, &key, submission).await.unwrap();
    assert_eq!(first.disposition, MutationDisposition::First);
    let first_job = accepted(first);

    let replay = service
        .submit(
            &owner,
            &key,
            request("https://downloads.example/archive?token=private".into()),
        )
        .await
        .unwrap();
    assert_eq!(replay.disposition, MutationDisposition::Replay);
    let replay_job = accepted(replay);
    assert_eq!(first_job.job_id, replay_job.job_id);
    assert_eq!(first_job.execution_grant_id, replay_job.execution_grant_id);

    let counts: (i64, i64, i64) = sqlx::query_as(
        "SELECT
            (SELECT count(*) FROM jobs WHERE job_id = $1),
            (SELECT count(*) FROM execution_grants WHERE grant_id = $2),
            (SELECT count(*) FROM outbox_messages WHERE topic = 'job.dispatch' AND payload->>'job_id' = $1)",
    )
    .bind(first_job.job_id.as_str())
    .bind(first_job.execution_grant_id.as_str())
    .fetch_one(&pool)
    .await
    .unwrap();
    assert_eq!(counts, (1, 1, 1));
}

#[tokio::test]
async fn changed_fingerprint_conflicts_and_current_authorization_is_rechecked_for_new_keys() {
    let Some(pool) = pool().await else { return };
    let service = service(pool);
    let owner_id = unique_id("owner");
    let owner = principal(owner_id.clone(), true, false);
    let key = unique_id("key").to_string();
    service
        .submit(
            &owner,
            &key,
            request("https://downloads.example/one".into()),
        )
        .await
        .unwrap();
    assert!(matches!(
        service
            .submit(
                &owner,
                &key,
                request("https://downloads.example/two".into())
            )
            .await,
        Err(AdmissionError::IdempotencyConflict)
    ));

    let lost_authorization = principal(owner_id, false, false);
    let replay = service
        .submit(
            &lost_authorization,
            &key,
            request("https://downloads.example/one".into()),
        )
        .await
        .unwrap();
    assert_eq!(replay.disposition, MutationDisposition::Replay);
    let fresh = service
        .submit(
            &lost_authorization,
            unique_id("key").as_str(),
            request("https://downloads.example/one".into()),
        )
        .await
        .unwrap();
    assert!(matches!(fresh.decision, AdmissionDecision::Rejected { .. }));
}

#[tokio::test]
async fn public_projection_and_outbox_do_not_leak_sensitive_url() {
    let Some(pool) = pool().await else { return };
    let service = service(pool.clone());
    let owner = principal(unique_id("owner"), true, false);
    let secret = unique_id("secret").to_string();
    let accepted = accepted(
        service
            .submit(
                &owner,
                unique_id("key").as_str(),
                request(format!(
                    "https://downloads.example/archive?token={secret}&signature=hidden"
                )),
            )
            .await
            .unwrap(),
    );
    let public_text: String = sqlx::query_scalar(
        "SELECT concat_ws(' ',
            (SELECT source_projection::text FROM jobs WHERE job_id = $1),
            (SELECT payload::text FROM outbox_messages WHERE topic = 'job.dispatch' AND payload->>'job_id' = $1),
            (SELECT response_body::text FROM mutation_records WHERE response_body->'accepted'->>'job_id' = $1)
         )",
    )
    .bind(accepted.job_id.as_str())
    .fetch_one(&pool)
    .await
    .unwrap();
    assert!(!public_text.contains(&secret));
    assert!(!public_text.contains("signature"));
    let private_text: String = sqlx::query_scalar(
        "SELECT source_input::text FROM sensitive_source_inputs
         WHERE sensitive_source_ref = (SELECT sensitive_source_ref FROM jobs WHERE job_id = $1)",
    )
    .bind(accepted.job_id.as_str())
    .fetch_one(&pool)
    .await
    .unwrap();
    assert!(private_text.contains(&secret));
}

#[tokio::test]
async fn credential_pin_and_accepted_grant_revoke_atomically_without_reactivation() {
    let Some(pool) = pool().await else { return };
    let service = service(pool.clone());
    let owner_id = unique_id("owner");
    let pin = insert_credential(&pool, &owner_id, "bearer_token").await;
    let credential_id = pin.credential_id.clone();
    let version_id = pin.version_id.clone();

    let mut request = request("https://downloads.example/private".into());
    let SourceSubmissionDto::Http(source) = &mut request.source else {
        unreachable!()
    };
    source.credential_id = Some(credential_id.clone());
    let accepted = accepted(
        service
            .submit(
                &principal(owner_id, true, false),
                unique_id("key").as_str(),
                request,
            )
            .await
            .unwrap(),
    );
    let binding: (String, String, String) = sqlx::query_as(
        "SELECT credential_id, version_id, status
         FROM execution_grant_credential_pins WHERE grant_id = $1",
    )
    .bind(accepted.execution_grant_id.as_str())
    .fetch_one(&pool)
    .await
    .unwrap();
    assert_eq!(
        binding,
        (
            credential_id.to_string(),
            version_id.to_string(),
            "active".into()
        )
    );

    let admin = principal(unique_id("admin"), false, true);
    let (left, right) = tokio::join!(
        service.revoke_execution_grant(&admin, &accepted.execution_grant_id),
        service.revoke_execution_grant(&admin, &accepted.execution_grant_id)
    );
    left.unwrap();
    right.unwrap();
    let statuses: (String, String, i64) = sqlx::query_as(
        "SELECT
            (SELECT status FROM execution_grants WHERE grant_id = $1),
            (SELECT status FROM execution_grant_credential_pins WHERE grant_id = $1),
            (SELECT count(*) FROM outbox_messages WHERE topic = 'execution_grant.revoked' AND payload->>'grant_id' = $1)",
    )
    .bind(accepted.execution_grant_id.as_str())
    .fetch_one(&pool)
    .await
    .unwrap();
    assert_eq!(statuses, ("revoked".into(), "revoked".into(), 1));

    let mut tx = pool.begin().await.unwrap();
    assert!(
        CredentialRegistry::bind_execution_grant_in_tx(
            &mut tx,
            &accepted.execution_grant_id,
            &pin,
            time::OffsetDateTime::now_utc(),
        )
        .await
        .is_err()
    );
    tx.rollback().await.unwrap();
}

#[tokio::test]
async fn authoritative_credential_tuple_rejects_both_binding_substitution_paths() {
    let Some(pool) = pool().await else { return };
    let service = service(pool.clone());
    let owner_id = unique_id("owner");
    let owner = principal(owner_id.clone(), true, false);
    let arbitrary_pin = insert_credential(&pool, &owner_id, "bearer_token").await;
    let credentialless = accepted(
        service
            .submit(
                &owner,
                unique_id("key").as_str(),
                request("https://downloads.example/public".into()),
            )
            .await
            .unwrap(),
    );
    let (registry, resolver) = registry_and_resolver(pool.clone());

    assert!(
        registry
            .bind_execution_grant(&credentialless.execution_grant_id, &arbitrary_pin)
            .await
            .is_err()
    );
    assert_eq!(
        registry
            .load_active_grant_pin(&credentialless.execution_grant_id)
            .await
            .unwrap(),
        None
    );
    let access =
        ActiveCredentialGrant::for_execution_grant(credentialless.execution_grant_id.clone());
    assert!(matches!(
        resolver.resolve_access(&access).await,
        Err(CredentialError::GrantInactive)
    ));

    let mut tx = pool.begin().await.unwrap();
    CredentialRegistry::bind_execution_grant_in_tx(
        &mut tx,
        &credentialless.execution_grant_id,
        &arbitrary_pin,
        time::OffsetDateTime::now_utc(),
    )
    .await
    .unwrap();
    assert!(tx.commit().await.is_err());
    assert_eq!(
        registry
            .load_active_grant_pin(&credentialless.execution_grant_id)
            .await
            .unwrap(),
        None
    );

    let accepted_pin = insert_credential(&pool, &owner_id, "bearer_token").await;
    let substitute_pin = insert_credential(&pool, &owner_id, "bearer_token").await;
    let mut credentialed_request = request("https://downloads.example/private".into());
    let SourceSubmissionDto::Http(source) = &mut credentialed_request.source else {
        unreachable!()
    };
    source.credential_id = Some(accepted_pin.credential_id.clone());
    let credentialed = accepted(
        service
            .submit(&owner, unique_id("key").as_str(), credentialed_request)
            .await
            .unwrap(),
    );
    assert_eq!(
        registry
            .load_active_grant_pin(&credentialed.execution_grant_id)
            .await
            .unwrap(),
        Some(accepted_pin.clone())
    );
    assert!(
        registry
            .bind_execution_grant(&credentialed.execution_grant_id, &accepted_pin)
            .await
            .is_ok()
    );
    assert!(
        registry
            .bind_execution_grant(&credentialed.execution_grant_id, &substitute_pin)
            .await
            .is_err()
    );
    assert_eq!(
        registry
            .load_active_grant_pin(&credentialed.execution_grant_id)
            .await
            .unwrap(),
        Some(accepted_pin)
    );
}

#[tokio::test]
async fn admission_transaction_is_invisible_to_concurrent_revoke_until_binding_commits() {
    let Some(pool) = pool().await else { return };
    let service = service(pool.clone());
    let owner_id = unique_id("owner");
    let pin = insert_credential(&pool, &owner_id, "bearer_token").await;
    let job_id = unique_id("job");
    let grant_id = unique_id("grant");
    let sensitive_source_ref = unique_id("source");
    let now = time::OffsetDateTime::now_utc();
    let mut admission_tx = pool.begin().await.unwrap();
    sqlx::query(
        "INSERT INTO sensitive_source_inputs (sensitive_source_ref, source_input, created_at)
         VALUES ($1, '{\"kind\":\"http\"}'::jsonb, $2)",
    )
    .bind(sensitive_source_ref.as_str())
    .bind(now)
    .execute(&mut *admission_tx)
    .await
    .unwrap();
    sqlx::query(
        "INSERT INTO execution_grants (
            grant_id, owner_id, target_store, target_bucket, target_prefix,
            policy_version, credential_id, credential_version_id, status, created_at, revoked_at
         ) VALUES ($1, $2, 'primary', 'remora-results', '[\"users\",\"alice\"]'::jsonb,
                   'source-policy-v1', $3, $4, 'active', $5, NULL)",
    )
    .bind(grant_id.as_str())
    .bind(owner_id.as_str())
    .bind(pin.credential_id.as_str())
    .bind(pin.version_id.as_str())
    .bind(now)
    .execute(&mut *admission_tx)
    .await
    .unwrap();
    sqlx::query(
        "INSERT INTO jobs (
            job_id, owner_id, label, revision, phase, lifecycle, source_projection,
            sensitive_source_ref, target_store, target_bucket, target_prefix,
            execution_grant_id, dispatch_needed, created_at, updated_at
         ) VALUES ($1, $2, NULL, 1, 'queued', 'visible', '{\"kind\":\"http\"}'::jsonb,
                   $3, 'primary', 'remora-results', '[\"users\",\"alice\"]'::jsonb,
                   $4, TRUE, $5, $5)",
    )
    .bind(job_id.as_str())
    .bind(owner_id.as_str())
    .bind(sensitive_source_ref.as_str())
    .bind(grant_id.as_str())
    .bind(now)
    .execute(&mut *admission_tx)
    .await
    .unwrap();
    CredentialRegistry::bind_execution_grant_in_tx(&mut admission_tx, &grant_id, &pin, now)
        .await
        .unwrap();

    let invisible: (i64, i64) = sqlx::query_as(
        "SELECT
            (SELECT count(*) FROM execution_grants WHERE grant_id = $1),
            (SELECT count(*) FROM execution_grant_credential_pins WHERE grant_id = $1)",
    )
    .bind(grant_id.as_str())
    .fetch_one(&pool)
    .await
    .unwrap();
    assert_eq!(invisible, (0, 0));
    let admin = principal(unique_id("admin"), false, true);
    assert!(matches!(
        service.revoke_execution_grant(&admin, &grant_id).await,
        Err(AdmissionError::GrantNotFound)
    ));

    admission_tx.commit().await.unwrap();
    assert_eq!(
        registry(pool.clone())
            .load_active_grant_pin(&grant_id)
            .await
            .unwrap(),
        Some(pin)
    );
    service
        .revoke_execution_grant(&admin, &grant_id)
        .await
        .unwrap();
    service
        .revoke_execution_grant(&admin, &grant_id)
        .await
        .unwrap();
    assert_eq!(
        registry(pool)
            .load_active_grant_pin(&grant_id)
            .await
            .unwrap(),
        None
    );
}