MarketTakers/arbiter
5
0
Fork 0
Code Issues 27 Pull Requests 4 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
hdbg
2026-02-16 21:30:23 +01:00
parent 8f5d4cc385
commit c6e13dc476
2 changed files with 277 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
server/crates/arbiter-server/src/actors/keyholder.rs
View File

@@ -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)]

265
server/crates/arbiter-server/src/actors/user_agent/tests.rs
View File

@@ -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<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()),
);
}
}