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

use std::time::Duration;

use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use jsonwebtoken::{Algorithm, EncodingKey, Header, encode};
use remora_durable::{PoolConfig, connect};
use remora_identity::{OidcConfig, OidcIdentityProvider, resolve_principal};
use remora_schema::migrate as migrate_schema;
use serde_json::{Value, json};
use sqlx::PgPool;

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

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

fn require_env() -> bool {
    std::env::var("REMORA_IDENTITY_REQUIRE").is_ok()
}

fn env_required(name: &str) -> String {
    std::env::var(name)
        .unwrap_or_else(|_| panic!("{name} is required when REMORA_IDENTITY_REQUIRE is set"))
}

async fn test_pool() -> Option<PgPool> {
    let Some(url) = database_url() else {
        if require_env() {
            panic!("DATABASE_URL is required when REMORA_IDENTITY_REQUIRE is set");
        }
        eprintln!("skipping identity durable test: DATABASE_URL is unset");
        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");
    migrate_schema(&pool).await.expect("schema migrate");
    Some(pool)
}

fn oidc_provider() -> OidcIdentityProvider {
    let issuer = env_required("REMORA_OIDC_ISSUER");
    let provider_kind = std::env::var("REMORA_OIDC_PROVIDER").unwrap_or_default();
    // Keycloak ROPC tokens use the OIDC client id as `aud` (via audience mapper).
    // ZITADEL JWT-profile access tokens assert the project resource as `aud`.
    let audience = if provider_kind == "zitadel" {
        env_required("REMORA_OIDC_PROJECT_ID")
    } else {
        env_required("REMORA_OIDC_CLIENT_ID")
    };
    let mut config = OidcConfig::new(issuer, audience).expect("oidc config");
    if provider_kind == "zitadel" {
        let project_id = env_required("REMORA_OIDC_PROJECT_ID");
        config = config
            .with_zitadel_project_roles()
            // Preferred project-scoped claim shape (ZITADEL docs).
            .with_project_role_claim(format!("urn:zitadel:iam:org:project:{project_id}:roles"))
            .with_resource_access_client(env_required("REMORA_OIDC_CLIENT_ID"));
    }
    OidcIdentityProvider::new(config).expect("provider")
}

#[tokio::test]
async fn principal_mapping_is_stable_and_idempotent() {
    let Some(pool) = test_pool().await else {
        return;
    };
    let external = remora_identity::ExternalIdentity {
        issuer: "https://issuer.example/lane".into(),
        subject: "subject-stable-1".into(),
        roles: remora_identity::DecodedRoles {
            administrator: false,
            scopes: vec![],
        },
    };
    let first = resolve_principal(&pool, external.clone())
        .await
        .expect("first resolve");
    let second = resolve_principal(&pool, external)
        .await
        .expect("second resolve");
    assert_eq!(first.id, second.id);
    assert_eq!(first.issuer, "https://issuer.example/lane");
    assert_eq!(first.subject, "subject-stable-1");

    let loaded = remora_identity::load_by_id(&pool, &first.id)
        .await
        .expect("load")
        .expect("present");
    assert_eq!(loaded.subject, "subject-stable-1");
}

#[tokio::test]
async fn real_oidc_jwks_rejects_invalid_tokens() {
    if oidc_issuer().is_none() {
        if require_env() {
            panic!("REMORA_OIDC_ISSUER is required when REMORA_IDENTITY_REQUIRE is set");
        }
        eprintln!("skipping oidc invalid-token test: REMORA_OIDC_ISSUER is unset");
        return;
    }
    let provider = oidc_provider();
    let minted = mint_unsigned_looking_token();
    let cases = [
        "",
        "not-a-jwt",
        "Bearer ",
        "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.e30.signature",
        minted.as_str(),
    ];
    for token in cases {
        let result = provider.authenticate_token(token).await;
        assert!(
            matches!(result, Err(remora_identity::IdentityError::InvalidToken)),
            "expected InvalidToken for {token:?}, got {result:?}"
        );
    }
}

// Sequential mint then revoke: one fixture principal, one nextest process.
#[tokio::test]
async fn provider_issued_token_and_grant_loss() {
    if oidc_issuer().is_none() {
        if require_env() {
            panic!("REMORA_OIDC_ISSUER is required when REMORA_IDENTITY_REQUIRE is set");
        }
        eprintln!("skipping provider-issued token test: REMORA_OIDC_ISSUER is unset");
        return;
    }
    ensure_remora_grants()
        .await
        .expect("ensure remora grants before mint");
    let provider = oidc_provider();
    let before = mint_provider_token().await.expect("mint with grants");
    let with_grants = match provider.authenticate_token(&before).await {
        Ok(identity) => identity,
        Err(error) => panic!(
            "authenticate provider token: {error}; claims={}",
            jwt_payload_preview(&before)
        ),
    };
    assert!(!with_grants.subject.is_empty());
    assert_eq!(with_grants.issuer, env_required("REMORA_OIDC_ISSUER"));
    assert!(
        with_grants.roles.administrator || !with_grants.roles.scopes.is_empty(),
        "expected remora roles on provider token, got {with_grants:?}; claims={}",
        jwt_payload_preview(&before)
    );
    let storage_role = std::env::var("REMORA_OIDC_STORAGE_ROLE").unwrap_or_default();
    if !storage_role.is_empty() {
        assert!(
            !with_grants.roles.scopes.is_empty(),
            "expected storage scope from fixture role {storage_role}"
        );
        assert_eq!(with_grants.roles.scopes[0].store.as_str(), "target-main");
        assert_eq!(with_grants.roles.scopes[0].bucket.as_str(), "target-bucket");
        assert_eq!(with_grants.roles.scopes[0].prefix.canonical(), "models");
    }

    revoke_remora_grants().await.expect("revoke remora grants");

    let after = mint_provider_token().await.expect("mint after revoke");
    let without_grants = provider
        .authenticate_token(&after)
        .await
        .expect("auth after revoke");
    assert!(
        !without_grants.roles.administrator,
        "new token must not be administrator after grant loss"
    );
    assert!(
        without_grants.roles.scopes.is_empty(),
        "new token must not carry storage scopes after grant loss"
    );
}

async fn mint_provider_token() -> Result<String, String> {
    match env_required("REMORA_OIDC_PROVIDER").as_str() {
        "keycloak" => mint_keycloak_password_token().await,
        "zitadel" => mint_zitadel_jwt_bearer_token().await,
        other => Err(format!("unsupported REMORA_OIDC_PROVIDER={other}")),
    }
}

async fn mint_keycloak_password_token() -> Result<String, String> {
    let issuer = env_required("REMORA_OIDC_ISSUER");
    let client_id = env_required("REMORA_OIDC_CLIENT_ID");
    let username = env_required("REMORA_OIDC_USERNAME");
    let password = env_required("REMORA_OIDC_PASSWORD");
    let token_url = format!(
        "{}/protocol/openid-connect/token",
        issuer.trim_end_matches('/')
    );
    let client = reqwest::Client::new();
    let response = client
        .post(token_url)
        .form(&[
            ("grant_type", "password"),
            ("client_id", client_id.as_str()),
            ("username", username.as_str()),
            ("password", password.as_str()),
            ("scope", "openid"),
        ])
        .send()
        .await
        .map_err(|error| error.to_string())?;
    if !response.status().is_success() {
        return Err(format!(
            "keycloak token endpoint: {} {}",
            response.status(),
            response.text().await.unwrap_or_default()
        ));
    }
    let body: Value = response.json().await.map_err(|error| error.to_string())?;
    body.get("access_token")
        .and_then(Value::as_str)
        .map(str::to_owned)
        .ok_or_else(|| "keycloak token response missing access_token".into())
}

/// Mint a provider-issued ZITADEL access token via JWT Profile (user application key).
async fn mint_zitadel_jwt_bearer_token() -> Result<String, String> {
    let issuer = env_required("REMORA_OIDC_ISSUER");
    let user_id = env_required("REMORA_OIDC_USER_ID");
    let project_id = env_required("REMORA_OIDC_PROJECT_ID");
    let key_b64 = env_required("REMORA_OIDC_USER_KEY_B64");
    let key_json = URL_SAFE_NO_PAD
        .decode(key_b64.as_bytes())
        .or_else(|_| base64::engine::general_purpose::STANDARD.decode(key_b64.as_bytes()))
        .map_err(|error| format!("decode user key: {error}"))?;
    let key: Value =
        serde_json::from_slice(&key_json).map_err(|error| format!("parse user key: {error}"))?;
    let key_id = key
        .get("keyId")
        .or_else(|| key.get("key_id"))
        .and_then(Value::as_str)
        .ok_or("user key missing keyId")?
        .to_owned();
    let private_key = key
        .get("key")
        .and_then(Value::as_str)
        .ok_or("user key missing private key PEM")?
        .replace("\\n", "\n");
    let issuer = issuer.trim_end_matches('/').to_owned();
    let token_url = format!("{issuer}/oauth/v2/token");
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map_err(|error| error.to_string())?
        .as_secs() as i64;
    // Prefer userId from the key document when present.
    let subject = key
        .get("userId")
        .or_else(|| key.get("userID"))
        .and_then(Value::as_str)
        .unwrap_or(user_id.as_str())
        .to_owned();
    let encoding = EncodingKey::from_rsa_pem(private_key.as_bytes()).map_err(|error| {
        format!(
            "user key pem: {error}; fields={:?}",
            key.as_object().map(|o| o.keys().collect::<Vec<_>>())
        )
    })?;
    // RFC 7523 assertion aud is the token endpoint; some builds also accept issuer.
    let audiences = [token_url.clone(), issuer.clone()];
    // Non-default port must appear in Host so issued iss matches REMORA_OIDC_ISSUER.
    let host = issuer
        .strip_prefix("http://")
        .or_else(|| issuer.strip_prefix("https://"))
        .unwrap_or(issuer.as_str());
    let scope = format!(
        "openid profile \
         urn:zitadel:iam:org:project:id:{project_id}:aud \
         urn:zitadel:iam:org:project:roles \
         urn:zitadel:iam:org:projects:roles \
         urn:zitadel:iam:org:project:{project_id}:roles"
    );
    let client = reqwest::Client::new();
    let mut last_error = String::new();
    for audience in audiences {
        let mut header = Header::new(Algorithm::RS256);
        header.kid = Some(key_id.clone());
        let assertion = match encode(
            &header,
            &json!({
                "iss": subject,
                "sub": subject,
                "aud": audience,
                "iat": now,
                "exp": now + 3600,
            }),
            &encoding,
        ) {
            Ok(value) => value,
            Err(error) => {
                last_error = format!("sign assertion: {error}");
                continue;
            }
        };
        let response = client
            .post(&token_url)
            .header("Host", host)
            .form(&[
                ("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"),
                ("assertion", assertion.as_str()),
                ("scope", scope.as_str()),
            ])
            .send()
            .await
            .map_err(|error| error.to_string())?;
        if response.status().is_success() {
            let body: Value = response.json().await.map_err(|error| error.to_string())?;
            // Prefer whichever JWT carries remora role keys; fall back to access_token.
            let access = body.get("access_token").and_then(Value::as_str);
            let id_token = body.get("id_token").and_then(Value::as_str);
            for token in [access, id_token].into_iter().flatten() {
                if jwt_payload_preview(token).contains("remora:v1:") {
                    return Ok(token.to_owned());
                }
            }
            if let Some(token) = access.or(id_token) {
                return Ok(token.to_owned());
            }
            return Err(format!(
                "zitadel token response missing tokens; body_keys={:?}",
                body.as_object().map(|o| o.keys().collect::<Vec<_>>())
            ));
        }
        last_error = format!(
            "aud={audience}: {} {}",
            response.status(),
            response.text().await.unwrap_or_default()
        );
    }
    Err(format!("zitadel jwt-bearer token failed: {last_error}"))
}

async fn ensure_remora_grants() -> Result<(), String> {
    match env_required("REMORA_OIDC_PROVIDER").as_str() {
        "keycloak" => ensure_keycloak_remora_client_roles().await,
        "zitadel" => {
            let storage_role = env_required("REMORA_OIDC_STORAGE_ROLE");
            set_zitadel_user_grant_roles(&["remora:v1:admin", storage_role.as_str()]).await
        }
        other => Err(format!("unsupported REMORA_OIDC_PROVIDER={other}")),
    }
}

async fn revoke_remora_grants() -> Result<(), String> {
    match env_required("REMORA_OIDC_PROVIDER").as_str() {
        "keycloak" => revoke_keycloak_remora_client_roles().await,
        "zitadel" => set_zitadel_user_grant_roles(&["test-user"]).await,
        other => Err(format!("unsupported REMORA_OIDC_PROVIDER={other}")),
    }
}

async fn revoke_keycloak_remora_client_roles() -> Result<(), String> {
    let (admin_token, realm, user_id, client_uuid, storage_role) =
        keycloak_admin_role_context().await?;
    let client = reqwest::Client::new();
    let roles_url = format!(
        "http://keycloak.remora.test:8080/admin/realms/{realm}/users/{user_id}/role-mappings/clients/{client_uuid}"
    );
    let mapped: Value = client
        .get(&roles_url)
        .bearer_auth(&admin_token)
        .send()
        .await
        .map_err(|error| error.to_string())?
        .json()
        .await
        .map_err(|error| error.to_string())?;
    let remora_roles: Vec<Value> = mapped
        .as_array()
        .cloned()
        .unwrap_or_default()
        .into_iter()
        .filter(|role| {
            role.get("name")
                .and_then(Value::as_str)
                .is_some_and(|name| {
                    name == "remora:v1:admin"
                        || name == storage_role
                        || name.starts_with("remora:v1:")
                })
        })
        .collect();
    if remora_roles.is_empty() {
        return Err("keycloak user has no remora client roles to revoke".into());
    }
    let delete = client
        .delete(&roles_url)
        .bearer_auth(&admin_token)
        .json(&remora_roles)
        .send()
        .await
        .map_err(|error| error.to_string())?;
    if !delete.status().is_success() {
        return Err(format!(
            "keycloak delete client roles: {} {}",
            delete.status(),
            delete.text().await.unwrap_or_default()
        ));
    }
    Ok(())
}

async fn set_zitadel_user_grant_roles(role_keys: &[&str]) -> Result<(), String> {
    let issuer = env_required("REMORA_OIDC_ISSUER");
    let org_id = env_required("REMORA_OIDC_TENANT");
    let user_id = env_required("REMORA_OIDC_USER_ID");
    let project_id = env_required("REMORA_OIDC_PROJECT_ID");
    let mgmt_token = env_required("REMORA_OIDC_MGMT_TOKEN");
    let client = reqwest::Client::new();

    let search_url = format!(
        "{}/management/v1/users/grants/_search",
        issuer.trim_end_matches('/')
    );
    let host = issuer
        .trim_end_matches('/')
        .strip_prefix("http://")
        .or_else(|| issuer.trim_end_matches('/').strip_prefix("https://"))
        .unwrap_or(issuer.trim_end_matches('/'));
    let search = client
        .post(search_url)
        .header("Host", host)
        .header("Authorization", format!("Bearer {mgmt_token}"))
        .header("Content-Type", "application/json")
        .header("x-zitadel-orgid", &org_id)
        .json(&json!({
            "queries": [{
                "userIdQuery": { "userId": user_id }
            }]
        }))
        .send()
        .await
        .map_err(|error| error.to_string())?;
    if !search.status().is_success() {
        return Err(format!(
            "zitadel search grants: {} {}",
            search.status(),
            search.text().await.unwrap_or_default()
        ));
    }
    let body: Value = search.json().await.map_err(|error| error.to_string())?;
    let grants = body
        .get("result")
        .and_then(Value::as_array)
        .cloned()
        .unwrap_or_default();
    let grant = grants
        .iter()
        .find(|row| {
            row.get("projectId")
                .and_then(Value::as_str)
                .is_some_and(|id| id == project_id)
        })
        .ok_or("zitadel user project grant not found")?;
    let grant_id = grant
        .get("id")
        .and_then(Value::as_str)
        .ok_or("zitadel grant missing id")?;

    let update_url = format!(
        "{}/management/v1/users/{user_id}/grants/{grant_id}",
        issuer.trim_end_matches('/')
    );
    let update = client
        .put(update_url)
        .header("Host", host)
        .header("Authorization", format!("Bearer {mgmt_token}"))
        .header("Content-Type", "application/json")
        .header("x-zitadel-orgid", &org_id)
        .json(&json!({ "roleKeys": role_keys }))
        .send()
        .await
        .map_err(|error| error.to_string())?;
    if !update.status().is_success() {
        let status = update.status();
        let body = update.text().await.unwrap_or_default();
        // Idempotent ensure: already-correct grants are rejected as unchanged.
        if status.as_u16() == 400 && body.contains("User grant has not been changed") {
            return Ok(());
        }
        return Err(format!("zitadel update grant: {status} {body}"));
    }
    Ok(())
}

async fn ensure_keycloak_remora_client_roles() -> Result<(), String> {
    let (admin_token, realm, user_id, client_uuid, storage_role) =
        keycloak_admin_role_context().await?;
    let client = reqwest::Client::new();
    let available_url = format!(
        "http://keycloak.remora.test:8080/admin/realms/{realm}/users/{user_id}/role-mappings/clients/{client_uuid}/available"
    );
    let available: Value = client
        .get(available_url)
        .bearer_auth(&admin_token)
        .send()
        .await
        .map_err(|error| error.to_string())?
        .json()
        .await
        .map_err(|error| error.to_string())?;
    let wanted: Vec<Value> = available
        .as_array()
        .cloned()
        .unwrap_or_default()
        .into_iter()
        .filter(|role| {
            role.get("name")
                .and_then(Value::as_str)
                .is_some_and(|name| name == "remora:v1:admin" || name == storage_role)
        })
        .collect();
    if wanted.is_empty() {
        // Roles already mapped (or missing from the client entirely).
        return Ok(());
    }
    let roles_url = format!(
        "http://keycloak.remora.test:8080/admin/realms/{realm}/users/{user_id}/role-mappings/clients/{client_uuid}"
    );
    let post = client
        .post(roles_url)
        .bearer_auth(&admin_token)
        .json(&wanted)
        .send()
        .await
        .map_err(|error| error.to_string())?;
    if !post.status().is_success() {
        return Err(format!(
            "keycloak assign client roles: {} {}",
            post.status(),
            post.text().await.unwrap_or_default()
        ));
    }
    Ok(())
}

async fn keycloak_admin_role_context() -> Result<(String, String, String, String, String), String> {
    let realm = env_required("REMORA_OIDC_TENANT");
    let client_id = env_required("REMORA_OIDC_CLIENT_ID");
    let username = env_required("REMORA_OIDC_USERNAME");
    let storage_role = env_required("REMORA_OIDC_STORAGE_ROLE");
    let client = reqwest::Client::new();

    let admin_token_url =
        "http://keycloak.remora.test:8080/realms/master/protocol/openid-connect/token";
    let admin_token_response = client
        .post(admin_token_url)
        .form(&[
            ("grant_type", "password"),
            ("client_id", "admin-cli"),
            ("username", "remora-admin"),
            ("password", "synthetic-keycloak"),
        ])
        .send()
        .await
        .map_err(|error| error.to_string())?;
    if !admin_token_response.status().is_success() {
        return Err(format!(
            "keycloak admin token: {} {}",
            admin_token_response.status(),
            admin_token_response.text().await.unwrap_or_default()
        ));
    }
    let admin_body: Value = admin_token_response
        .json()
        .await
        .map_err(|error| error.to_string())?;
    let admin_token = admin_body
        .get("access_token")
        .and_then(Value::as_str)
        .ok_or("keycloak admin token missing access_token")?
        .to_owned();

    let users_url = format!(
        "http://keycloak.remora.test:8080/admin/realms/{realm}/users?username={username}&exact=true"
    );
    let users: Value = client
        .get(users_url)
        .bearer_auth(&admin_token)
        .send()
        .await
        .map_err(|error| error.to_string())?
        .json()
        .await
        .map_err(|error| error.to_string())?;
    let user_id = users
        .as_array()
        .and_then(|rows| rows.first())
        .and_then(|row| row.get("id"))
        .and_then(Value::as_str)
        .ok_or("keycloak user not found")?
        .to_owned();

    let clients_url = format!(
        "http://keycloak.remora.test:8080/admin/realms/{realm}/clients?clientId={client_id}"
    );
    let clients: Value = client
        .get(clients_url)
        .bearer_auth(&admin_token)
        .send()
        .await
        .map_err(|error| error.to_string())?
        .json()
        .await
        .map_err(|error| error.to_string())?;
    let client_uuid = clients
        .as_array()
        .and_then(|rows| rows.first())
        .and_then(|row| row.get("id"))
        .and_then(Value::as_str)
        .ok_or("keycloak client not found")?
        .to_owned();

    Ok((admin_token, realm, user_id, client_uuid, storage_role))
}

fn jwt_payload_preview(token: &str) -> String {
    let parts: Vec<&str> = token.split('.').collect();
    if parts.len() != 3 {
        return format!("<not-a-jwt parts={} len={}>", parts.len(), token.len());
    }
    let Ok(bytes) = URL_SAFE_NO_PAD
        .decode(parts[1].as_bytes())
        .or_else(|_| base64::engine::general_purpose::STANDARD.decode(parts[1].as_bytes()))
    else {
        return "<payload-decode-failed>".into();
    };
    match String::from_utf8(bytes) {
        Ok(text) => text,
        Err(_) => "<payload-not-utf8>".into(),
    }
}

fn mint_unsigned_looking_token() -> String {
    let header = Header::new(Algorithm::HS256);
    let claims = json!({
        "iss": "https://attacker.example",
        "sub": "attacker",
        "aud": "remora-client",
        "exp": 4_102_444_800_i64
    });
    encode(
        &header,
        &claims,
        &EncodingKey::from_secret(b"not-the-fixture-secret"),
    )
    .expect("mint")
}