arbiter/server/crates/arbiter-server/tests/vault/lifecycle.rs

use crate::common;
use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
use arbiter_server::{
    actors::{
        GlobalActors,
        vault::{Error, Vault},
        vault_coordinator::{Error as CoordinatorError, StartBootstrap, VaultCoordinator},
    },
    crypto::{KeyCell, encryption::v1::{Nonce, ROOT_KEY_TAG}},
    db::{self, models, schema},
};

use diesel::{QueryDsl, SelectableHelper};
use diesel_async::RunQueryDsl;
use kameo::actor::Spawn as _;

#[tokio::test]
#[test_log::test]
async fn test_bootstrap() {
    let db = db::create_test_pool().await;
    let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus())
        .await
        .unwrap();

    let seal_key = KeyCell::from([0u8; 32]);
    actor.bootstrap(seal_key).await.unwrap();

    let mut conn = db.get().await.unwrap();
    let row: models::RootKeyHistory = schema::root_key_history::table
        .select(models::RootKeyHistory::as_select())
        .first(&mut conn)
        .await
        .unwrap();

    assert_eq!(row.schema_version, 1);
    assert_eq!(row.tag, ROOT_KEY_TAG);
    assert!(!row.ciphertext.is_empty());
    assert!(!row.salt.is_empty());
    assert_eq!(row.data_encryption_nonce, Nonce::default().to_vec());
}

#[tokio::test]
#[test_log::test]
async fn test_bootstrap_rejects_double() {
    let db = db::create_test_pool().await;
    let mut actor = common::bootstrapped_vault(&db).await;

    let seal_key2 = KeyCell::from([0u8; 32]);
    let err = actor.bootstrap(seal_key2).await.unwrap_err();
    assert!(matches!(err, Error::AlreadyBootstrapped));
}

#[tokio::test]
#[test_log::test]
async fn test_create_new_before_bootstrap_fails() {
    let db = db::create_test_pool().await;
    let mut actor = Vault::new(db, GlobalActors::spawn_message_bus())
        .await
        .unwrap();

    let err = actor
        .create_new(SafeCell::new(b"data".to_vec()))
        .await
        .unwrap_err();
    assert!(matches!(err, Error::NotBootstrapped));
}

#[tokio::test]
#[test_log::test]
async fn test_decrypt_before_bootstrap_fails() {
    let db = db::create_test_pool().await;
    let mut actor = Vault::new(db, GlobalActors::spawn_message_bus())
        .await
        .unwrap();

    let err = actor.decrypt(1).await.unwrap_err();
    assert!(matches!(err, Error::NotBootstrapped));
}

#[tokio::test]
#[test_log::test]
async fn test_new_restores_sealed_state() {
    let db = db::create_test_pool().await;
    let actor = common::bootstrapped_vault(&db).await;
    drop(actor);

    let mut actor2 = Vault::new(db, GlobalActors::spawn_message_bus())
        .await
        .unwrap();
    let err = actor2.decrypt(1).await.unwrap_err();
    assert!(matches!(err, Error::Sealed));
}

#[tokio::test]
#[test_log::test]
async fn test_unseal_correct_password() {
    let db = db::create_test_pool().await;
    let mut actor = common::bootstrapped_vault(&db).await;

    let plaintext = b"survive a restart";
    let aead_id = actor
        .create_new(SafeCell::new(plaintext.to_vec()))
        .await
        .unwrap();
    drop(actor);

    let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus())
        .await
        .unwrap();
    let seal_key = KeyCell::from([0u8; 32]);
    actor.try_unseal(seal_key).await.unwrap();

    let mut decrypted = actor.decrypt(aead_id).await.unwrap();
    assert_eq!(*decrypted.read(), plaintext);
}

#[tokio::test]
#[test_log::test]
async fn test_unseal_wrong_then_correct_password() {
    let db = db::create_test_pool().await;
    let mut actor = common::bootstrapped_vault(&db).await;

    let plaintext = b"important data";
    let aead_id = actor
        .create_new(SafeCell::new(plaintext.to_vec()))
        .await
        .unwrap();
    drop(actor);

    let mut actor = Vault::new(db.clone(), GlobalActors::spawn_message_bus())
        .await
        .unwrap();

    let bad_key = KeyCell::from([1u8; 32]);
    let err = actor.try_unseal(bad_key).await.unwrap_err();
    assert!(matches!(err, Error::InvalidKey));

    let good_key = KeyCell::from([0u8; 32]);
    actor.try_unseal(good_key).await.unwrap();

    let mut decrypted = actor.decrypt(aead_id).await.unwrap();
    assert_eq!(*decrypted.read(), plaintext);
}

#[tokio::test]
#[test_log::test]
async fn two_operator_vault_requires_recovery_share() {
    let db = db::create_test_pool().await;
    let bus = GlobalActors::spawn_message_bus();
    let vault_ref = Vault::spawn(Vault::new(db.clone(), bus).await.unwrap());
    let coordinator = VaultCoordinator::spawn(VaultCoordinator::new(db, vault_ref));

    let err = coordinator
        .ask(StartBootstrap {
            operator_id: 1,
            declared_count: 2,
            recovery_count: 0,
        })
        .await
        .unwrap_err();

    assert!(
        matches!(
            err,
            kameo::error::SendError::HandlerError(CoordinatorError::TwoOperatorsRequireRecovery)
        ),
        "expected TwoOperatorsRequireRecovery, got {err:?}"
    );
}