cipherstash-client 0.39.1

#![cfg(all(feature = "test-utils", feature = "tokio"))]

mod prelude;

use prelude::*;
use zerokms_protocol::IdentifiedBy;

async fn create_scoped_cipher(
    access_key: &str,
    workspace_id: WorkspaceId,
    client_key: &ClientKey,
    keyset_id: Option<Uuid>,
) -> Arc<ScopedCipher<AccessKeyStrategy>> {
    let strategy = build_access_key_strategy(access_key, workspace_id);
    let client = Arc::new(build_zerokms_with_client_key(strategy, client_key));
    let cipher = ScopedCipher::init(client, keyset_id.map(IdentifiedBy::Uuid))
        .await
        .expect("failed to create scoped cipher");
    Arc::new(cipher)
}

fn create_column_config_with_indexes(indexes: Vec<Index>) -> ColumnConfig {
    let mut config = ColumnConfig::build("test_column");

    for index in indexes {
        config.indexes.push(index);
    }

    config
}

#[tokio::test]
#[ignore = "e2e"]
async fn test_encrypt_and_decrypt_eql_basic() {
    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 zero_kms = build_zerokms(strategy);

    let keyset = create_keyset(&zero_kms).await;
    let client_key = create_client_key(&zero_kms, &keyset).await;
    let cipher =
        create_scoped_cipher(&access_key, workspace_id, &client_key, Some(keyset.id)).await;

    // Create test data
    let identifier = Identifier::new("users", "email");
    let plaintext_value = cipherstash_client::encryption::Plaintext::from("test@example.com");
    let column_config = create_column_config_with_indexes(vec![Index::new_match()]);

    let prepared = PreparedPlaintext::new(
        Cow::Owned(column_config),
        identifier.clone(),
        plaintext_value.clone(),
        EqlOperation::Store,
    );

    // Encrypt
    let encrypt_opts = EqlEncryptOpts::default();
    let encrypted = encrypt_eql(cipher.clone(), vec![prepared], &encrypt_opts)
        .await
        .expect("failed to encrypt");

    assert_eq!(encrypted.len(), 1);
    let stored = expect_store(&encrypted[0]);
    let payload = expect_encrypted(stored);
    assert_eq!(payload.identifier, identifier);
    assert_eq!(payload.version, EQL_SCHEMA_VERSION);

    // Decrypt
    let decrypt_opts = EqlDecryptOpts::default();
    let stored_ciphertexts: Vec<EqlCiphertext> = encrypted.into_iter().map(into_store).collect();
    let decrypted = decrypt_eql(cipher.clone(), stored_ciphertexts, &decrypt_opts)
        .await
        .expect("failed to decrypt");

    assert_eq!(decrypted.len(), 1);
    assert_eq!(
        String::try_from(decrypted[0].clone()).unwrap(),
        "test@example.com"
    );
}

#[tokio::test]
#[ignore = "e2e"]
async fn test_encrypt_eql_with_multiple_index_types() {
    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 zero_kms = build_zerokms(strategy);

    let keyset = create_keyset(&zero_kms).await;
    let client_key = create_client_key(&zero_kms, &keyset).await;
    let cipher =
        create_scoped_cipher(&access_key, workspace_id, &client_key, Some(keyset.id)).await;

    // Create column config with multiple index types
    let identifier = Identifier::new("cms", "doc_text");
    let plaintext_value = cipherstash_client::encryption::Plaintext::from("the quick brown fox");
    let column_config = create_column_config_with_indexes(vec![
        Index::new_match(),
        Index::new_ore(),
        Index::new_unique(),
    ]);

    let prepared = PreparedPlaintext::new(
        Cow::Owned(column_config),
        identifier.clone(),
        plaintext_value,
        EqlOperation::Store,
    );

    // Encrypt
    let encrypt_opts = EqlEncryptOpts::default();
    let encrypted = encrypt_eql(cipher.clone(), vec![prepared], &encrypt_opts)
        .await
        .expect("failed to encrypt");

    assert_eq!(encrypted.len(), 1);

    // Verify indexes were generated
    let payload = expect_encrypted(expect_store(&encrypted[0]));
    assert!(
        payload.bloom_filter.is_some(),
        "Match index should be present"
    );
    assert!(
        payload.ore_block_u64_8_256.is_some(),
        "ORE index should be present"
    );
    assert!(payload.hmac_256.is_some(), "Unique index should be present");
}

