From 3b828d5874747e0123f3de88e9b66a05bceb0ade Mon Sep 17 00:00:00 2001 From: Skipper Date: Thu, 16 Apr 2026 22:14:55 +0200 Subject: [PATCH] refactor(server::grpc::vault_gate): standard approach using / traits --- .../src/grpc/user_agent/vault_gate.rs | 201 ++---------------- .../src/grpc/user_agent/vault_gate/inbound.rs | 129 +++++++++++ .../grpc/user_agent/vault_gate/outbound.rs | 115 ++++++++++ 3 files changed, 257 insertions(+), 188 deletions(-) create mode 100644 server/crates/arbiter-server/src/grpc/user_agent/vault_gate/inbound.rs create mode 100644 server/crates/arbiter-server/src/grpc/user_agent/vault_gate/outbound.rs diff --git a/server/crates/arbiter-server/src/grpc/user_agent/vault_gate.rs b/server/crates/arbiter-server/src/grpc/user_agent/vault_gate.rs index 381fc45..37f3c5b 100644 --- a/server/crates/arbiter-server/src/grpc/user_agent/vault_gate.rs +++ b/server/crates/arbiter-server/src/grpc/user_agent/vault_gate.rs @@ -1,55 +1,16 @@ -use arbiter_proto::{ - proto::{ - shared::VaultState as ProtoVaultState, - user_agent::{ - user_agent_request::Payload as UserAgentRequestPayload, - 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, - request::Payload as UnsealRequestPayload, - response::Payload as UnsealResponsePayload, - }, - }, - }, - }, - transport::{Bi, Error as TransportError, Receiver, Sender}, -}; +use arbiter_proto::transport::{Bi, Error as TransportError, Receiver, Sender}; use async_trait::async_trait; use tonic::Status; use tracing::warn; use super::auth::AuthTransportAdapter; use crate::{ - actors::vault::VaultState, - peers::user_agent::vault_gate::{ - self as vault_gate, HandleBootstrapEncryptedKey, HandleHandshake, HandleUnsealEncryptedKey, - }, + grpc::TryConvert, + peers::user_agent::vault_gate::{self as vault_gate}, }; -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(), - })) -} - -impl AuthTransportAdapter<'_> {} +mod inbound; +mod outbound; #[async_trait] impl Receiver for AuthTransportAdapter<'_> { @@ -79,87 +40,12 @@ impl Receiver for AuthTransportAdapter<'_> { return None; }; - let vault_req = match payload { - UserAgentRequestPayload::Vault(req) => req, - _ => { - let _ = self - .bi_mut() - .send(Err(Status::permission_denied( - "Only vault operations are permitted before unsealing", - ))) - .await; + match payload.try_convert() { + Ok(inbound) => return Some(inbound), + Err(status) => { + let _ = self.bi_mut().send(Err(status)).await; return None; } - }; - - let Some(vault_payload) = vault_req.payload else { - let _ = self - .bi_mut() - .send(Err(Status::invalid_argument( - "Missing vault request payload", - ))) - .await; - return None; - }; - - match vault_payload { - VaultRequestPayload::QueryState(_) => { - return Some(vault_gate::Inbound::HandleVaultState); - } - VaultRequestPayload::Unseal(req) => { - let Some(unseal_payload) = req.payload else { - let _ = self - .bi_mut() - .send(Err(Status::invalid_argument( - "Missing unseal request payload", - ))) - .await; - return None; - }; - match unseal_payload { - UnsealRequestPayload::Start(start) => { - let Ok(bytes) = <[u8; 32]>::try_from(start.client_pubkey) else { - let _ = self - .bi_mut() - .send(Err(Status::invalid_argument( - "Invalid X25519 public key", - ))) - .await; - return None; - }; - return Some(vault_gate::Inbound::HandleHandshake(HandleHandshake { - client_pubkey: x25519_dalek::PublicKey::from(bytes), - })); - } - UnsealRequestPayload::EncryptedKey(key) => { - return Some(vault_gate::Inbound::HandleUnsealEncryptedKey( - HandleUnsealEncryptedKey { - nonce: key.nonce, - ciphertext: key.ciphertext, - associated_data: key.associated_data, - }, - )); - } - } - } - VaultRequestPayload::Bootstrap(req) => { - let Some(encrypted_key) = req.encrypted_key else { - let _ = self - .bi_mut() - .send(Err(Status::invalid_argument( - "Missing bootstrap encrypted key", - ))) - .await; - return None; - }; - return Some(vault_gate::Inbound::HandleBootstrapEncryptedKey( - HandleBootstrapEncryptedKey { - nonce: encrypted_key.nonce, - ciphertext: encrypted_key.ciphertext, - associated_data: encrypted_key.associated_data, - }, - )); - } } } } @@ -182,71 +68,10 @@ impl Sender> for AuthTransportAd } }; - let payload = match outbound { - vault_gate::Outbound::HandleVaultState(result) => match result { - Ok(state) => { - let state = match state { - VaultState::Unbootstrapped => ProtoVaultState::Unbootstrapped, - VaultState::Sealed => ProtoVaultState::Sealed, - VaultState::Unsealed => ProtoVaultState::Unsealed, - }; - - wrap_vault_response(VaultResponsePayload::State(state.into())) - } - Err(err) => { - warn!(?err, "vault state query failed"); - return self - .bi_mut() - .send(Err(Status::internal("Failed to query vault state"))) - .await; - } - }, - vault_gate::Outbound::HandleHandshake(Ok(response)) => wrap_unseal_response( - UnsealResponsePayload::Start(proto_unseal::UnsealStartResponse { - server_pubkey: response.server_pubkey.as_bytes().to_vec(), - }), - ), - vault_gate::Outbound::HandleHandshake(Err(err)) => { - warn!(?err, "handshake failed"); - return self - .bi_mut() - .send(Err(Status::internal("Failed to start unseal flow"))) - .await; - } - vault_gate::Outbound::HandleUnsealEncryptedKey(result) => { - let proto_result = match result { - Ok(()) => ProtoUnsealResult::Success, - Err(vault_gate::Error::InvalidKey) => ProtoUnsealResult::InvalidKey, - Err(err) => { - warn!(?err, "unseal failed"); - return self - .bi_mut() - .send(Err(Status::internal("Failed to unseal vault"))) - .await; - } - }; - wrap_unseal_response(UnsealResponsePayload::Result(proto_result.into())) - } - vault_gate::Outbound::HandleBootstrapEncryptedKey(result) => { - let proto_result = match result { - Ok(()) => ProtoBootstrapResult::Success, - Err(vault_gate::Error::InvalidKey) => ProtoBootstrapResult::InvalidKey, - Err(vault_gate::Error::AlreadyBootstrapped) => { - ProtoBootstrapResult::AlreadyBootstrapped - } - Err(err) => { - warn!(?err, "bootstrap failed"); - return self - .bi_mut() - .send(Err(Status::internal("Failed to bootstrap vault"))) - .await; - } - }; - wrap_bootstrap_response(proto_result) - } - }; - - self.send_response_payload(payload).await + match outbound.try_convert() { + Ok(payload) => self.send_response_payload(payload).await, + Err(status) => self.bi_mut().send(Err(status)).await, + } } } diff --git a/server/crates/arbiter-server/src/grpc/user_agent/vault_gate/inbound.rs b/server/crates/arbiter-server/src/grpc/user_agent/vault_gate/inbound.rs new file mode 100644 index 0000000..68761ae --- /dev/null +++ b/server/crates/arbiter-server/src/grpc/user_agent/vault_gate/inbound.rs @@ -0,0 +1,129 @@ +use arbiter_proto::proto::user_agent::{ + user_agent_request::Payload as UserAgentRequestPayload, + vault::{ + self as proto_vault, + bootstrap::{self as proto_bootstrap}, + request::Payload as VaultRequestPayload, + unseal::{self as proto_unseal, request::Payload as UnsealRequestPayload}, + }, +}; +use tonic::Status; + +use crate::{ + grpc::{Convert, TryConvert}, + peers::user_agent::vault_gate::{ + self as vault_gate, HandleBootstrapEncryptedKey, HandleHandshake, HandleUnsealEncryptedKey, + }, +}; + +impl TryConvert for UserAgentRequestPayload { + type Output = vault_gate::Inbound; + type Error = Status; + + fn try_convert(self) -> Result { + match self { + UserAgentRequestPayload::Vault(req) => req.try_convert(), + _ => Err(Status::permission_denied( + "Only vault operations are permitted before unsealing", + )), + } + } +} + +impl TryConvert for proto_vault::Request { + type Output = vault_gate::Inbound; + type Error = Status; + + fn try_convert(self) -> Result { + self.payload + .ok_or_else(|| Status::invalid_argument("Missing vault request payload"))? + .try_convert() + } +} + +impl TryConvert for VaultRequestPayload { + type Output = vault_gate::Inbound; + type Error = Status; + + fn try_convert(self) -> Result { + match self { + VaultRequestPayload::QueryState(_) => Ok(vault_gate::Inbound::HandleVaultState), + VaultRequestPayload::Unseal(req) => req.try_convert(), + VaultRequestPayload::Bootstrap(req) => req.try_convert(), + } + } +} + +impl TryConvert for proto_unseal::Request { + type Output = vault_gate::Inbound; + type Error = Status; + + fn try_convert(self) -> Result { + self.payload + .ok_or_else(|| Status::invalid_argument("Missing unseal request payload"))? + .try_convert() + } +} + +impl TryConvert for UnsealRequestPayload { + type Output = vault_gate::Inbound; + type Error = Status; + + fn try_convert(self) -> Result { + match self { + UnsealRequestPayload::Start(start) => start.try_convert(), + UnsealRequestPayload::EncryptedKey(key) => Ok(key.convert()), + } + } +} + +impl TryConvert for proto_unseal::UnsealStart { + type Output = vault_gate::Inbound; + type Error = Status; + + fn try_convert(self) -> Result { + let bytes = <[u8; 32]>::try_from(self.client_pubkey) + .map_err(|_| Status::invalid_argument("Invalid X25519 public key"))?; + Ok(vault_gate::Inbound::HandleHandshake(HandleHandshake { + client_pubkey: x25519_dalek::PublicKey::from(bytes), + })) + } +} + +impl Convert for proto_unseal::UnsealEncryptedKey { + type Output = vault_gate::Inbound; + + fn convert(self) -> vault_gate::Inbound { + vault_gate::Inbound::HandleUnsealEncryptedKey(HandleUnsealEncryptedKey { + nonce: self.nonce, + ciphertext: self.ciphertext, + associated_data: self.associated_data, + }) + } +} + +impl TryConvert for proto_bootstrap::Request { + type Output = vault_gate::Inbound; + type Error = Status; + + fn try_convert(self) -> Result { + self.encrypted_key + .ok_or_else(|| Status::invalid_argument("Missing bootstrap encrypted key"))? + .try_convert() + } +} + +impl TryConvert for proto_bootstrap::BootstrapEncryptedKey { + type Output = vault_gate::Inbound; + type Error = Status; + + fn try_convert(self) -> Result { + Ok(vault_gate::Inbound::HandleBootstrapEncryptedKey( + HandleBootstrapEncryptedKey { + nonce: self.nonce, + ciphertext: self.ciphertext, + associated_data: self.associated_data, + }, + )) + } +} diff --git a/server/crates/arbiter-server/src/grpc/user_agent/vault_gate/outbound.rs b/server/crates/arbiter-server/src/grpc/user_agent/vault_gate/outbound.rs new file mode 100644 index 0000000..c594b33 --- /dev/null +++ b/server/crates/arbiter-server/src/grpc/user_agent/vault_gate/outbound.rs @@ -0,0 +1,115 @@ +use arbiter_proto::proto::{ + shared::VaultState as ProtoVaultState, + user_agent::{ + user_agent_response::Payload as UserAgentResponsePayload, + vault::{ + self as proto_vault, + bootstrap::{self as proto_bootstrap, BootstrapResult as ProtoBootstrapResult}, + response::Payload as VaultResponsePayload, + unseal::{ + self as proto_unseal, UnsealResult as ProtoUnsealResult, + response::Payload as UnsealResponsePayload, + }, + }, + }, +}; +use tonic::Status; +use tracing::warn; + +use crate::{ + actors::vault::VaultState, + grpc::{Convert, TryConvert}, + peers::user_agent::vault_gate::{self as vault_gate}, +}; + +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(), + })) +} + +impl Convert for VaultState { + type Output = UserAgentResponsePayload; + + fn convert(self) -> UserAgentResponsePayload { + let proto_state = match self { + VaultState::Unbootstrapped => ProtoVaultState::Unbootstrapped, + VaultState::Sealed => ProtoVaultState::Sealed, + VaultState::Unsealed => ProtoVaultState::Unsealed, + }; + wrap_vault_response(VaultResponsePayload::State(proto_state.into())) + } +} + +impl Convert for vault_gate::HandshakeResponse { + type Output = UserAgentResponsePayload; + + fn convert(self) -> UserAgentResponsePayload { + wrap_unseal_response(UnsealResponsePayload::Start( + proto_unseal::UnsealStartResponse { + server_pubkey: self.server_pubkey.as_bytes().to_vec(), + }, + )) + } +} + +impl TryConvert for vault_gate::Outbound { + type Output = UserAgentResponsePayload; + type Error = Status; + + fn try_convert(self) -> Result { + match self { + vault_gate::Outbound::HandleVaultState(result) => result + .map_err(|err| { + warn!(?err, "vault state query failed"); + Status::internal("Failed to query vault state") + }) + .map(VaultState::convert), + vault_gate::Outbound::HandleHandshake(result) => result + .map_err(|err| { + warn!(?err, "handshake failed"); + Status::internal("Failed to start unseal flow") + }) + .map(vault_gate::HandshakeResponse::convert), + vault_gate::Outbound::HandleUnsealEncryptedKey(result) => { + let proto_result = match result { + Ok(()) => ProtoUnsealResult::Success, + Err(vault_gate::Error::InvalidKey) => ProtoUnsealResult::InvalidKey, + Err(err) => { + warn!(?err, "unseal failed"); + return Err(Status::internal("Failed to unseal vault")); + } + }; + Ok(wrap_unseal_response(UnsealResponsePayload::Result( + proto_result.into(), + ))) + } + vault_gate::Outbound::HandleBootstrapEncryptedKey(result) => { + let proto_result = match result { + Ok(()) => ProtoBootstrapResult::Success, + Err(vault_gate::Error::InvalidKey) => ProtoBootstrapResult::InvalidKey, + Err(vault_gate::Error::AlreadyBootstrapped) => { + ProtoBootstrapResult::AlreadyBootstrapped + } + Err(err) => { + warn!(?err, "bootstrap failed"); + return Err(Status::internal("Failed to bootstrap vault")); + } + }; + Ok(wrap_bootstrap_response(proto_result)) + } + } + } +}