cipherstash-client 0.40.0

#![cfg(feature = "test-utils")]
#![allow(unused)]

pub(crate) use cipherstash_client::{
    cts_client::{OidcProvider, OidcVendor},
    encryption::ScopedCipher,
    zerokms::{ClientKey, EncryptPayload, ZeroKMSBuilder, ZeroKMSWithClientKey},
    AuthError, CtsClient, ZeroKMS,
};

#[cfg(feature = "tokio")]
pub(crate) use cipherstash_client::eql::*;

pub(crate) use cipherstash_config::{column::Index, ColumnConfig};
pub(crate) use cts_common::{claims::Role, protocol::CreateWorkspaceResponse, Region, WorkspaceId};
pub(crate) use fake::{Fake, Faker};
pub(crate) use mock_auth_server::{vendor::MockClerkClaims, MockClaimsBuilder, MockToken};
pub(crate) use serde_json::json;
pub(crate) use stack_auth::{
    AccessKeyStrategy, DeviceSessionStrategy, SecretToken, StaticTokenStrategy, Token,
};
pub(crate) use std::{
    borrow::Cow,
    path::Path,
    sync::{Arc, LazyLock},
    time::{SystemTime, UNIX_EPOCH},
};
pub(crate) use tempfile::{tempdir, TempDir};
pub(crate) use uuid::Uuid;
pub(crate) use zerokms_protocol::{Context, Keyset};

pub(crate) static CONFIG_DIR: LazyLock<TempDir> = LazyLock::new(|| {
    let temp_dir = tempdir().unwrap();
    std::env::set_var("CS_CONFIG_PATH", temp_dir.path());
    temp_dir
});

pub(crate) fn cts_base_url() -> url::Url {
    std::env::var("CS_CTS_HOST")
        .expect("CS_CTS_HOST should be set")
        .parse()
        .expect("CS_CTS_HOST should be a valid URL")
}

pub(crate) fn build_mock_token() -> MockToken {
    let iss = format!(
        "http://{}:3030",
        std::env::var("MOCK_AUTH_SERVER_HOST").unwrap()
    );

    MockClaimsBuilder::new()
        .sub("cs-client-e2e-user")
        .iss(Some(iss))
        .o(Some(mock_auth_server::token::Org {
            id: "test-org-id".to_string(),
            rol: mock_auth_server::token::OrgRole::Admin,
            slg: "test-org-slg".to_string(),
            name: "Test Org".to_string(),
        }))
        .exp(
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs()
                + 60 * 60,
        )
        .aud(std::env::var("CTS__AUTH__AUDIENCE").ok())
        .build_token::<MockClerkClaims>()
        .unwrap()
}

pub(crate) fn build_cts_client(token: MockToken) -> CtsClient<StaticTokenStrategy> {
    CtsClient::new(StaticTokenStrategy::new(token.0)).with_base_url(cts_base_url())
}

pub(crate) async fn create_workspace(cts_client: &CtsClient<StaticTokenStrategy>) -> WorkspaceId {
    let CreateWorkspaceResponse { id: workspace_id } = cts_client
        .create_workspace(
            Faker.fake::<String>().as_str(),
            Region::new("ap-southeast-2.aws").unwrap(),
        )
        .await
        .unwrap();

    workspace_id
}

pub(crate) async fn create_access_key(
    cts_client: &CtsClient<StaticTokenStrategy>,
    workspace_id: WorkspaceId,
) -> String {
    cts_client
        .create_access_key("test access key", workspace_id, Role::Admin)
        .await
        .unwrap()
}

/// Build an [`AccessKeyStrategy`] for the given (just-minted) access key
/// and the workspace it was created against.
///
/// The CRN passed to the strategy is constructed at runtime from the
/// fresh workspace — `AccessKeyStrategy` now verifies the issued
/// token's `workspace` claim against this CRN, so reading a static CRN
/// from the environment would fail every e2e test the moment
/// `create_workspace()` mints a new workspace ID.
///
/// The CTS base URL is resolved from `CS_CTS_HOST` (automatically
/// checked inside `AccessKeyStrategy::new`).
pub(crate) fn build_access_key_strategy(
    access_key: &str,
    workspace_id: WorkspaceId,
) -> AccessKeyStrategy {
    // Region must match the one passed to `create_workspace` above.
    let region = Region::new("ap-southeast-2.aws").unwrap();
    let crn = cts_common::Crn::new(region, workspace_id);
    let key: stack_auth::AccessKey = access_key
        .parse()
        .expect("access key should be a valid CSAK key");
    AccessKeyStrategy::new(crn, key).expect("failed to build AccessKeyStrategy")
}