#[tokio::test]
#[ignore = "e2e"]
async fn test_encrypt_eql_in_query_mode() {
    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 zero_kms = build_zerokms(strategy);

    let keyset = create_keyset(&zero_kms).await;
    let client_key = create_client_key(&zero_kms, &keyset).await;
    let cipher =
        create_scoped_cipher(&access_key, workspace_id, &client_key, Some(keyset.id)).await;

    // Create query mode encryption
    let identifier = Identifier::new("users", "name");
    let plaintext_value = cipherstash_client::encryption::Plaintext::from("Alice");
    let match_index = Index::new_match();
    let column_config = create_column_config_with_indexes(vec![match_index.clone()]);

    let prepared = PreparedPlaintext::new(
        Cow::Owned(column_config),
        identifier.clone(),
        plaintext_value,
        EqlOperation::Query(
            &match_index.index_type,
            cipherstash_client::encryption::QueryOp::Default,
        ),
    );

    // Encrypt in query mode
    let encrypt_opts = EqlEncryptOpts::default();
    let encrypted = encrypt_eql(cipher.clone(), vec![prepared], &encrypt_opts)
        .await
        .expect("failed to encrypt in query mode");

    assert_eq!(encrypted.len(), 1);
    let query = expect_query(&encrypted[0]);
    let encrypted_query = expect_query_encrypted(query);
    assert_eq!(encrypted_query.identifier, identifier);

    // In query mode there's no ciphertext field — just the matching index term.
    assert!(
        matches!(encrypted_query.term, RootQueryTerm::BloomFilter { .. }),
        "Query mode should produce a bloom-filter root query term, got {:?}",
        encrypted_query.term
    );
}

#[tokio::test]
#[ignore = "e2e"]
async fn test_decrypt_eql_with_wrong_keyset() {
    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 zero_kms = build_zerokms(strategy);

    // Create first keyset and encrypt data
    let keyset1 = create_keyset(&zero_kms).await;
    let client_key1 = create_client_key(&zero_kms, &keyset1).await;
    let cipher1 =
        create_scoped_cipher(&access_key, workspace_id, &client_key1, Some(keyset1.id)).await;

    let identifier = Identifier::new("users", "secret");
    let plaintext_value = cipherstash_client::encryption::Plaintext::from("secret_data");
    let column_config = create_column_config_with_indexes(vec![Index::new_unique()]);

    let prepared = PreparedPlaintext::new(
        Cow::Owned(column_config),
        identifier.clone(),
        plaintext_value,
        EqlOperation::Store,
    );

    let encrypt_opts = EqlEncryptOpts::default();
    let encrypted = encrypt_eql(cipher1.clone(), vec![prepared], &encrypt_opts)
        .await
        .expect("failed to encrypt");

    // Create second keyset and try to decrypt with it
    let keyset2 = create_keyset(&zero_kms).await;
    let client_key2 = create_client_key(&zero_kms, &keyset2).await;
    let cipher2 =
        create_scoped_cipher(&access_key, workspace_id, &client_key2, Some(keyset2.id)).await;

    // Try to decrypt with wrong keyset (should fail)
    let decrypt_opts = EqlDecryptOpts {
        keyset_id: Some(keyset2.id),
        ..Default::default()
    };
    let stored: Vec<EqlCiphertext> = encrypted.into_iter().map(into_store).collect();
    let result = decrypt_eql(cipher2.clone(), stored, &decrypt_opts).await;

    assert!(result.is_err(), "Decryption with wrong keyset should fail");

    // Verify error message mentions the keyset
    let err = result.unwrap_err();
    let err_str = err.to_string();
    assert!(
        err_str.contains("decrypt") || err_str.contains("keyset"),
        "Error should mention decryption or keyset failure: {err_str}"
    );
}

