remora/crates/remora-api/tests/support/mod.rs

use base64::{
    Engine as _,
    engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD},
};
use jsonwebtoken::{Algorithm, EncodingKey, Header, encode};
use remora_identity::{OidcConfig, OidcIdentityProvider};
use reqwest::Client;
use serde_json::{Value, json};

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

fn keycloak_admin_url() -> String {
    std::env::var("REMORA_KEYCLOAK_ADMIN_URL")
        .unwrap_or_else(|_| "http://keycloak.remora.test:8080".to_owned())
        .trim_end_matches('/')
        .to_owned()
}

pub fn oidc_provider() -> OidcIdentityProvider {
    let provider = required_env("REMORA_OIDC_PROVIDER");
    let audience = if provider == "zitadel" {
        required_env("REMORA_OIDC_PROJECT_ID")
    } else {
        required_env("REMORA_OIDC_CLIENT_ID")
    };
    let mut config = OidcConfig::new(required_env("REMORA_OIDC_ISSUER"), audience).unwrap();
    if provider == "zitadel" {
        let project_id = required_env("REMORA_OIDC_PROJECT_ID");
        config = config
            .with_zitadel_project_roles()
            .with_project_role_claim(format!("urn:zitadel:iam:org:project:{project_id}:roles"))
            .with_resource_access_client(required_env("REMORA_OIDC_CLIENT_ID"));
    } else {
        config = config.with_resource_access_client(required_env("REMORA_OIDC_CLIENT_ID"));
    }
    OidcIdentityProvider::new(config).unwrap()
}

pub async fn mint_provider_token() -> Result<String, String> {
    match required_env("REMORA_OIDC_PROVIDER").as_str() {
        "keycloak" => mint_keycloak_token().await,
        "zitadel" => mint_zitadel_token().await,
        provider => Err(format!("unsupported OIDC provider {provider}")),
    }
}

pub async fn revoke_remora_grants() -> Result<(), String> {
    match required_env("REMORA_OIDC_PROVIDER").as_str() {
        "keycloak" => revoke_keycloak_remora_roles().await,
        "zitadel" => revoke_zitadel_remora_roles().await,
        provider => Err(format!("unsupported OIDC provider {provider}")),
    }
}

pub async fn restore_remora_grants() -> Result<(), String> {
    match required_env("REMORA_OIDC_PROVIDER").as_str() {
        "keycloak" => restore_keycloak_remora_roles().await,
        "zitadel" => restore_zitadel_remora_roles().await,
        provider => Err(format!("unsupported OIDC provider {provider}")),
    }
}

async fn mint_keycloak_token() -> Result<String, String> {
    let response = Client::new()
        .post(format!(
            "{}/protocol/openid-connect/token",
            required_env("REMORA_OIDC_ISSUER").trim_end_matches('/')
        ))
        .form(&[
            ("grant_type", "password"),
            ("client_id", required_env("REMORA_OIDC_CLIENT_ID").as_str()),
            ("username", required_env("REMORA_OIDC_USERNAME").as_str()),
            ("password", required_env("REMORA_OIDC_PASSWORD").as_str()),
            ("scope", "openid"),
        ])
        .send()
        .await
        .map_err(|error| error.to_string())?;
    token_from_response(response).await
}

