feat(user-agent): add VaultGate for sealed vault authentication

This commit is contained in:
hdbg
2026-04-08 18:29:52 +02:00
parent 205227a3df
commit 87ee0fe87b
24 changed files with 900 additions and 625 deletions

View File

@@ -10,7 +10,6 @@ use tonic::{Request, Response, Status, async_trait};
use tracing::info;
use crate::{
grpc::user_agent::start,
peers::{client::ClientConnection, user_agent::UserAgentConnection},
};
@@ -63,7 +62,7 @@ impl arbiter_proto::proto::arbiter_service_server::ArbiterService for super::Ser
let (bi, rx) = GrpcBi::from_bi_stream(req_stream);
tokio::spawn(start(
tokio::spawn(user_agent::start(
UserAgentConnection {
db: self.context.db.clone(),
actors: self.context.actors.clone(),

View File

@@ -1,4 +1,4 @@
use tokio::sync::mpsc;
use tokio::sync::{mpsc, oneshot};
use arbiter_proto::{
proto::user_agent::{
@@ -14,8 +14,12 @@ use tonic::Status;
use tracing::{error, info, warn};
use crate::{
crypto::integrity,
grpc::request_tracker::RequestTracker,
peers::user_agent::{OutOfBand, UserAgentConnection, UserAgentSession},
peers::user_agent::{
Credentials, OutOfBand, UserAgentConnection, UserAgentSession,
vault_gate::VaultGate,
},
};
mod auth;
@@ -24,6 +28,7 @@ mod inbound;
mod outbound;
mod sdk_client;
mod vault;
mod vault_gate;
pub struct OutOfBandAdapter(mpsc::Sender<OutOfBand>);
@@ -124,27 +129,115 @@ pub async fn start(
) {
let mut request_tracker = RequestTracker::default();
let (id, pubkey) = match auth::start(&mut conn, &mut bi, &mut request_tracker).await {
Ok(pubkey) => pubkey,
let auth_creds = match auth::start(&mut conn, &mut bi, &mut request_tracker).await {
Ok(creds) => creds,
Err(e) => {
warn!(error = ?e, "Authentication failed");
return;
}
};
info!(?pubkey, "User authenticated successfully");
info!(pubkey = ?auth_creds.creds.pubkey, "User authenticated successfully");
let creds = if integrity::is_signing_available(&conn.actors.vault)
.await
.unwrap_or(false)
{
// Vault is unsealed; integrity was verified during auth — promote directly.
auth_creds.creds
} else {
// Vault is sealed/unbootstrapped; run the VaultGate phase.
let (promotion_tx, promotion_rx) = oneshot::channel();
let gate = VaultGate::spawn(VaultGate::new(
auth_creds,
conn.actors.clone(),
conn.db.clone(),
promotion_tx,
));
let result = vault_gate_loop(&mut bi, &gate, &mut request_tracker, promotion_rx).await;
gate.kill();
match result {
Some(creds) => creds,
None => return,
}
};
let (oob_sender, oob_receiver) = mpsc::channel(16);
let oob_adapter = OutOfBandAdapter(oob_sender);
let actor = UserAgentSession::spawn(UserAgentSession::new(
conn,
id,
pubkey,
Box::new(oob_adapter),
));
let actor = UserAgentSession::spawn(UserAgentSession::new(conn, creds, Box::new(oob_adapter)));
let actor_for_cleanup = actor.clone();
dispatch_loop(bi, actor, oob_receiver, request_tracker).await;
actor_for_cleanup.kill();
}
async fn vault_gate_loop(
bi: &mut GrpcBi<UserAgentRequest, UserAgentResponse>,
gate: &ActorRef<VaultGate>,
request_tracker: &mut RequestTracker,
mut promotion_rx: oneshot::Receiver<Result<Credentials, crate::peers::user_agent::vault_gate::Error>>,
) -> Option<Credentials> {
loop {
tokio::select! {
result = &mut promotion_rx => {
return match result {
Ok(Ok(creds)) => Some(creds),
Ok(Err(e)) => {
warn!(error = ?e, "VaultGate promotion failed");
None
}
Err(_) => {
warn!("VaultGate promotion channel closed unexpectedly");
None
}
};
}
message = bi.recv() => {
let Some(message) = message else { return None; };
let conn = match message {
Ok(conn) => conn,
Err(err) => {
warn!(error = ?err, "Failed to receive request during vault gate phase");
return None;
}
};
let request_id = match request_tracker.request(conn.id) {
Ok(id) => id,
Err(err) => {
let _ = bi.send(Err(err)).await;
return None;
}
};
let Some(payload) = conn.payload else {
let _ = bi.send(Err(Status::invalid_argument("Missing request payload"))).await;
return None;
};
let response = match payload {
UserAgentRequestPayload::Vault(req) => vault_gate::dispatch(gate, req).await,
_ => Err(Status::permission_denied("Only vault operations are permitted before unsealing")),
};
match response {
Ok(Some(payload)) => {
if bi.send(Ok(UserAgentResponse { id: Some(request_id), payload: Some(payload) })).await.is_err() {
return None;
}
}
Ok(None) => {}
Err(status) => {
let _ = bi.send(Err(status)).await;
return None;
}
}
}
}
}
}

View File

@@ -19,7 +19,7 @@ use tracing::warn;
use crate::{
grpc::request_tracker::RequestTracker,
peers::user_agent::{UserAgentConnection, auth},
peers::user_agent::{AuthCredentials, UserAgentConnection, auth},
};
pub struct AuthTransportAdapter<'a> {
@@ -167,7 +167,7 @@ pub async fn start(
conn: &mut UserAgentConnection,
bi: &mut GrpcBi<UserAgentRequest, UserAgentResponse>,
request_tracker: &mut RequestTracker,
) -> Result<(i32, authn::PublicKey), auth::Error> {
) -> Result<AuthCredentials, auth::Error> {
let transport = AuthTransportAdapter::new(bi, request_tracker);
auth::authenticate(conn, transport).await
}

View File

@@ -29,7 +29,7 @@ use crate::{
},
peers::user_agent::{
UserAgentSession,
session::connection::{
session::handlers::{
GrantMutationError, HandleEvmWalletCreate, HandleEvmWalletList, HandleGrantCreate,
HandleGrantDelete, HandleGrantList, HandleSignTransaction,
SignTransactionError as SessionSignTransactionError,

View File

@@ -25,7 +25,7 @@ use crate::{
grpc::Convert,
peers::user_agent::{
OutOfBand, UserAgentSession,
session::connection::{
session::handlers::{
HandleGrantEvmWalletAccess, HandleListWalletAccess, HandleNewClientApprove,
HandleRevokeEvmWalletAccess, HandleSdkClientList,
},

View File

@@ -1,34 +1,15 @@
use arbiter_proto::proto::shared::VaultState as ProtoVaultState;
use arbiter_proto::proto::user_agent::{
user_agent_response::Payload as UserAgentResponsePayload,
vault::{
self as proto_vault,
bootstrap::{
self as proto_bootstrap, BootstrapEncryptedKey as ProtoBootstrapEncryptedKey,
BootstrapResult as ProtoBootstrapResult,
},
request::Payload as VaultRequestPayload,
response::Payload as VaultResponsePayload,
unseal::{
self as proto_unseal, UnsealEncryptedKey as ProtoUnsealEncryptedKey,
UnsealResult as ProtoUnsealResult, UnsealStart,
request::Payload as UnsealRequestPayload, response::Payload as UnsealResponsePayload,
},
},
vault::{self as proto_vault, request::Payload as VaultRequestPayload, response::Payload as VaultResponsePayload},
};
use kameo::{actor::ActorRef, error::SendError};
use kameo::actor::ActorRef;
use tonic::Status;
use tracing::warn;
use crate::{
actors::vault::VaultState,
peers::user_agent::{
UserAgentSession,
session::connection::{
BootstrapError, HandleBootstrapEncryptedKey, HandleQueryVaultState,
HandleUnsealEncryptedKey, HandleUnsealRequest, UnsealError,
},
},
peers::user_agent::{UserAgentSession, session::handlers::HandleQueryVaultState},
};
fn wrap_vault_response(payload: VaultResponsePayload) -> UserAgentResponsePayload {
@@ -37,18 +18,6 @@ fn wrap_vault_response(payload: VaultResponsePayload) -> UserAgentResponsePayloa
})
}
fn wrap_unseal_response(payload: UnsealResponsePayload) -> UserAgentResponsePayload {
wrap_vault_response(VaultResponsePayload::Unseal(proto_unseal::Response {
payload: Some(payload),
}))
}
fn wrap_bootstrap_response(result: ProtoBootstrapResult) -> UserAgentResponsePayload {
wrap_vault_response(VaultResponsePayload::Bootstrap(proto_bootstrap::Response {
result: result.into(),
}))
}
pub(super) async fn dispatch(
actor: &ActorRef<UserAgentSession>,
req: proto_vault::Request,
@@ -59,109 +28,14 @@ pub(super) async fn dispatch(
match payload {
VaultRequestPayload::QueryState(_) => handle_query_vault_state(actor).await,
VaultRequestPayload::Unseal(req) => dispatch_unseal_request(actor, req).await,
VaultRequestPayload::Bootstrap(req) => handle_bootstrap_request(actor, req).await,
VaultRequestPayload::Unseal(_) | VaultRequestPayload::Bootstrap(_) => {
Err(Status::permission_denied(
"Vault is already unsealed; unseal/bootstrap not permitted in session",
))
}
}
}
async fn dispatch_unseal_request(
actor: &ActorRef<UserAgentSession>,
req: proto_unseal::Request,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let Some(payload) = req.payload else {
return Err(Status::invalid_argument("Missing unseal request payload"));
};
match payload {
UnsealRequestPayload::Start(req) => handle_unseal_start(actor, req).await,
UnsealRequestPayload::EncryptedKey(req) => handle_unseal_encrypted_key(actor, req).await,
}
}
async fn handle_unseal_start(
actor: &ActorRef<UserAgentSession>,
req: UnsealStart,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let client_pubkey = <[u8; 32]>::try_from(req.client_pubkey)
.map(x25519_dalek::PublicKey::from)
.map_err(|_| Status::invalid_argument("Invalid X25519 public key"))?;
let response = actor
.ask(HandleUnsealRequest { client_pubkey })
.await
.map_err(|err| {
warn!(error = ?err, "Failed to handle unseal start request");
Status::internal("Failed to start unseal flow")
})?;
Ok(Some(wrap_unseal_response(UnsealResponsePayload::Start(
proto_unseal::UnsealStartResponse {
server_pubkey: response.server_pubkey.as_bytes().to_vec(),
},
))))
}
async fn handle_unseal_encrypted_key(
actor: &ActorRef<UserAgentSession>,
req: ProtoUnsealEncryptedKey,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match actor
.ask(HandleUnsealEncryptedKey {
nonce: req.nonce,
ciphertext: req.ciphertext,
associated_data: req.associated_data,
})
.await
{
Ok(()) => ProtoUnsealResult::Success,
Err(SendError::HandlerError(UnsealError::InvalidKey)) => ProtoUnsealResult::InvalidKey,
Err(err) => {
warn!(error = ?err, "Failed to handle unseal request");
return Err(Status::internal("Failed to unseal vault"));
}
};
Ok(Some(wrap_unseal_response(UnsealResponsePayload::Result(
result.into(),
))))
}
async fn handle_bootstrap_request(
actor: &ActorRef<UserAgentSession>,
req: proto_bootstrap::Request,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let encrypted_key = req
.encrypted_key
.ok_or_else(|| Status::invalid_argument("Missing bootstrap encrypted key"))?;
handle_bootstrap_encrypted_key(actor, encrypted_key).await
}
async fn handle_bootstrap_encrypted_key(
actor: &ActorRef<UserAgentSession>,
req: ProtoBootstrapEncryptedKey,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match actor
.ask(HandleBootstrapEncryptedKey {
nonce: req.nonce,
ciphertext: req.ciphertext,
associated_data: req.associated_data,
})
.await
{
Ok(()) => ProtoBootstrapResult::Success,
Err(SendError::HandlerError(BootstrapError::InvalidKey)) => {
ProtoBootstrapResult::InvalidKey
}
Err(SendError::HandlerError(BootstrapError::AlreadyBootstrapped)) => {
ProtoBootstrapResult::AlreadyBootstrapped
}
Err(err) => {
warn!(error = ?err, "Failed to handle bootstrap request");
return Err(Status::internal("Failed to bootstrap vault"));
}
};
Ok(Some(wrap_bootstrap_response(result)))
}
async fn handle_query_vault_state(
actor: &ActorRef<UserAgentSession>,
) -> Result<Option<UserAgentResponsePayload>, Status> {

View File

@@ -0,0 +1,151 @@
use arbiter_proto::proto::user_agent::{
user_agent_response::Payload as UserAgentResponsePayload,
vault::{
self as proto_vault,
bootstrap::{
self as proto_bootstrap, BootstrapResult as ProtoBootstrapResult,
},
request::Payload as VaultRequestPayload,
response::Payload as VaultResponsePayload,
unseal::{
self as proto_unseal, UnsealResult as ProtoUnsealResult, UnsealStart,
request::Payload as UnsealRequestPayload, response::Payload as UnsealResponsePayload,
},
},
};
use kameo::{actor::ActorRef, error::SendError};
use tonic::Status;
use tracing::warn;
use crate::peers::user_agent::vault_gate::{
self as vault_gate, HandleBootstrapEncryptedKey, HandleHandshake, HandleUnsealEncryptedKey,
VaultGate,
};
fn wrap_vault_response(payload: VaultResponsePayload) -> UserAgentResponsePayload {
UserAgentResponsePayload::Vault(proto_vault::Response {
payload: Some(payload),
})
}
fn wrap_unseal_response(payload: UnsealResponsePayload) -> UserAgentResponsePayload {
wrap_vault_response(VaultResponsePayload::Unseal(proto_unseal::Response {
payload: Some(payload),
}))
}
fn wrap_bootstrap_response(result: ProtoBootstrapResult) -> UserAgentResponsePayload {
wrap_vault_response(VaultResponsePayload::Bootstrap(proto_bootstrap::Response {
result: result.into(),
}))
}
pub(super) async fn dispatch(
gate: &ActorRef<VaultGate>,
req: proto_vault::Request,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let Some(payload) = req.payload else {
return Err(Status::invalid_argument("Missing vault request payload"));
};
match payload {
VaultRequestPayload::QueryState(_) => {
use arbiter_proto::proto::shared::VaultState as ProtoVaultState;
Ok(Some(wrap_vault_response(VaultResponsePayload::State(
ProtoVaultState::Sealed.into(),
))))
}
VaultRequestPayload::Unseal(req) => dispatch_unseal(gate, req).await,
VaultRequestPayload::Bootstrap(req) => dispatch_bootstrap(gate, req).await,
}
}
async fn dispatch_unseal(
gate: &ActorRef<VaultGate>,
req: proto_unseal::Request,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let Some(payload) = req.payload else {
return Err(Status::invalid_argument("Missing unseal request payload"));
};
match payload {
UnsealRequestPayload::Start(req) => handle_unseal_start(gate, req).await,
UnsealRequestPayload::EncryptedKey(req) => handle_unseal_encrypted_key(gate, req).await,
}
}
async fn handle_unseal_start(
gate: &ActorRef<VaultGate>,
req: UnsealStart,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let client_pubkey = <[u8; 32]>::try_from(req.client_pubkey)
.map(x25519_dalek::PublicKey::from)
.map_err(|_| Status::invalid_argument("Invalid X25519 public key"))?;
let response = gate
.ask(HandleHandshake { client_pubkey })
.await
.map_err(|err| {
warn!(error = ?err, "Failed to handle unseal start");
Status::internal("Failed to start unseal flow")
})?;
Ok(Some(wrap_unseal_response(UnsealResponsePayload::Start(
proto_unseal::UnsealStartResponse {
server_pubkey: response.server_pubkey.as_bytes().to_vec(),
},
))))
}
async fn handle_unseal_encrypted_key(
gate: &ActorRef<VaultGate>,
req: arbiter_proto::proto::user_agent::vault::unseal::UnsealEncryptedKey,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let result = match gate
.ask(HandleUnsealEncryptedKey {
nonce: req.nonce,
ciphertext: req.ciphertext,
associated_data: req.associated_data,
})
.await
{
Ok(()) => ProtoUnsealResult::Success,
Err(SendError::HandlerError(vault_gate::Error::InvalidKey)) => ProtoUnsealResult::InvalidKey,
Err(err) => {
warn!(error = ?err, "Failed to handle unseal request");
return Err(Status::internal("Failed to unseal vault"));
}
};
Ok(Some(wrap_unseal_response(UnsealResponsePayload::Result(
result.into(),
))))
}
async fn dispatch_bootstrap(
gate: &ActorRef<VaultGate>,
req: proto_bootstrap::Request,
) -> Result<Option<UserAgentResponsePayload>, Status> {
let encrypted_key = req
.encrypted_key
.ok_or_else(|| Status::invalid_argument("Missing bootstrap encrypted key"))?;
let result = match gate
.ask(HandleBootstrapEncryptedKey {
nonce: encrypted_key.nonce,
ciphertext: encrypted_key.ciphertext,
associated_data: encrypted_key.associated_data,
})
.await
{
Ok(()) => ProtoBootstrapResult::Success,
Err(SendError::HandlerError(vault_gate::Error::InvalidKey)) => ProtoBootstrapResult::InvalidKey,
Err(SendError::HandlerError(vault_gate::Error::AlreadyBootstrapped)) => {
ProtoBootstrapResult::AlreadyBootstrapped
}
Err(err) => {
warn!(error = ?err, "Failed to handle bootstrap request");
return Err(Status::internal("Failed to bootstrap vault"));
}
};
Ok(Some(wrap_bootstrap_response(result)))
}