#[tokio::test]
#[ignore = "e2e"]
async fn test_encrypt_eql_multiple_plaintexts() {
    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 zero_kms = build_zerokms(strategy);

    let keyset = create_keyset(&zero_kms).await;
    let client_key = create_client_key(&zero_kms, &keyset).await;
    let cipher =
        create_scoped_cipher(&access_key, workspace_id, &client_key, Some(keyset.id)).await;

    // Create multiple plaintexts
    let column_config = create_column_config_with_indexes(vec![Index::new_match()]);

    let prepared_plaintexts = vec![
        PreparedPlaintext::new(
            Cow::Borrowed(&column_config),
            Identifier::new("users", "email"),
            cipherstash_client::encryption::Plaintext::from("alice@example.com"),
            EqlOperation::Store,
        ),
        PreparedPlaintext::new(
            Cow::Borrowed(&column_config),
            Identifier::new("users", "email"),
            cipherstash_client::encryption::Plaintext::from("bob@example.com"),
            EqlOperation::Store,
        ),
        PreparedPlaintext::new(
            Cow::Borrowed(&column_config),
            Identifier::new("users", "email"),
            cipherstash_client::encryption::Plaintext::from("charlie@example.com"),
            EqlOperation::Store,
        ),
    ];

    // Encrypt multiple values
    let encrypt_opts = EqlEncryptOpts::default();
    let encrypted = encrypt_eql(cipher.clone(), prepared_plaintexts, &encrypt_opts)
        .await
        .expect("failed to encrypt multiple plaintexts");

    assert_eq!(encrypted.len(), 3);

    // Verify all were encrypted
    for eql in &encrypted {
        let payload = expect_encrypted(expect_store(eql));
        assert_eq!(payload.identifier.table, "users");
        assert_eq!(payload.identifier.column, "email");
    }

    // Decrypt all
    let decrypt_opts = EqlDecryptOpts::default();
    let stored: Vec<EqlCiphertext> = encrypted.into_iter().map(into_store).collect();
    let decrypted = decrypt_eql(cipher.clone(), stored, &decrypt_opts)
        .await
        .expect("failed to decrypt multiple values");

    assert_eq!(decrypted.len(), 3);
    assert_eq!(
        String::try_from(decrypted[0].clone()).unwrap(),
        "alice@example.com"
    );
    assert_eq!(
        String::try_from(decrypted[1].clone()).unwrap(),
        "bob@example.com"
    );
    assert_eq!(
        String::try_from(decrypted[2].clone()).unwrap(),
        "charlie@example.com"
    );
}

