From eb9cbc88e9e830694af274938e980f068a294217 Mon Sep 17 00:00:00 2001 From: hdbg Date: Mon, 16 Feb 2026 21:07:09 +0100 Subject: [PATCH] feat(server::user-agent): Unseal implemented --- protobufs/unseal.proto | 25 ------- protobufs/user_agent.proto | 40 +++++++++-- server/crates/arbiter-proto/src/lib.rs | 3 - .../arbiter-server/src/actors/user_agent.rs | 69 ++++++++++++++----- .../src/actors/user_agent/tests.rs | 10 +++ 5 files changed, 96 insertions(+), 51 deletions(-) delete mode 100644 protobufs/unseal.proto diff --git a/protobufs/unseal.proto b/protobufs/unseal.proto deleted file mode 100644 index 8c006a0..0000000 --- a/protobufs/unseal.proto +++ /dev/null @@ -1,25 +0,0 @@ -syntax = "proto3"; - -package arbiter.unseal; - -import "google/protobuf/empty.proto"; - -message UnsealStart { - bytes client_pubkey = 1; -} - -message UnsealStartResponse { - bytes server_pubkey = 1; -} -message UnsealEncryptedKey { - bytes nonce = 1; - bytes ciphertext = 2; - bytes associated_data = 3; -} - -enum UnsealResult { - UNSEAL_RESULT_UNSPECIFIED = 0; - UNSEAL_RESULT_SUCCESS = 1; - UNSEAL_RESULT_INVALID_KEY = 2; - UNSEAL_RESULT_UNBOOTSTRAPPED = 3; -} diff --git a/protobufs/user_agent.proto b/protobufs/user_agent.proto index a3d7fd5..b9958ee 100644 --- a/protobufs/user_agent.proto +++ b/protobufs/user_agent.proto @@ -3,19 +3,49 @@ syntax = "proto3"; package arbiter; import "auth.proto"; -import "unseal.proto"; +import "google/protobuf/empty.proto"; + +message UnsealStart { + bytes client_pubkey = 1; +} + +message UnsealStartResponse { + bytes server_pubkey = 1; +} +message UnsealEncryptedKey { + bytes nonce = 1; + bytes ciphertext = 2; + bytes associated_data = 3; +} + +enum UnsealResult { + UNSEAL_RESULT_UNSPECIFIED = 0; + UNSEAL_RESULT_SUCCESS = 1; + UNSEAL_RESULT_INVALID_KEY = 2; + UNSEAL_RESULT_UNBOOTSTRAPPED = 3; +} + +enum VaultState { + VAULT_STATE_UNSPECIFIED = 0; + VAULT_STATE_UNBOOTSTRAPPED = 1; + VAULT_STATE_SEALED = 2; + VAULT_STATE_UNSEALED = 3; + VAULT_STATE_ERROR = 4; +} message UserAgentRequest { oneof payload { arbiter.auth.ClientMessage auth_message = 1; - arbiter.unseal.UnsealStart unseal_start = 2; - arbiter.unseal.UnsealEncryptedKey unseal_encrypted_key = 3; + UnsealStart unseal_start = 2; + UnsealEncryptedKey unseal_encrypted_key = 3; + google.protobuf.Empty query_vault_state = 4; } } message UserAgentResponse { oneof payload { arbiter.auth.ServerMessage auth_message = 1; - arbiter.unseal.UnsealStartResponse unseal_start_response = 2; - arbiter.unseal.UnsealResult unseal_result = 3; + UnsealStartResponse unseal_start_response = 2; + UnsealResult unseal_result = 3; + VaultState vault_state = 4; } } diff --git a/server/crates/arbiter-proto/src/lib.rs b/server/crates/arbiter-proto/src/lib.rs index 2a40151..bce8e36 100644 --- a/server/crates/arbiter-proto/src/lib.rs +++ b/server/crates/arbiter-proto/src/lib.rs @@ -6,9 +6,6 @@ pub mod proto { pub mod auth { tonic::include_proto!("arbiter.auth"); } - pub mod unseal { - tonic::include_proto!("arbiter.unseal"); - } } pub mod transport; diff --git a/server/crates/arbiter-server/src/actors/user_agent.rs b/server/crates/arbiter-server/src/actors/user_agent.rs index ec788af..c180b28 100644 --- a/server/crates/arbiter-server/src/actors/user_agent.rs +++ b/server/crates/arbiter-server/src/actors/user_agent.rs @@ -1,25 +1,18 @@ -use std::{ - ops::DerefMut, - sync::Mutex, -}; +use std::{ops::DerefMut, sync::Mutex}; use arbiter_proto::proto::{ - UserAgentResponse, + UnsealEncryptedKey, UnsealResult, UnsealStart, UnsealStartResponse, UserAgentResponse, auth::{ self, AuthChallengeRequest, AuthOk, ServerMessage as AuthServerMessage, server_message::Payload as ServerAuthPayload, }, - unseal::{UnsealEncryptedKey, UnsealResult, UnsealStart, UnsealStartResponse}, user_agent_response::Payload as UserAgentResponsePayload, }; -use chacha20poly1305::{ - AeadInPlace, XChaCha20Poly1305, XNonce, - aead::KeyInit, -}; +use chacha20poly1305::{AeadInPlace, XChaCha20Poly1305, XNonce, aead::KeyInit}; use diesel::{ExpressionMethods as _, OptionalExtension as _, QueryDsl, dsl::update}; use diesel_async::{AsyncConnection, RunQueryDsl}; use ed25519_dalek::VerifyingKey; -use kameo::{Actor, actor::ActorRef, messages}; +use kameo::{Actor, actor::ActorRef, error::SendError, messages}; use memsafe::MemSafe; use tokio::sync::mpsc::Sender; use tonic::Status; @@ -30,6 +23,7 @@ use crate::{ ServerContext, actors::{ bootstrap::{Bootstrapper, ConsumeToken}, + keyholder::{self, KeyHolder, TryUnseal}, user_agent::state::{ AuthRequestContext, ChallengeContext, DummyContext, UnsealContext, UserAgentEvents, UserAgentStateMachine, UserAgentStates, @@ -50,6 +44,7 @@ pub(crate) use transport::handle_user_agent; pub struct UserAgentActor { db: db::DatabasePool, bootstapper: ActorRef, + keyholder: ActorRef, state: UserAgentStateMachine, // will be used in future _tx: Sender>, @@ -63,6 +58,7 @@ impl UserAgentActor { Self { db: context.db.clone(), bootstapper: context.bootstrapper.clone(), + keyholder: context.keyholder.clone(), state: UserAgentStateMachine::new(DummyContext), _tx: tx, } @@ -72,11 +68,13 @@ impl UserAgentActor { pub(crate) fn new_manual( db: db::DatabasePool, bootstapper: ActorRef, + keyholder: ActorRef, tx: Sender>, ) -> Self { Self { db, bootstapper, + keyholder, state: UserAgentStateMachine::new(DummyContext), _tx: tx, } @@ -280,22 +278,57 @@ impl UserAgentActor { let shared_secret = ephemeral_secret.diffie_hellman(&unseal_context.client_public_key); let cipher = XChaCha20Poly1305::new(shared_secret.as_bytes().into()); - let mut root_key_buffer = MemSafe::new(req.ciphertext.clone()).unwrap(); - let mut write_handle = root_key_buffer.write().unwrap(); - let write_handle = write_handle.deref_mut(); + let mut seal_key_buffer = MemSafe::new(req.ciphertext.clone()).unwrap(); - let decryption_result = cipher - .decrypt_in_place(nonce, &req.associated_data, write_handle); + let decryption_result = { + let mut write_handle = seal_key_buffer.write().unwrap(); + let write_handle = write_handle.deref_mut(); + cipher.decrypt_in_place(nonce, &req.associated_data, write_handle) + }; match decryption_result { - Ok(_) => todo!("Send key to the keyguarding"), + Ok(_) => { + match self + .keyholder + .ask(TryUnseal { + seal_key_raw: seal_key_buffer, + }) + .await + { + Ok(_) => { + info!("Successfully unsealed key with client-provided key"); + self.transition(UserAgentEvents::ReceivedValidKey)?; + Ok(unseal_response(UserAgentResponsePayload::UnsealResult( + UnsealResult::Success.into(), + ))) + } + Err(SendError::HandlerError(keyholder::Error::InvalidKey)) => { + self.transition(UserAgentEvents::ReceivedInvalidKey)?; + Ok(unseal_response(UserAgentResponsePayload::UnsealResult( + UnsealResult::InvalidKey.into(), + ))) + } + Err(SendError::HandlerError(err)) => { + error!(?err, "Keyholder failed to unseal key"); + self.transition(UserAgentEvents::ReceivedInvalidKey)?; + Ok(unseal_response(UserAgentResponsePayload::UnsealResult( + UnsealResult::InvalidKey.into(), + ))) + } + Err(err) => { + error!(?err, "Failed to send unseal request to keyholder"); + self.transition(UserAgentEvents::ReceivedInvalidKey)?; + Err(Status::internal("Vault is not available")) + } + } + } Err(err) => { error!(?err, "Failed to decrypt unseal key"); self.transition(UserAgentEvents::ReceivedInvalidKey)?; return Ok(unseal_response(UserAgentResponsePayload::UnsealResult( UnsealResult::InvalidKey.into(), ))); - }, + } } } 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 2e63eeb..eec285e 100644 --- a/server/crates/arbiter-server/src/actors/user_agent/tests.rs +++ b/server/crates/arbiter-server/src/actors/user_agent/tests.rs @@ -11,6 +11,7 @@ use kameo::actor::Spawn; use crate::{ actors::{ bootstrap::Bootstrapper, + keyholder::{self, KeyHolder}, user_agent::{HandleAuthChallengeRequest, HandleAuthChallengeSolution}, }, db::{self, schema}, @@ -24,12 +25,15 @@ pub async fn test_bootstrap_token_auth() { let db = db::create_test_pool().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(); let token = bootstrapper.get_token().unwrap(); let bootstrapper_ref = Bootstrapper::spawn(bootstrapper); + let keyholder_ref = KeyHolder::spawn(keyholder); let user_agent = UserAgentActor::new_manual( db.clone(), bootstrapper_ref, + keyholder_ref, tokio::sync::mpsc::channel(1).0, // dummy channel, we won't actually send responses in this test ); let user_agent_ref = UserAgentActor::spawn(user_agent); @@ -78,11 +82,15 @@ pub async fn test_bootstrap_invalid_token_auth() { let db = db::create_test_pool().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(); let bootstrapper_ref = Bootstrapper::spawn(bootstrapper); + let keyholder_ref = KeyHolder::spawn(keyholder); + let user_agent = UserAgentActor::new_manual( db.clone(), bootstrapper_ref, + keyholder_ref, tokio::sync::mpsc::channel(1).0, // dummy channel, we won't actually send responses in this test ); let user_agent_ref = UserAgentActor::spawn(user_agent); @@ -126,9 +134,11 @@ pub async fn test_challenge_auth() { let db = db::create_test_pool().await; let bootstrapper_ref = Bootstrapper::spawn(Bootstrapper::new(&db).await.unwrap()); + let keyholder_ref = KeyHolder::spawn(KeyHolder::new(db.clone()).await.unwrap()); let user_agent = UserAgentActor::new_manual( db.clone(), bootstrapper_ref, + keyholder_ref, tokio::sync::mpsc::channel(1).0, // dummy channel, we won't actually send responses in this test ); let user_agent_ref = UserAgentActor::spawn(user_agent);