From c6e13dc4761cfc9101149a51f739da62bc32d01d Mon Sep 17 00:00:00 2001 From: hdbg Date: Mon, 16 Feb 2026 21:30:23 +0100 Subject: [PATCH] feat(keyholder): add seal method and unseal integration tests --- .../arbiter-server/src/actors/keyholder.rs | 19 +- .../src/actors/user_agent/tests.rs | 265 +++++++++++++++++- 2 files changed, 277 insertions(+), 7 deletions(-) diff --git a/server/crates/arbiter-server/src/actors/keyholder.rs b/server/crates/arbiter-server/src/actors/keyholder.rs index bb820ca..4a0bfbf 100644 --- a/server/crates/arbiter-server/src/actors/keyholder.rs +++ b/server/crates/arbiter-server/src/actors/keyholder.rs @@ -213,8 +213,8 @@ impl KeyHolder { let mut conn = self.db.get().await?; schema::root_key_history::table .filter(schema::root_key_history::id.eq(*root_key_history_id)) - .select(schema::root_key_history::data_encryption_nonce ) - .select(RootKeyHistory::as_select() ) + .select(schema::root_key_history::data_encryption_nonce) + .select(RootKeyHistory::as_select()) .first(&mut conn) .await? }; @@ -327,6 +327,21 @@ impl KeyHolder { pub fn get_state(&self) -> StateDiscriminants { self.state.discriminant() } + + #[message] + pub fn seal(&mut self) -> Result<(), Error> { + let State::Unsealed { + root_key_history_id, + .. + } = &self.state + else { + return Err(Error::NotBootstrapped); + }; + self.state = State::Sealed { + root_key_history_id: *root_key_history_id, + }; + Ok(()) + } } #[cfg(test)] diff --git a/server/crates/arbiter-server/src/actors/user_agent/tests.rs b/server/crates/arbiter-server/src/actors/user_agent/tests.rs index eec285e..2dbcdd3 100644 --- a/server/crates/arbiter-server/src/actors/user_agent/tests.rs +++ b/server/crates/arbiter-server/src/actors/user_agent/tests.rs @@ -1,28 +1,133 @@ use arbiter_proto::proto::{ - UserAgentResponse, + UnsealEncryptedKey, UnsealResult, UnsealStart, UserAgentResponse, auth::{self, AuthChallengeRequest, AuthOk}, user_agent_response::Payload as UserAgentResponsePayload, }; +use chacha20poly1305::{AeadInPlace, XChaCha20Poly1305, XNonce, aead::KeyInit}; use diesel::{ExpressionMethods as _, QueryDsl, insert_into}; use diesel_async::RunQueryDsl; use ed25519_dalek::Signer as _; -use kameo::actor::Spawn; +use kameo::actor::{ActorRef, Spawn}; +use memsafe::MemSafe; +use x25519_dalek::{EphemeralSecret, PublicKey}; use crate::{ actors::{ bootstrap::Bootstrapper, - keyholder::{self, KeyHolder}, - user_agent::{HandleAuthChallengeRequest, HandleAuthChallengeSolution}, + keyholder::KeyHolder, + user_agent::{ + HandleAuthChallengeRequest, HandleAuthChallengeSolution, HandleUnsealEncryptedKey, + HandleUnsealRequest, + }, }, - db::{self, schema}, + db::{self, models::ArbiterSetting, schema}, }; use super::UserAgentActor; +async fn seed_settings(db: &db::DatabasePool) { + let mut conn = db.get().await.unwrap(); + insert_into(schema::arbiter_settings::table) + .values(&ArbiterSetting { + id: 1, + root_key_id: None, + cert_key: vec![], + cert: vec![], + }) + .execute(&mut conn) + .await + .unwrap(); +} + +/// Bootstrap keyholder with `seal_key`, and Seal it +/// then create and authenticate a user agent (reaching Idle state). +async fn setup_authenticated_user_agent( + seal_key: &[u8], +) -> (db::DatabasePool, ActorRef) { + let db = db::create_test_pool().await; + seed_settings(&db).await; + + let mut keyholder = KeyHolder::new(db.clone()).await.unwrap(); + keyholder + .bootstrap(MemSafe::new(seal_key.to_vec()).unwrap()) + .await + .unwrap(); + keyholder.seal().unwrap(); + let keyholder_ref = KeyHolder::spawn(keyholder); + + let bootstrapper = Bootstrapper::new(&db).await.unwrap(); + let token = bootstrapper.get_token().unwrap(); + let bootstrapper_ref = Bootstrapper::spawn(bootstrapper); + + let user_agent = UserAgentActor::new_manual( + db.clone(), + bootstrapper_ref, + keyholder_ref, + tokio::sync::mpsc::channel(1).0, + ); + let user_agent_ref = UserAgentActor::spawn(user_agent); + + let auth_key = ed25519_dalek::SigningKey::generate(&mut rand::rng()); + user_agent_ref + .ask(HandleAuthChallengeRequest { + req: AuthChallengeRequest { + pubkey: auth_key.verifying_key().to_bytes().to_vec(), + bootstrap_token: Some(token), + }, + }) + .await + .unwrap(); + + (db, user_agent_ref) +} + +/// Client side of the DH unseal exchange: +/// sends UnsealStart, derives shared secret, encrypts `key_to_send`. +async fn client_dh_encrypt( + user_agent_ref: &ActorRef, + key_to_send: &[u8], +) -> UnsealEncryptedKey { + let client_secret = EphemeralSecret::random(); + let client_public = PublicKey::from(&client_secret); + + let response = user_agent_ref + .ask(HandleUnsealRequest { + req: UnsealStart { + client_pubkey: client_public.as_bytes().to_vec(), + }, + }) + .await + .unwrap(); + + let server_pubkey = match response.payload.unwrap() { + UserAgentResponsePayload::UnsealStartResponse(resp) => resp.server_pubkey, + other => panic!("Expected UnsealStartResponse, got {other:?}"), + }; + let server_public = PublicKey::from( + <[u8; 32]>::try_from(server_pubkey.as_slice()).unwrap(), + ); + + let shared_secret = client_secret.diffie_hellman(&server_public); + let cipher = XChaCha20Poly1305::new(shared_secret.as_bytes().into()); + let nonce = XNonce::from([0u8; 24]); + let associated_data = b"unseal"; + let mut ciphertext = key_to_send.to_vec(); + cipher + .encrypt_in_place(&nonce, associated_data, &mut ciphertext) + .unwrap(); + + UnsealEncryptedKey { + nonce: nonce.to_vec(), + ciphertext, + associated_data: associated_data.to_vec(), + } +} + #[tokio::test] #[test_log::test] pub async fn test_bootstrap_token_auth() { let db = db::create_test_pool().await; + seed_settings(&db).await; // explicitly not installing any user_agent pubkeys let bootstrapper = Bootstrapper::new(&db).await.unwrap(); // this will create bootstrap token let keyholder = KeyHolder::new(db.clone()).await.unwrap(); @@ -80,6 +185,7 @@ pub async fn test_bootstrap_token_auth() { #[test_log::test] pub async fn test_bootstrap_invalid_token_auth() { let db = db::create_test_pool().await; + seed_settings(&db).await; // explicitly not installing any user_agent pubkeys let bootstrapper = Bootstrapper::new(&db).await.unwrap(); // this will create bootstrap token let keyholder = KeyHolder::new(db.clone()).await.unwrap(); @@ -132,6 +238,7 @@ pub async fn test_bootstrap_invalid_token_auth() { #[test_log::test] pub async fn test_challenge_auth() { let db = db::create_test_pool().await; + seed_settings(&db).await; let bootstrapper_ref = Bootstrapper::spawn(Bootstrapper::new(&db).await.unwrap()); let keyholder_ref = KeyHolder::spawn(KeyHolder::new(db.clone()).await.unwrap()); @@ -206,3 +313,151 @@ pub async fn test_challenge_auth() { } ); } + +#[tokio::test] +#[test_log::test] +pub async fn test_unseal_success() { + let seal_key = b"test-seal-key"; + let (_db, user_agent_ref) = setup_authenticated_user_agent(seal_key).await; + + let encrypted_key = client_dh_encrypt(&user_agent_ref, seal_key).await; + + let response = user_agent_ref + .ask(HandleUnsealEncryptedKey { req: encrypted_key }) + .await + .unwrap(); + + assert_eq!( + response.payload.unwrap(), + UserAgentResponsePayload::UnsealResult(UnsealResult::Success.into()), + ); +} + +#[tokio::test] +#[test_log::test] +pub async fn test_unseal_wrong_seal_key() { + let (_db, user_agent_ref) = setup_authenticated_user_agent(b"correct-key").await; + + // Encrypt a different key through the DH channel + let encrypted_key = client_dh_encrypt(&user_agent_ref, b"wrong-key").await; + + let response = user_agent_ref + .ask(HandleUnsealEncryptedKey { req: encrypted_key }) + .await + .unwrap(); + + assert_eq!( + response.payload.unwrap(), + UserAgentResponsePayload::UnsealResult(UnsealResult::InvalidKey.into()), + ); +} + +#[tokio::test] +#[test_log::test] +pub async fn test_unseal_corrupted_ciphertext() { + let (_db, user_agent_ref) = setup_authenticated_user_agent(b"test-key").await; + + // Do UnsealStart to reach WaitingForUnsealKey state + let client_secret = EphemeralSecret::random(); + let client_public = PublicKey::from(&client_secret); + + user_agent_ref + .ask(HandleUnsealRequest { + req: UnsealStart { + client_pubkey: client_public.as_bytes().to_vec(), + }, + }) + .await + .unwrap(); + + // Send garbage that wasn't encrypted with the DH shared secret + let response = user_agent_ref + .ask(HandleUnsealEncryptedKey { + req: UnsealEncryptedKey { + nonce: vec![0u8; 24], + ciphertext: vec![0u8; 32], + associated_data: vec![], + }, + }) + .await + .unwrap(); + + assert_eq!( + response.payload.unwrap(), + UserAgentResponsePayload::UnsealResult(UnsealResult::InvalidKey.into()), + ); +} + +#[tokio::test] +#[test_log::test] +pub async fn test_unseal_start_without_auth_fails() { + let db = db::create_test_pool().await; + seed_settings(&db).await; + + let keyholder_ref = KeyHolder::spawn( KeyHolder::new(db.clone()).await.unwrap()); + let bootstrapper_ref = Bootstrapper::spawn(Bootstrapper::new(&db).await.unwrap()); + + let user_agent = UserAgentActor::new_manual( + db.clone(), + bootstrapper_ref, + keyholder_ref, + tokio::sync::mpsc::channel(1).0, + ); + let user_agent_ref = UserAgentActor::spawn(user_agent); + + // Try unseal from Init state (not authenticated) + let client_secret = EphemeralSecret::random(); + let client_public = PublicKey::from(&client_secret); + + let result = user_agent_ref + .ask(HandleUnsealRequest { + req: UnsealStart { + client_pubkey: client_public.as_bytes().to_vec(), + }, + }) + .await; + + match result { + Err(kameo::error::SendError::HandlerError(status)) => { + assert_eq!(status.code(), tonic::Code::Internal); + } + other => panic!("Expected state machine error, got {other:?}"), + } +} + +#[tokio::test] +#[test_log::test] +pub async fn test_unseal_retry_after_invalid_key() { + let seal_key = b"real-seal-key"; + let (_db, user_agent_ref) = setup_authenticated_user_agent(seal_key).await; + + // First attempt: wrong key -> InvalidKey, state goes back to Idle + { + let encrypted_key = client_dh_encrypt(&user_agent_ref, b"wrong-key").await; + + let response = user_agent_ref + .ask(HandleUnsealEncryptedKey { req: encrypted_key }) + .await + .unwrap(); + + assert_eq!( + response.payload.unwrap(), + UserAgentResponsePayload::UnsealResult(UnsealResult::InvalidKey.into()), + ); + } + + // Second attempt: correct key -> Success + { + let encrypted_key = client_dh_encrypt(&user_agent_ref, seal_key).await; + + let response = user_agent_ref + .ask(HandleUnsealEncryptedKey { req: encrypted_key }) + .await + .unwrap(); + + assert_eq!( + response.payload.unwrap(), + UserAgentResponsePayload::UnsealResult(UnsealResult::Success.into()), + ); + } +}