#[tokio::test]
#[ignore = "e2e"]
async fn test_encrypt_json_document_with_stevec_index() {
    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 zero_kms = build_zerokms(strategy);

    let keyset = create_keyset(&zero_kms).await;
    let client_key = create_client_key(&zero_kms, &keyset).await;
    let cipher =
        create_scoped_cipher(&access_key, workspace_id, &client_key, Some(keyset.id)).await;

    // Create a JSON document with data suitable for testing all SteVec query types
    let json_document = json!({
        "user": {
            "name": "Alice Johnson",
            "email": "alice@example.com",
            "age": 28,
            "department": "Engineering"
        },
        "metrics": {
            "login_count": 42,
            "score": 95.5
        },
        "tags": ["premium", "active", "verified"]
    });

    let identifier = Identifier::new("users", "profile_data");
    let plaintext_value =
        cipherstash_client::encryption::Plaintext::Json(Some(json_document.clone()));

    // Create SteVec index configuration
    let ste_vec_index = Index::new(cipherstash_config::column::IndexType::SteVec {
        prefix: "cs_ste_vec_v1".to_string(),
        term_filters: Vec::new(),
        array_index_mode: cipherstash_config::column::ArrayIndexMode::ALL,
        mode: cipherstash_config::column::SteVecMode::Compat,
    });
    let column_config = create_column_config_with_indexes(vec![ste_vec_index.clone()]);

    let prepared = PreparedPlaintext::new(
        Cow::Owned(column_config.clone()),
        identifier.clone(),
        plaintext_value.clone(),
        EqlOperation::Store,
    );

    // Encrypt the JSON document
    let encrypt_opts = EqlEncryptOpts {
        keyset_id: Some(keyset.id),
        ..Default::default()
    };
    let encrypted = encrypt_eql(cipher.clone(), vec![prepared], &encrypt_opts)
        .await
        .expect("failed to encrypt JSON document");

    assert_eq!(encrypted.len(), 1);
    let stevec_payload = expect_ste_vec(expect_store(&encrypted[0]));
    assert_eq!(stevec_payload.identifier, identifier);
    assert!(
        !stevec_payload.ste_vec.is_empty(),
        "SteVec should carry per-selector entries"
    );

    // Decrypt and verify
    let decrypt_opts = EqlDecryptOpts::default();
    let stored: Vec<EqlCiphertext> = encrypted.into_iter().map(into_store).collect();
    let decrypted = decrypt_eql(cipher.clone(), stored, &decrypt_opts)
        .await
        .expect("failed to decrypt");

    assert_eq!(decrypted.len(), 1);
    let decrypted_json = decrypted[0]
        .clone_as_json()
        .expect("failed to convert to JSON");
    assert_eq!(decrypted_json, json_document);

    // Test Query 1: Match query - query with full JSON document (containment)
    let query_json = json!({
        "user": {
            "name": "Alice Johnson"
        }
    });
    let query_plaintext = cipherstash_client::encryption::Plaintext::Json(Some(query_json));

    let query_prepared = PreparedPlaintext::new(
        Cow::Borrowed(&column_config),
        identifier.clone(),
        query_plaintext,
        EqlOperation::Query(
            &ste_vec_index.index_type,
            cipherstash_client::encryption::QueryOp::Default,
        ),
    );

    let query_encrypted = encrypt_eql(cipher.clone(), vec![query_prepared], &encrypt_opts)
        .await
        .expect("failed to encrypt query");

    // A default SteVec match (JSON containment) query produces a Containment
    // term carrying the full STE query vector.
    let stevec_query = expect_query_ste_vec(expect_query(&query_encrypted[0]));
    assert!(
        matches!(stevec_query.term, SteVecQueryTerm::Containment { .. }),
        "SteVec containment query should produce a Containment term, got {:?}",
        stevec_query.term
    );

    // Test Query 2: Selector query - query by JSON path
    let selector_path = "$.user.email";
    let selector_plaintext =
        cipherstash_client::encryption::Plaintext::Text(Some(selector_path.to_string()));

    let selector_prepared = PreparedPlaintext::new(
        Cow::Borrowed(&column_config),
        identifier.clone(),
        selector_plaintext,
        EqlOperation::Query(
            &ste_vec_index.index_type,
            cipherstash_client::encryption::QueryOp::SteVecSelector,
        ),
    );

    let selector_encrypted = encrypt_eql(cipher.clone(), vec![selector_prepared], &encrypt_opts)
        .await
        .expect("failed to encrypt selector query");

    let selector_query = expect_query_ste_vec(expect_query(&selector_encrypted[0]));
    assert!(
        matches!(selector_query.term, SteVecQueryTerm::Selector { .. }),
        "Selector query should produce a Selector term, got {:?}",
        selector_query.term
    );

    // Test Query 3: Term query with string (for equality) - using OPE
    let term_string_value = "alice@example.com";
    let term_string_plaintext =
        cipherstash_client::encryption::Plaintext::from(term_string_value.to_string());

    let term_string_prepared = PreparedPlaintext::new(
        Cow::Borrowed(&column_config),
        identifier.clone(),
        term_string_plaintext,
        EqlOperation::Query(
            &ste_vec_index.index_type,
            cipherstash_client::encryption::QueryOp::SteVecTerm,
        ),
    );

    let term_string_encrypted =
        encrypt_eql(cipher.clone(), vec![term_string_prepared], &encrypt_opts)
            .await
            .expect("failed to encrypt string term query");

    let term_string_query = expect_query_ste_vec(expect_query(&term_string_encrypted[0]));
    assert!(
        matches!(term_string_query.term, SteVecQueryTerm::OreCllw { .. }),
        "String term query should produce an OreCllw term, got {:?}",
        term_string_query.term
    );

    // Test Query 4: Term query with number (for OPE range comparison)
    let term_number_value = 28_f64;
    let term_number_plaintext = cipherstash_client::encryption::Plaintext::from(term_number_value);

    let term_number_prepared = PreparedPlaintext::new(
        Cow::Borrowed(&column_config),
        identifier.clone(),
        term_number_plaintext,
        EqlOperation::Query(
            &ste_vec_index.index_type,
            cipherstash_client::encryption::QueryOp::SteVecTerm,
        ),
    );

    let term_number_encrypted =
        encrypt_eql(cipher.clone(), vec![term_number_prepared], &encrypt_opts)
            .await
            .expect("failed to encrypt number term query");

    let term_number_query = expect_query_ste_vec(expect_query(&term_number_encrypted[0]));
    assert!(
        matches!(term_number_query.term, SteVecQueryTerm::OreCllw { .. }),
        "Number term query should produce an OreCllw term, got {:?}",
        term_number_query.term
    );
}