/// Build a [`ZeroKMS`] client for the given auth strategy.
///
/// Reads `CS_TEST_ZEROKMS_HOST` from the environment to determine the ZeroKMS base URL.
pub(crate) fn build_zerokms<S>(strategy: S) -> ZeroKMS<S>
where
    S: Send + Sync + 'static,
    for<'a> &'a S: stack_auth::AuthStrategy,
{
    let base_url: url::Url = std::env::var("CS_TEST_ZEROKMS_HOST")
        .expect("var CS_TEST_ZEROKMS_HOST should be set")
        .parse()
        .expect("CS_TEST_ZEROKMS_HOST should be a valid URL");
    ZeroKMSBuilder::new(strategy)
        .with_base_url(base_url)
        .build()
        .expect("failed to build zero_kms client")
}

/// Build a [`ZeroKMSWithClientKey`] client for the given auth strategy and client key.
///
/// Reads `CS_TEST_ZEROKMS_HOST` from the environment to determine the ZeroKMS base URL.
pub(crate) fn build_zerokms_with_client_key<S>(
    strategy: S,
    client_key: &ClientKey,
) -> ZeroKMSWithClientKey<S>
where
    S: Send + Sync + 'static,
    for<'a> &'a S: stack_auth::AuthStrategy,
{
    let base_url: url::Url = std::env::var("CS_TEST_ZEROKMS_HOST")
        .expect("var CS_TEST_ZEROKMS_HOST should be set")
        .parse()
        .expect("CS_TEST_ZEROKMS_HOST should be a valid URL");
    ZeroKMSBuilder::new(strategy)
        .with_base_url(base_url)
        .with_client_key(client_key.clone())
        .build()
        .expect("failed to build zero_kms client with client key")
}

/// Create a keyset on ZeroKMS using the given client.
pub(crate) async fn create_keyset<C>(zero_kms: &ZeroKMS<C>) -> Keyset
where
    C: Send + Sync + 'static,
    for<'a> &'a C: stack_auth::AuthStrategy,
{
    let keyset_name: String = Faker.fake();
    zero_kms
        .create_keyset(&keyset_name, "test-description")
        .await
        .expect("failed to create keyset for zero_kms client")
}

/// Create a client key on ZeroKMS using the given client and keyset.
pub(crate) async fn create_client_key<C>(zero_kms: &ZeroKMS<C>, keyset: &Keyset) -> ClientKey
where
    C: Send + Sync + 'static,
    for<'a> &'a C: stack_auth::AuthStrategy,
{
    let client = zero_kms
        .create_client("test-client", "test-description", Some(keyset.id))
        .await
        .expect("failed to create client for zero_kms client");

    ClientKey::from_bytes(client.id, &client.client_key[..])
        .expect("failed to create proxy key set from client response")
}

/// Set up a ZeroKMS client from scratch: creates a workspace, access key, and builds the client.
///
/// Returns `(zero_kms_client, access_key, workspace_id)` so callers can create additional
/// clients (e.g. with a client key) from the same workspace.
pub(crate) async fn setup_zerokms_client() -> (ZeroKMS<AccessKeyStrategy>, String, WorkspaceId) {
    let _temp_config_dir = &*CONFIG_DIR;

    let token = build_mock_token();
    let cts_client = build_cts_client(token);
    let workspace_id = create_workspace(&cts_client).await;
    let access_key = create_access_key(&cts_client, workspace_id).await;

    let strategy = build_access_key_strategy(&access_key, workspace_id);
    let client = build_zerokms(strategy);

    (client, access_key, workspace_id)
}

/// Set up a ZeroKMS client with a client key: creates workspace, access key, keyset, client key,
/// and builds a `ZeroKMSWithClientKey`.
///
/// Returns `(zerokms_with_client_key, zerokms_admin, keyset, client_key)` — the admin client is
/// useful for keyset management operations (disable, enable, grant).
pub(crate) async fn setup_zerokms_client_with_client_key() -> (
    cipherstash_client::zerokms::ZeroKMSWithClientKey<AccessKeyStrategy>,
    ZeroKMS<AccessKeyStrategy>,
    Keyset,
    ClientKey,
) {
    let (zero_kms, access_key, workspace_id) = setup_zerokms_client().await;

    let keyset = create_keyset(&zero_kms).await;
    let client_key = create_client_key(&zero_kms, &keyset).await;

    let strategy = build_access_key_strategy(&access_key, workspace_id);
    let zerokms_client = build_zerokms_with_client_key(strategy, &client_key);

    (zerokms_client, zero_kms, keyset, client_key)
}

