From a55221573b0c5f5625b4fa1a339c45f4b427db7c Mon Sep 17 00:00:00 2001 From: hdbg Date: Sat, 14 Feb 2026 23:44:37 +0100 Subject: [PATCH] feat(unseal): add unseal protocol support for user agents --- protobufs/arbiter.proto | 25 +--------- protobufs/client.proto | 17 +++++++ protobufs/unseal.proto | 19 +++++--- protobufs/user_agent.proto | 21 +++++++++ server/crates/arbiter-proto/src/lib.rs | 3 ++ server/crates/arbiter-server/src/actors.rs | 1 + .../src/{context => actors}/bootstrap.rs | 0 .../arbiter-server/src/actors/user_agent.rs | 47 +++++++++++++------ .../arbiter-server/src/context/lease.rs | 41 ---------------- 9 files changed, 90 insertions(+), 84 deletions(-) create mode 100644 protobufs/client.proto create mode 100644 protobufs/user_agent.proto rename server/crates/arbiter-server/src/{context => actors}/bootstrap.rs (100%) delete mode 100644 server/crates/arbiter-server/src/context/lease.rs diff --git a/protobufs/arbiter.proto b/protobufs/arbiter.proto index 67bf6d0..9d991c5 100644 --- a/protobufs/arbiter.proto +++ b/protobufs/arbiter.proto @@ -3,29 +3,8 @@ syntax = "proto3"; package arbiter; import "auth.proto"; - -message ClientRequest { - oneof payload { - arbiter.auth.ClientMessage auth_message = 1; - } -} - -message ClientResponse { - oneof payload { - arbiter.auth.ServerMessage auth_message = 1; - } -} - -message UserAgentRequest { - oneof payload { - arbiter.auth.ClientMessage auth_message = 1; - } -} -message UserAgentResponse { - oneof payload { - arbiter.auth.ServerMessage auth_message = 1; - } -} +import "client.proto"; +import "user_agent.proto"; message ServerInfo { string version = 1; diff --git a/protobufs/client.proto b/protobufs/client.proto new file mode 100644 index 0000000..ebe0422 --- /dev/null +++ b/protobufs/client.proto @@ -0,0 +1,17 @@ +syntax = "proto3"; + +package arbiter; + +import "auth.proto"; + +message ClientRequest { + oneof payload { + arbiter.auth.ClientMessage auth_message = 1; + } +} + +message ClientResponse { + oneof payload { + arbiter.auth.ServerMessage auth_message = 1; + } +} diff --git a/protobufs/unseal.proto b/protobufs/unseal.proto index 9ba0837..6ca7cca 100644 --- a/protobufs/unseal.proto +++ b/protobufs/unseal.proto @@ -2,13 +2,20 @@ syntax = "proto3"; package arbiter.unseal; -message UserAgentKeyRequest {} +import "google/protobuf/empty.proto"; -message ServerKeyResponse { +message UnsealStart {} + +message UnsealStartResponse { bytes pubkey = 1; } -message UserAgentSealedKey { - bytes sealed_key = 1; - bytes pubkey = 2; - bytes nonce = 3; +message UnsealEncryptedKey { + bytes key = 1; +} + +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 new file mode 100644 index 0000000..a3d7fd5 --- /dev/null +++ b/protobufs/user_agent.proto @@ -0,0 +1,21 @@ +syntax = "proto3"; + +package arbiter; + +import "auth.proto"; +import "unseal.proto"; + +message UserAgentRequest { + oneof payload { + arbiter.auth.ClientMessage auth_message = 1; + arbiter.unseal.UnsealStart unseal_start = 2; + arbiter.unseal.UnsealEncryptedKey unseal_encrypted_key = 3; + } +} +message UserAgentResponse { + oneof payload { + arbiter.auth.ServerMessage auth_message = 1; + arbiter.unseal.UnsealStartResponse unseal_start_response = 2; + arbiter.unseal.UnsealResult unseal_result = 3; + } +} diff --git a/server/crates/arbiter-proto/src/lib.rs b/server/crates/arbiter-proto/src/lib.rs index bce8e36..2a40151 100644 --- a/server/crates/arbiter-proto/src/lib.rs +++ b/server/crates/arbiter-proto/src/lib.rs @@ -6,6 +6,9 @@ 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.rs b/server/crates/arbiter-server/src/actors.rs index 691101e..2cb579e 100644 --- a/server/crates/arbiter-server/src/actors.rs +++ b/server/crates/arbiter-server/src/actors.rs @@ -1,2 +1,3 @@ pub mod user_agent; pub mod client; +pub(crate) mod bootstrap; \ No newline at end of file diff --git a/server/crates/arbiter-server/src/context/bootstrap.rs b/server/crates/arbiter-server/src/actors/bootstrap.rs similarity index 100% rename from server/crates/arbiter-server/src/context/bootstrap.rs rename to server/crates/arbiter-server/src/actors/bootstrap.rs diff --git a/server/crates/arbiter-server/src/actors/user_agent.rs b/server/crates/arbiter-server/src/actors/user_agent.rs index ecd8431..59f7616 100644 --- a/server/crates/arbiter-server/src/actors/user_agent.rs +++ b/server/crates/arbiter-server/src/actors/user_agent.rs @@ -1,5 +1,5 @@ use arbiter_proto::proto::{ - UserAgentRequest, UserAgentResponse, + UserAgentResponse, auth::{ self, AuthChallenge, AuthChallengeRequest, AuthOk, ClientMessage, ServerMessage as AuthServerMessage, client_message::Payload as ClientAuthPayload, @@ -11,22 +11,19 @@ use arbiter_proto::proto::{ use diesel::{ExpressionMethods as _, OptionalExtension as _, QueryDsl, dsl::update}; use diesel_async::{AsyncConnection, RunQueryDsl}; use ed25519_dalek::VerifyingKey; -use futures::StreamExt; use kameo::{ Actor, actor::{ActorRef, Spawn}, - error::SendError, messages, prelude::Context, }; -use tokio::sync::mpsc; use tokio::sync::mpsc::Sender; use tonic::Status; use tracing::{error, info}; use crate::{ ServerContext, - context::bootstrap::{BootstrapActor, ConsumeToken}, + actors::bootstrap::{BootstrapActor, ConsumeToken}, db::{self, schema}, errors::GrpcStatusExt, }; @@ -54,12 +51,15 @@ smlang::statemachine!( custom_error: false, transitions: { *Init + AuthRequest(AuthRequestContext) / auth_request_context = ReceivedAuthRequest(AuthRequestContext), - ReceivedAuthRequest(AuthRequestContext) + ReceivedBootstrapToken = Authenticated, + ReceivedAuthRequest(AuthRequestContext) + ReceivedBootstrapToken = Idle, ReceivedAuthRequest(AuthRequestContext) + SentChallenge(ChallengeContext) / move_challenge = WaitingForChallengeSolution(ChallengeContext), - WaitingForChallengeSolution(ChallengeContext) + ReceivedGoodSolution = Authenticated, + WaitingForChallengeSolution(ChallengeContext) + ReceivedGoodSolution = Idle, WaitingForChallengeSolution(ChallengeContext) + ReceivedBadSolution = AuthError, // block further transitions, but connection should close anyway + + Idle + UnsealRequest / generate_temp_keypair = UnsealStarted(ed25519_dalek::SigningKey), + UnsealStarted(ed25519_dalek::SigningKey) + SentTempKeypair / move_keypair = WaitingForUnsealKey(ed25519_dalek::SigningKey), } ); @@ -69,7 +69,7 @@ impl UserAgentStateMachineContext for DummyContext { #[allow(clippy::unused_unit)] fn move_challenge( &mut self, - state_data: &AuthRequestContext, + _state_data: &AuthRequestContext, event_data: ChallengeContext, ) -> Result { Ok(event_data) @@ -83,6 +83,21 @@ impl UserAgentStateMachineContext for DummyContext { ) -> Result { Ok(event_data) } + + #[allow(missing_docs)] + #[allow(clippy::unused_unit)] + fn move_keypair( + &mut self, + state_data: &ed25519_dalek::SigningKey, + ) -> Result { + Ok(state_data.clone()) + } + + #[allow(missing_docs)] + #[allow(clippy::unused_unit)] + fn generate_temp_keypair(&mut self) -> Result { + Ok(ed25519_dalek::SigningKey::generate(&mut rand::rng())) + } } #[derive(Actor)] @@ -90,7 +105,8 @@ pub struct UserAgentActor { db: db::DatabasePool, bootstapper: ActorRef, state: UserAgentStateMachine, - tx: Sender>, + // will be used in future + _tx: Sender>, } impl UserAgentActor { @@ -102,10 +118,11 @@ impl UserAgentActor { db: context.db.clone(), bootstapper: context.bootstrapper.clone(), state: UserAgentStateMachine::new(DummyContext), - tx, + _tx: tx, } } + #[cfg(test)] pub(crate) fn new_manual( db: db::DatabasePool, bootstapper: ActorRef, @@ -115,7 +132,7 @@ impl UserAgentActor { db, bootstapper, state: UserAgentStateMachine::new(DummyContext), - tx, + _tx: tx, } } @@ -319,8 +336,10 @@ mod tests { use kameo::actor::Spawn; use crate::{ - actors::user_agent::{HandleAuthChallengeRequest, HandleAuthChallengeSolution}, - context::bootstrap::BootstrapActor, + actors::{ + bootstrap::BootstrapActor, + user_agent::{HandleAuthChallengeRequest, HandleAuthChallengeSolution}, + }, db::{self, schema}, }; @@ -449,7 +468,7 @@ mod tests { { let mut conn = db.get().await.unwrap(); insert_into(schema::useragent_client::table) - .values((schema::useragent_client::public_key.eq(pubkey_bytes.clone()))) + .values(schema::useragent_client::public_key.eq(pubkey_bytes.clone())) .execute(&mut conn) .await .unwrap(); diff --git a/server/crates/arbiter-server/src/context/lease.rs b/server/crates/arbiter-server/src/context/lease.rs deleted file mode 100644 index 2b0f2bd..0000000 --- a/server/crates/arbiter-server/src/context/lease.rs +++ /dev/null @@ -1,41 +0,0 @@ -use std::sync::Arc; - -use dashmap::DashSet; - -#[derive(Clone, Default)] -struct LeaseStorage(Arc>); - -// A lease that automatically releases the item when dropped -pub struct Lease { - item: T, - storage: LeaseStorage, -} -impl Drop for Lease { - fn drop(&mut self) { - self.storage.0.remove(&self.item); - } -} - -#[derive(Clone, Default)] -pub struct LeaseHandler { - storage: LeaseStorage, -} - -impl LeaseHandler { - pub fn new() -> Self { - Self { - storage: LeaseStorage(Arc::new(DashSet::new())), - } - } - - pub fn acquire(&self, item: T) -> Result, ()> { - if self.storage.0.insert(item.clone()) { - Ok(Lease { - item, - storage: self.storage.clone(), - }) - } else { - Err(()) - } - } -}