diff --git a/keep-enclave/enclave/src/vsock_server.rs b/keep-enclave/enclave/src/vsock_server.rs index 94ab2dbf..a74fd006 100644 --- a/keep-enclave/enclave/src/vsock_server.rs +++ b/keep-enclave/enclave/src/vsock_server.rs @@ -3,7 +3,7 @@ use crate::error::{EnclaveError, Result}; use crate::kms::{EncryptedWallet, EnclaveKms}; use crate::policy::{PolicyConfig, PolicyDecision, PolicyEngine, SigningContext}; -use crate::signer::EnclaveSigner; +use crate::signer::{EnclaveSigner, PsbtAnalysis}; use bitcoin::Network; use serde::{Deserialize, Serialize}; use std::time::{SystemTime, UNIX_EPOCH}; @@ -448,13 +448,44 @@ impl VsockServer { } }; + if let Some(denied) = + self.enforce_destination_policies(&analysis, &req.key_id, total_spend, timestamp) + { + return denied; + } + + match self.signer.sign_psbt(&req.key_id, &req.psbt, network) { + Ok((signed_psbt, signed_inputs, _)) => { + self.policy_engine.record_operation(&req.key_id); + EnclaveResponse::SignedPsbt { + psbt: signed_psbt, + signed_inputs, + } + } + Err(e) => EnclaveResponse::Error { + code: ErrorCode::SigningFailed, + message: e.to_string(), + }, + } + } + + /// Evaluate the policy engine against each non-change destination. Returns the + /// error response to send back if any destination is denied or requires + /// approval, or `None` when every destination is allowed. + fn enforce_destination_policies( + &mut self, + analysis: &PsbtAnalysis, + key_id: &str, + total_spend: u64, + timestamp: u64, + ) -> Option { for dest in &analysis.destinations { if dest.is_change { continue; } let addr = dest.address.as_deref(); let ctx = SigningContext { - key_id: &req.key_id, + key_id, amount_sats: Some(total_spend), destination: addr, event_kind: None, @@ -464,33 +495,20 @@ impl VsockServer { match self.policy_engine.evaluate(&ctx) { PolicyDecision::Allow => {} PolicyDecision::Deny(reason) => { - return EnclaveResponse::Error { + return Some(EnclaveResponse::Error { code: ErrorCode::PolicyDenied, message: format!("Destination {:?} denied: {}", addr, reason), - }; + }); } PolicyDecision::RequireApproval => { - return EnclaveResponse::Error { + return Some(EnclaveResponse::Error { code: ErrorCode::PolicyDenied, message: format!("Destination {:?} requires approval", addr), - }; - } - } - } - - match self.signer.sign_psbt(&req.key_id, &req.psbt, network) { - Ok((signed_psbt, signed_inputs, _)) => { - self.policy_engine.record_operation(&req.key_id); - EnclaveResponse::SignedPsbt { - psbt: signed_psbt, - signed_inputs, + }); } } - Err(e) => EnclaveResponse::Error { - code: ErrorCode::SigningFailed, - message: e.to_string(), - }, } + None } fn handle_set_policy(&mut self, config: PolicyConfig) -> EnclaveResponse {