/// Exchange an access key for a CTS service token and return a workspace-scoped [`CtsClient`].
///
/// The service token's JWT claims include the workspace ID, so no `x-cs-workspace-id`
/// header is needed for workspace-scoped operations.
pub(crate) async fn build_service_cts_client(access_key: &str) -> CtsClient<StaticTokenStrategy> {
    #[derive(serde::Deserialize)]
    #[serde(rename_all = "camelCase")]
    struct AuthorizeResponse {
        access_token: String,
    }

    let base_url = cts_base_url();
    let resp: AuthorizeResponse = reqwest::Client::new()
        .post(base_url.join("/api/authorize").unwrap())
        .json(&json!({ "accessKey": access_key }))
        .send()
        .await
        .expect("authorize request failed")
        .error_for_status()
        .expect("authorize returned error")
        .json()
        .await
        .expect("failed to parse authorize response");

    CtsClient::new(StaticTokenStrategy::new(resp.access_token)).with_base_url(base_url)
}

pub(crate) async fn create_oidc_provider(
    cts_client: &CtsClient<StaticTokenStrategy>,
) -> OidcProvider {
    cts_client
        .create_oidc_provider(
            &url::Url::parse("http://mock-auth-server2:3030").unwrap(),
            OidcVendor::Auth0,
        )
        .await
        .unwrap()
}

pub(crate) fn build_orgless_token(
    oidc_provider: &OidcProvider,
    workspace_id: WorkspaceId,
) -> MockToken {
    let mut builder = MockClaimsBuilder::new()
        .sub("cs-client-e2e-user")
        .azp(format!("OIDC|{}", oidc_provider.id))
        .o(None::<mock_auth_server::token::Org>)
        .iss(Some("http://mock-auth-server2:3030"))
        .scopes(vec![
            "offline_access".to_string(),
            "dataset:create".to_string(),
            "client:create".to_string(),
        ])
        .exp(
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs()
                + 60 * 60,
        )
        .aud(std::env::var("ZEROKMS__IDP__AUDIENCE").ok());

    builder
        .claims
        .insert("workspace".into(), json!(workspace_id.to_string()));

    builder.build_token::<MockClerkClaims>().unwrap()
}

/// Build a [`DeviceSessionStrategy`] from a [`MockToken`].
///
/// Constructs a `stack_auth::Token` from the raw JWT string and configures the
/// strategy with the CTS base URL from the environment.
pub(crate) fn build_device_session_strategy(mock_token: &MockToken) -> DeviceSessionStrategy {
    let region =
        Region::new(&std::env::var("CS_REGION").unwrap_or("ap-southeast-2.aws".into())).unwrap();

    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs();

    let token: Token = serde_json::from_value(json!({
        "access_token": mock_token.0,
        "token_type": "Bearer",
        "expires_at": now + 3600,
    }))
    .expect("failed to deserialize mock token into stack_auth::Token");

    let base_url = cts_base_url();

    DeviceSessionStrategy::with_token(region, "test-client", token)
        .base_url(base_url)
        .build()
        .expect("failed to build DeviceSessionStrategy")
}

/// Federate an OIDC token through CTS to obtain a CTS-signed service token.
///
/// Sends the OIDC token to CTS's `/api/authorize` endpoint, which validates it
/// against the registered OIDC provider and returns a service token that ZeroKMS trusts.
pub(crate) async fn federate_oidc_token(
    mock_token: &MockToken,
    workspace_id: WorkspaceId,
) -> StaticTokenStrategy {
    #[derive(serde::Deserialize)]
    #[serde(rename_all = "camelCase")]
    struct AuthorizeResponse {
        access_token: String,
    }

    let base_url = cts_base_url();
    let resp: AuthorizeResponse = reqwest::Client::new()
        .post(base_url.join("/api/authorize").unwrap())
        .json(&json!({
            "oidcToken": mock_token.0,
            "workspaceId": workspace_id.to_string(),
        }))
        .send()
        .await
        .expect("authorize request failed")
        .error_for_status()
        .expect("authorize returned error")
        .json()
        .await
        .expect("failed to parse authorize response");

    StaticTokenStrategy::new(resp.access_token)
}