fn expect_store(output: &EqlOutput) -> &EqlCiphertext {
    match output {
        EqlOutput::Store(c) => c,
        EqlOutput::Query(_) => panic!("expected EqlOutput::Store, got Query"),
    }
}

fn into_store(output: EqlOutput) -> EqlCiphertext {
    match output {
        EqlOutput::Store(c) => c,
        EqlOutput::Query(_) => panic!("expected EqlOutput::Store, got Query"),
    }
}

fn expect_query(output: &EqlOutput) -> &EqlQueryPayload {
    match output {
        EqlOutput::Query(q) => q,
        EqlOutput::Store(_) => panic!("expected EqlOutput::Query, got Store"),
    }
}

fn expect_encrypted(eql: &EqlCiphertext) -> &EncryptedPayload {
    match eql {
        EqlCiphertext::Encrypted(p) => p,
        EqlCiphertext::SteVec(_) => panic!("expected EqlCiphertext::Encrypted, got SteVec"),
    }
}

fn expect_ste_vec(eql: &EqlCiphertext) -> &SteVecPayload {
    match eql {
        EqlCiphertext::SteVec(p) => p,
        EqlCiphertext::Encrypted(_) => panic!("expected EqlCiphertext::SteVec, got Encrypted"),
    }
}

fn expect_query_encrypted(eql: &EqlQueryPayload) -> &EncryptedQueryPayload {
    match eql {
        EqlQueryPayload::Encrypted(p) => p,
        EqlQueryPayload::SteVec(_) => panic!("expected EqlQueryPayload::Encrypted, got SteVec"),
    }
}

fn expect_query_ste_vec(eql: &EqlQueryPayload) -> &SteVecQueryPayload {
    match eql {
        EqlQueryPayload::SteVec(p) => p,
        EqlQueryPayload::Encrypted(_) => panic!("expected EqlQueryPayload::SteVec, got Encrypted"),
    }
}