async fn mint_zitadel_token() -> Result<String, String> {
    let issuer = required_env("REMORA_OIDC_ISSUER")
        .trim_end_matches('/')
        .to_owned();
    let token_url = format!("{issuer}/oauth/v2/token");
    let encoded_key = required_env("REMORA_OIDC_USER_KEY_B64");
    let key_json = URL_SAFE_NO_PAD
        .decode(encoded_key.as_bytes())
        .or_else(|_| STANDARD.decode(encoded_key.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_else(|| "user key missing keyId".to_owned())?
        .to_owned();
    let subject = key
        .get("userId")
        .or_else(|| key.get("userID"))
        .and_then(Value::as_str)
        .map(str::to_owned)
        .unwrap_or_else(|| required_env("REMORA_OIDC_USER_ID"));
    let private_key = key["key"]
        .as_str()
        .ok_or_else(|| "user key missing private key".to_owned())?
        .replace("\\n", "\n");
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map_err(|error| error.to_string())?
        .as_secs() as i64;
    let encoding =
        EncodingKey::from_rsa_pem(private_key.as_bytes()).map_err(|error| error.to_string())?;
    let project_id = required_env("REMORA_OIDC_PROJECT_ID");
    let host = issuer
        .strip_prefix("http://")
        .or_else(|| issuer.strip_prefix("https://"))
        .unwrap_or(&issuer);
    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 = Client::new();
    let mut last_error = String::new();
    for audience in [&token_url, &issuer] {
        let mut header = Header::new(Algorithm::RS256);
        header.kid = Some(key_id.clone());
        let assertion = encode(
            &header,
            &json!({
                "iss": subject,
                "sub": subject,
                "aud": audience,
                "iat": now,
                "exp": now + 3600
            }),
            &encoding,
        )
        .map_err(|error| error.to_string())?;
        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() {
            return token_from_response(response).await;
        }
        last_error = format!(
            "token endpoint: {} {}",
            response.status(),
            response.text().await.unwrap_or_default()
        );
    }
    Err(last_error)
}

async fn token_from_response(response: reqwest::Response) -> Result<String, String> {
    if !response.status().is_success() {
        return Err(format!(
            "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(|| "token response missing access_token".into())
}

async fn revoke_keycloak_remora_roles() -> Result<(), String> {
    let realm = required_env("REMORA_OIDC_TENANT");
    let client_id = required_env("REMORA_OIDC_CLIENT_ID");
    let username = required_env("REMORA_OIDC_USERNAME");
    let keycloak_url = keycloak_admin_url();
    let client = Client::new();
    let admin_response = client
        .post(format!(
            "{keycloak_url}/realms/master/protocol/openid-connect/token"
        ))
        .form(&[
            ("grant_type", "password"),
            ("client_id", "admin-cli"),
            ("username", "remora-admin"),
            ("password", "synthetic-keycloak"),
        ])
        .send()
        .await
        .map_err(|error| error.to_string())?;
    let admin_body: Value = admin_response
        .json()
        .await
        .map_err(|error| error.to_string())?;
    let admin_token = admin_body["access_token"]
        .as_str()
        .ok_or_else(|| "keycloak admin token missing".to_owned())?;
    let users: Value = client
        .get(format!(
            "{keycloak_url}/admin/realms/{realm}/users?username={username}&exact=true"
        ))
        .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["id"].as_str())
        .ok_or_else(|| "keycloak user missing".to_owned())?;
    let clients: Value = client
        .get(format!(
            "{keycloak_url}/admin/realms/{realm}/clients?clientId={client_id}"
        ))
        .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["id"].as_str())
        .ok_or_else(|| "keycloak client missing".to_owned())?;
    let roles_url = format!(
        "{keycloak_url}/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 roles = mapped
        .as_array()
        .cloned()
        .unwrap_or_default()
        .into_iter()
        .filter(|role| {
            role["name"]
                .as_str()
                .is_some_and(|name| name.starts_with("remora:v1:"))
        })
        .collect::<Vec<_>>();
    let response = client
        .delete(roles_url)
        .bearer_auth(admin_token)
        .json(&roles)
        .send()
        .await
        .map_err(|error| error.to_string())?;
    if response.status().is_success() {
        Ok(())
    } else {
        Err(format!(
            "keycloak role revocation failed: {}",
            response.status()
        ))
    }
}

async fn restore_keycloak_remora_roles() -> Result<(), String> {
    let realm = required_env("REMORA_OIDC_TENANT");
    let client_id = required_env("REMORA_OIDC_CLIENT_ID");
    let username = required_env("REMORA_OIDC_USERNAME");
    let keycloak_url = keycloak_admin_url();
    let client = Client::new();
    let admin_response = client
        .post(format!(
            "{keycloak_url}/realms/master/protocol/openid-connect/token"
        ))
        .form(&[
            ("grant_type", "password"),
            ("client_id", "admin-cli"),
            ("username", "remora-admin"),
            ("password", "synthetic-keycloak"),
        ])
        .send()
        .await
        .map_err(|error| error.to_string())?;
    let admin_body: Value = admin_response
        .json()
        .await
        .map_err(|error| error.to_string())?;
    let admin_token = admin_body["access_token"]
        .as_str()
        .ok_or_else(|| "keycloak admin token missing".to_owned())?;
    let users: Value = client
        .get(format!(
            "{keycloak_url}/admin/realms/{realm}/users?username={username}&exact=true"
        ))
        .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["id"].as_str())
        .ok_or_else(|| "keycloak user missing".to_owned())?;
    let clients: Value = client
        .get(format!(
            "{keycloak_url}/admin/realms/{realm}/clients?clientId={client_id}"
        ))
        .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["id"].as_str())
        .ok_or_else(|| "keycloak client missing".to_owned())?;
    let roles: Value = client
        .get(format!(
            "{keycloak_url}/admin/realms/{realm}/clients/{client_uuid}/roles"
        ))
        .bearer_auth(admin_token)
        .send()
        .await
        .map_err(|error| error.to_string())?
        .json()
        .await
        .map_err(|error| error.to_string())?;
    let storage_role = required_env("REMORA_OIDC_STORAGE_ROLE");
    let roles = roles
        .as_array()
        .cloned()
        .unwrap_or_default()
        .into_iter()
        .filter(|role| {
            role["name"]
                .as_str()
                .is_some_and(|name| name == "remora:v1:admin" || name == storage_role)
        })
        .collect::<Vec<_>>();
    let response = client
        .post(format!(
            "{keycloak_url}/admin/realms/{realm}/users/{user_id}/role-mappings/clients/{client_uuid}"
        ))
        .bearer_auth(admin_token)
        .json(&roles)
        .send()
        .await
        .map_err(|error| error.to_string())?;
    if response.status().is_success() {
        Ok(())
    } else {
        Err(format!(
            "keycloak role restore failed: {}",
            response.status()
        ))
    }
}

async fn revoke_zitadel_remora_roles() -> Result<(), String> {
    let issuer = required_env("REMORA_OIDC_ISSUER");
    let organization = required_env("REMORA_OIDC_TENANT");
    let user_id = required_env("REMORA_OIDC_USER_ID");
    let project_id = required_env("REMORA_OIDC_PROJECT_ID");
    let token = required_env("REMORA_OIDC_MGMT_TOKEN");
    let host = issuer
        .strip_prefix("http://")
        .or_else(|| issuer.strip_prefix("https://"))
        .unwrap_or(&issuer);
    let client = Client::new();
    let grants: Value = client
        .post(format!("{issuer}/management/v1/users/grants/_search"))
        .header("Host", host)
        .bearer_auth(&token)
        .header("x-zitadel-orgid", &organization)
        .json(&json!({ "queries": [{ "userIdQuery": { "userId": user_id } }] }))
        .send()
        .await
        .map_err(|error| error.to_string())?
        .json()
        .await
        .map_err(|error| error.to_string())?;
    let grant_id = grants["result"]
        .as_array()
        .and_then(|rows| {
            rows.iter().find(|row| {
                row["projectId"]
                    .as_str()
                    .is_some_and(|value| value == project_id)
            })
        })
        .and_then(|row| row["id"].as_str())
        .ok_or_else(|| "zitadel project grant missing".to_owned())?;
    let response = client
        .put(format!(
            "{issuer}/management/v1/users/{user_id}/grants/{grant_id}"
        ))
        .header("Host", host)
        .bearer_auth(token)
        .header("x-zitadel-orgid", organization)
        .json(&json!({ "roleKeys": ["test-user"] }))
        .send()
        .await
        .map_err(|error| error.to_string())?;
    if response.status().is_success() {
        Ok(())
    } else {
        Err(format!(
            "zitadel role revocation failed: {}",
            response.status()
        ))
    }
}

async fn restore_zitadel_remora_roles() -> Result<(), String> {
    let issuer = required_env("REMORA_OIDC_ISSUER");
    let organization = required_env("REMORA_OIDC_TENANT");
    let user_id = required_env("REMORA_OIDC_USER_ID");
    let project_id = required_env("REMORA_OIDC_PROJECT_ID");
    let token = required_env("REMORA_OIDC_MGMT_TOKEN");
    let host = issuer
        .strip_prefix("http://")
        .or_else(|| issuer.strip_prefix("https://"))
        .unwrap_or(&issuer);
    let client = Client::new();
    let grants: Value = client
        .post(format!("{issuer}/management/v1/users/grants/_search"))
        .header("Host", host)
        .bearer_auth(&token)
        .header("x-zitadel-orgid", &organization)
        .json(&json!({ "queries": [{ "userIdQuery": { "userId": user_id } }] }))
        .send()
        .await
        .map_err(|error| error.to_string())?
        .json()
        .await
        .map_err(|error| error.to_string())?;
    let grant_id = grants["result"]
        .as_array()
        .and_then(|rows| {
            rows.iter().find(|row| {
                row["projectId"]
                    .as_str()
                    .is_some_and(|value| value == project_id)
            })
        })
        .and_then(|row| row["id"].as_str())
        .ok_or_else(|| "zitadel project grant missing".to_owned())?;
    let response = client
        .put(format!(
            "{issuer}/management/v1/users/{user_id}/grants/{grant_id}"
        ))
        .header("Host", host)
        .bearer_auth(token)
        .header("x-zitadel-orgid", organization)
        .json(&json!({
            "roleKeys": [
                "test-user",
                "remora:v1:admin",
                required_env("REMORA_OIDC_STORAGE_ROLE")
            ]
        }))
        .send()
        .await
        .map_err(|error| error.to_string())?;
    if response.status().is_success() {
        Ok(())
    } else {
        Err(format!(
            "zitadel role restore failed: {}",
            response.status()
        ))
    }
}