feat(keyholder): add seal method and unseal integration tests

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(schema::root_key_history::data_encryption_nonce)
-                .select(RootKeyHistory::as_select() )
+                .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)]

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},
+        keyholder::KeyHolder,
-        user_agent::{HandleAuthChallengeRequest, HandleAuthChallengeSolution},
+        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<UserAgentActor>) {
+    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<UserAgentActor>,
+    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()),
+        );
+    }
+}