From d9d08bf7668931fd6616bd7d27cbf406b1218ee8 Mon Sep 17 00:00:00 2001 From: Siddharth Sharma Date: Thu, 23 Apr 2026 04:54:48 -0700 Subject: [PATCH 1/8] [guardian] in-process e2e harness + TestNetworksBuilder.with_guardian() --- Cargo.lock | 3 + crates/e2e-tests/Cargo.toml | 6 + crates/e2e-tests/src/e2e_flow.rs | 100 ++++ crates/e2e-tests/src/guardian_harness.rs | 177 +++++++ crates/e2e-tests/src/hashi_network.rs | 10 + crates/e2e-tests/src/lib.rs | 68 ++- crates/hashi-guardian/Cargo.toml | 8 +- crates/hashi-guardian/src/enclave.rs | 480 ++++++++++++++++++ crates/hashi-guardian/src/getters.rs | 2 +- crates/hashi-guardian/src/lib.rs | 24 + crates/hashi-guardian/src/main.rs | 592 +---------------------- crates/hashi-guardian/src/test_utils.rs | 161 ++++++ crates/hashi/src/config.rs | 9 + crates/hashi/src/grpc/guardian_client.rs | 40 ++ crates/hashi/src/grpc/mod.rs | 1 + crates/hashi/src/lib.rs | 31 ++ 16 files changed, 1121 insertions(+), 591 deletions(-) create mode 100644 crates/e2e-tests/src/guardian_harness.rs create mode 100644 crates/hashi-guardian/src/enclave.rs create mode 100644 crates/hashi-guardian/src/test_utils.rs create mode 100644 crates/hashi/src/grpc/guardian_client.rs diff --git a/Cargo.lock b/Cargo.lock index b612f7a20..1235928c2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2089,6 +2089,7 @@ dependencies = [ "fastcrypto-tbls", "futures", "hashi", + "hashi-guardian", "hashi-screener", "hashi-types", "nix", @@ -2105,6 +2106,8 @@ dependencies = [ "sui-transaction-builder", "tempfile", "tokio", + "tokio-stream", + "tonic", "tracing", "tracing-subscriber", ] diff --git a/crates/e2e-tests/Cargo.toml b/crates/e2e-tests/Cargo.toml index 7d744d360..e3c945147 100644 --- a/crates/e2e-tests/Cargo.toml +++ b/crates/e2e-tests/Cargo.toml @@ -5,6 +5,10 @@ edition = "2024" [dependencies] hashi = { path = "../hashi" } +hashi-guardian = { path = "../hashi-guardian", features = [ + "test-utils", + "non-enclave-dev", +] } hashi-screener = { path = "../hashi-screener", features = ["test-utils"] } hashi-types = { path = "../hashi-types" } @@ -32,6 +36,8 @@ clap.workspace = true colored.workspace = true rand.workspace = true nix = { version = "0.26.4", features = ["signal"] } +tokio-stream = "0.1" +tonic.workspace = true tracing-subscriber.workspace = true [[bin]] diff --git a/crates/e2e-tests/src/e2e_flow.rs b/crates/e2e-tests/src/e2e_flow.rs index 06f9bb91e..7ebe650be 100644 --- a/crates/e2e-tests/src/e2e_flow.rs +++ b/crates/e2e-tests/src/e2e_flow.rs @@ -339,6 +339,106 @@ mod tests { Ok(()) } + /// Same flow as `test_bitcoin_withdrawal_e2e_flow`, but with an + /// in-process guardian in the loop. + #[tokio::test] + async fn test_bitcoin_withdrawal_with_guardian_e2e_flow() -> Result<()> { + init_test_logging(); + info!("=== Starting Bitcoin Withdrawal E2E Test (with guardian) ==="); + + let mut networks = TestNetworksBuilder::new() + .with_nodes(4) + .with_guardian() + .build() + .await?; + + info!("Waiting for MPC key to be ready..."); + networks.hashi_network.nodes()[0] + .wait_for_mpc_key(Duration::from_secs(60)) + .await?; + info!("MPC key ready"); + + for node in networks.hashi_network.nodes() { + assert!( + node.hashi().config.guardian_endpoint().is_some(), + "guardian endpoint should be configured on every test node" + ); + assert!( + node.hashi().guardian_client().is_some(), + "guardian client should be populated after Hashi::start()" + ); + } + let harness = networks + .guardian_harness + .as_ref() + .expect("harness present when .with_guardian() is set"); + assert!( + harness.enclave().is_fully_initialized(), + "guardian harness should have reached fully-initialized state" + ); + + let deposit_amount_sats = 100_000u64; + let hbtc_recipient = create_deposit_and_wait(&mut networks, deposit_amount_sats).await?; + + let hashi = networks.hashi_network.nodes()[0].hashi().clone(); + let user_key = networks.sui_network.user_keys.first().unwrap(); + let withdrawal_amount_sats = 30_000u64; + let btc_destination = networks.bitcoin_node.get_new_address()?; + let destination_bytes = extract_witness_program(&btc_destination)?; + info!( + "Requesting withdrawal of {} sats to {}", + withdrawal_amount_sats, btc_destination + ); + + let mut withdrawal_executor = + SuiTxExecutor::from_config(&hashi.config, hashi.onchain_state())? + .with_signer(user_key.clone()); + withdrawal_executor + .execute_create_withdrawal_request(withdrawal_amount_sats, destination_bytes) + .await?; + + let miner = BackgroundMiner::start(&networks.bitcoin_node); + + let confirmed_event = wait_for_withdrawal_confirmation( + &mut networks.sui_network.client, + Duration::from_secs(60), + ) + .await?; + info!("Withdrawal confirmed on Sui"); + + drop(miner); + + let hbtc_balance_after = get_hbtc_balance( + &mut networks.sui_network.client, + networks.hashi_network.ids().package_id, + hbtc_recipient, + ) + .await?; + let expected_remaining = deposit_amount_sats - withdrawal_amount_sats; + assert_eq!( + hbtc_balance_after, expected_remaining, + "Expected remaining hBTC after withdrawal" + ); + + let withdrawal_txid: Txid = confirmed_event.txid.into(); + let max_output = Amount::from_sat(withdrawal_amount_sats); + let min_output = Amount::from_sat( + withdrawal_amount_sats.saturating_sub(hashi.onchain_state().worst_case_network_fee()), + ); + wait_for_withdrawal_tx_success( + &networks.bitcoin_node, + &withdrawal_txid, + &btc_destination, + max_output, + min_output, + Duration::from_secs(30), + ) + .await?; + + info!("=== Bitcoin Withdrawal E2E Test (with guardian) Passed ==="); + Ok(()) + } + async fn withdraw_and_confirm( networks: &mut TestNetworks, hashi: &hashi::Hashi, diff --git a/crates/e2e-tests/src/guardian_harness.rs b/crates/e2e-tests/src/guardian_harness.rs new file mode 100644 index 000000000..656d4a07d --- /dev/null +++ b/crates/e2e-tests/src/guardian_harness.rs @@ -0,0 +1,177 @@ +// Copyright (c) Mysten Labs, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//! In-process `hashi-guardian` for integration tests. Started in two +//! stages because hashi-side DKG is a prerequisite for provisioner-init: +//! [`GuardianHarness::start`] serves gRPC immediately; [`GuardianHarness::finalize`] +//! completes provisioner-init once DKG output is on chain. + +use anyhow::Context; +use anyhow::Result; +use bitcoin::Network; +use hashi_guardian::Enclave; +use hashi_guardian::OperatorInitTestArgs; +use hashi_guardian::create_operator_initialized_enclave; +use hashi_guardian::rpc::GuardianGrpc; +use hashi_types::committee::Committee as HashiCommittee; +use hashi_types::guardian::BitcoinPubkey; +use hashi_types::guardian::LimiterState; +use hashi_types::guardian::ProvisionerInitState; +use hashi_types::guardian::WithdrawalConfig; +use hashi_types::proto::guardian_service_server::GuardianServiceServer; +use std::net::SocketAddr; +use std::sync::Arc; +use tokio::net::TcpListener; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; +use tonic::transport::Server; + +/// In-process guardian reachable over gRPC on a local TCP socket. Drop +/// shuts the server down. +pub struct GuardianHarness { + enclave: Arc, + endpoint: String, + network: Network, + shutdown_tx: Option>, + server_handle: Option>, +} + +impl GuardianHarness { + /// Start an operator-init'd guardian. Withdrawal RPCs stay gated + /// on [`Self::finalize`] completing provisioner-init. + pub async fn start(network: Network) -> Result { + let enclave = create_operator_initialized_enclave( + OperatorInitTestArgs::default().with_network(network), + ) + .await; + + let listener = TcpListener::bind("127.0.0.1:0") + .await + .context("bind guardian harness listener")?; + let addr: SocketAddr = listener.local_addr()?; + let endpoint = format!("http://{addr}"); + + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + let svc = GuardianGrpc { + enclave: enclave.clone(), + setup_mode: false, + }; + let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener); + let server_handle = tokio::spawn(async move { + let result = Server::builder() + .add_service(GuardianServiceServer::new(svc)) + .serve_with_incoming_shutdown(incoming, async move { + let _ = shutdown_rx.await; + }) + .await; + if let Err(e) = result { + tracing::warn!("guardian harness server exited: {e}"); + } + }); + + Ok(Self { + enclave, + endpoint, + network, + shutdown_tx: Some(shutdown_tx), + server_handle: Some(server_handle), + }) + } + + /// Complete provisioner-init with the committee + BTC master pubkey + /// from hashi's DKG. + pub async fn finalize( + &self, + committee: HashiCommittee, + master_pubkey: BitcoinPubkey, + withdrawal_config: WithdrawalConfig, + limiter_state: LimiterState, + ) -> Result<()> { + use bitcoin::secp256k1::Keypair; + use bitcoin::secp256k1::Secp256k1; + use bitcoin::secp256k1::SecretKey; + use rand::RngCore; + + let secp = Secp256k1::new(); + let mut sk_bytes = [0u8; 32]; + rand::thread_rng().fill_bytes(&mut sk_bytes); + let enclave_btc_keypair = Keypair::from_secret_key( + &secp, + &SecretKey::from_slice(&sk_bytes).expect("random bytes form a valid secp256k1 key"), + ); + self.enclave + .config + .set_btc_keypair(enclave_btc_keypair) + .context("set enclave btc keypair")?; + self.enclave + .config + .set_hashi_btc_pk(master_pubkey) + .context("set hashi btc master pubkey")?; + self.enclave + .config + .set_withdrawal_config(withdrawal_config) + .context("set withdrawal config")?; + + let init_state = + ProvisionerInitState::new(committee, withdrawal_config, limiter_state, master_pubkey) + .context("valid ProvisionerInitState")?; + self.enclave + .state + .init(init_state) + .context("init enclave state")?; + + self.enclave + .scratchpad + .provisioner_init_logging_complete + .set(()) + .map_err(|_| anyhow::anyhow!("provisioner_init already finalized"))?; + + anyhow::ensure!( + self.enclave.is_fully_initialized(), + "guardian did not reach fully-initialized state" + ); + Ok(()) + } + + pub fn endpoint(&self) -> &str { + &self.endpoint + } + + pub fn enclave(&self) -> &Arc { + &self.enclave + } + + #[allow(dead_code)] + pub fn network(&self) -> Network { + self.network + } +} + +impl Drop for GuardianHarness { + fn drop(&mut self) { + if let Some(tx) = self.shutdown_tx.take() { + let _ = tx.send(()); + } + if let Some(handle) = self.server_handle.take() { + handle.abort(); + } + } +} + +pub fn default_test_withdrawal_config(committee: &HashiCommittee) -> WithdrawalConfig { + let total_weight = committee.total_weight(); + let committee_threshold = total_weight.div_ceil(3) * 2; + WithdrawalConfig { + committee_threshold, + refill_rate_sats_per_sec: 0, + max_bucket_capacity_sats: 100_000_000, + } +} + +pub fn full_bucket(config: &WithdrawalConfig) -> LimiterState { + LimiterState { + num_tokens_available: config.max_bucket_capacity_sats, + last_updated_at: 0, + next_seq: 0, + } +} diff --git a/crates/e2e-tests/src/hashi_network.rs b/crates/e2e-tests/src/hashi_network.rs index 7303703a9..dff4b5084 100644 --- a/crates/e2e-tests/src/hashi_network.rs +++ b/crates/e2e-tests/src/hashi_network.rs @@ -248,6 +248,7 @@ pub struct HashiNetworkBuilder { /// Node index whose shares should be corrupted by all other nodes, /// triggering the complaint recovery flow. pub test_corrupt_shares_target: Option, + pub guardian_endpoint: Option, } impl HashiNetworkBuilder { @@ -261,9 +262,15 @@ impl HashiNetworkBuilder { withdrawal_max_batch_size: None, max_mempool_chain_depth: None, test_corrupt_shares_target: None, + guardian_endpoint: None, } } + pub fn with_guardian_endpoint(mut self, endpoint: impl Into) -> Self { + self.guardian_endpoint = Some(endpoint.into()); + self + } + pub fn with_num_nodes(mut self, num_nodes: usize) -> Self { self.num_nodes = num_nodes; self @@ -363,6 +370,9 @@ impl HashiNetworkBuilder { config.bitcoin_chain_id = Some(hashi::constants::BITCOIN_REGTEST_CHAIN_ID.to_string()); config.sui_chain_id = service_info.chain_id.clone(); config.screener_endpoint = Some(screener_endpoint.clone()); + if let Some(ref guardian_endpoint) = self.guardian_endpoint { + config.guardian_endpoint = Some(guardian_endpoint.clone()); + } config.db = Some(dir.join(validator_address.to_string())); configs.push(config); } diff --git a/crates/e2e-tests/src/lib.rs b/crates/e2e-tests/src/lib.rs index b882b70f8..3dfae434f 100644 --- a/crates/e2e-tests/src/lib.rs +++ b/crates/e2e-tests/src/lib.rs @@ -21,6 +21,7 @@ use anyhow::Result; pub mod backup_restore; pub mod bitcoin_node; pub mod e2e_flow; +pub mod guardian_harness; pub mod hashi_network; mod publish; pub mod sui_network; @@ -45,6 +46,7 @@ pub struct TestNetworks { pub sui_network: SuiNetworkHandle, pub hashi_network: HashiNetwork, pub bitcoin_node: BitcoinNodeHandle, + pub guardian_harness: Option, } impl TestNetworks { @@ -95,6 +97,7 @@ pub struct TestNetworksBuilder { /// On-chain config overrides applied after DKG completes, before `build()` /// returns. Each entry is run through the full propose/vote/execute flow. onchain_config_overrides: Vec<(String, hashi_types::move_types::ConfigValue)>, + with_guardian: bool, } impl TestNetworksBuilder { @@ -104,9 +107,17 @@ impl TestNetworksBuilder { hashi_builder: HashiNetworkBuilder::new(), bitcoin_builder: BitcoinNodeBuilder::new(), onchain_config_overrides: Vec::new(), + with_guardian: false, } } + /// Spin up an in-process guardian, wire its endpoint into every + /// hashi node, and finalize it once DKG completes. + pub fn with_guardian(mut self) -> Self { + self.with_guardian = true; + self + } + pub fn with_nodes(mut self, num_nodes: usize) -> Self { self = self.with_hashi_nodes(num_nodes); self = self.with_sui_validators(num_nodes); @@ -221,8 +232,21 @@ impl TestNetworksBuilder { ) .await?; - let hashi_network = self - .hashi_builder + let mut hashi_builder = self.hashi_builder; + let guardian_harness = if self.with_guardian { + let harness = + guardian_harness::GuardianHarness::start(bitcoin::Network::Regtest).await?; + hashi_builder = hashi_builder.with_guardian_endpoint(harness.endpoint().to_string()); + tracing::info!( + endpoint = %harness.endpoint(), + "guardian harness started (operator-init)" + ); + Some(harness) + } else { + None + }; + + let hashi_network = hashi_builder .build( &dir.path().join("hashi"), &sui_network, @@ -236,6 +260,7 @@ impl TestNetworksBuilder { sui_network, hashi_network, bitcoin_node, + guardian_harness, }; tracing::info!("rpc url: {}", test_networks.sui_network().rpc_url); @@ -245,6 +270,10 @@ impl TestNetworksBuilder { .await?; } + if test_networks.guardian_harness.is_some() { + finalize_guardian_harness(&mut test_networks).await?; + } + Ok(test_networks) } @@ -265,6 +294,41 @@ impl TestNetworksBuilder { } } +async fn finalize_guardian_harness(networks: &mut TestNetworks) -> Result<()> { + use crate::guardian_harness::default_test_withdrawal_config; + use crate::guardian_harness::full_bucket; + + let nodes = networks.hashi_network.nodes(); + anyhow::ensure!( + !nodes.is_empty(), + "no hashi nodes to provision guardian from" + ); + + nodes[0] + .wait_for_mpc_key(std::time::Duration::from_secs(120)) + .await?; + + let hashi = nodes[0].hashi(); + let committee = hashi + .onchain_state() + .current_committee() + .ok_or_else(|| anyhow::anyhow!("no current committee after DKG"))?; + let master_pubkey = hashi.get_onchain_mpc_pubkey()?; + + let withdrawal_config = default_test_withdrawal_config(&committee); + let limiter_state = full_bucket(&withdrawal_config); + + let harness = networks + .guardian_harness + .as_ref() + .expect("guardian_harness set when finalize_guardian_harness is called"); + harness + .finalize(committee, master_pubkey, withdrawal_config, limiter_state) + .await?; + tracing::info!("guardian harness finalized"); + Ok(()) +} + /// Apply on-chain config overrides by running the full propose/vote/execute /// cycle for each `(key, value)` pair. Called from `TestNetworksBuilder::build` /// when overrides are present. diff --git a/crates/hashi-guardian/Cargo.toml b/crates/hashi-guardian/Cargo.toml index d7fb6ce23..70a132840 100644 --- a/crates/hashi-guardian/Cargo.toml +++ b/crates/hashi-guardian/Cargo.toml @@ -27,11 +27,15 @@ aws-config = "1.1.7" aws-sdk-s3 = "1.12.0" aws-credential-types = "1" +aws-smithy-mocks = { version = "0.1", optional = true } + [dev-dependencies] aws-sdk-s3 = { version = "1.12.0", features = ["test-util"] } aws-smithy-mocks = "0.1" [features] -# Stubs out the AWS Nitro NSM attestation call so the guardian can run outside -# a Nitro enclave (local dev, plain Kubernetes). Do not enable for prod builds. +# Stubs out the AWS Nitro NSM attestation call so the guardian can run +# outside a Nitro enclave. Do not enable for prod builds. non-enclave-dev = [] +# Exposes `test_utils` to external crates. +test-utils = ["dep:aws-smithy-mocks", "aws-sdk-s3/test-util"] diff --git a/crates/hashi-guardian/src/enclave.rs b/crates/hashi-guardian/src/enclave.rs new file mode 100644 index 000000000..7ad2a661e --- /dev/null +++ b/crates/hashi-guardian/src/enclave.rs @@ -0,0 +1,480 @@ +// Copyright (c) Mysten Labs, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//! Core enclave types: `Enclave` holds all guardian state (immutable config +//! set during operator/provisioner-init, mutable runtime state, and the +//! one-time init scratchpad). Lives in the library so external crates +//! (integration test harnesses, ops tooling) can construct and drive an +//! enclave without going through `main`. + +use bitcoin::secp256k1::Keypair; +use bitcoin::Network; +use bitcoin::Txid; +use hashi_types::guardian::bitcoin_utils::sign_btc_tx; +use hashi_types::guardian::bitcoin_utils::TxUTXOs; +use hashi_types::guardian::crypto::Share; +use hashi_types::guardian::GuardianError::InvalidInputs; +use hashi_types::guardian::*; +use hpke::Serializable; +use serde::Serialize; +use std::sync::Arc; +use std::sync::OnceLock; +use std::sync::RwLock; +use std::time::Duration; +use tracing::info; + +use crate::s3_logger::S3Logger; +use crate::withdraw::LimiterGuard; +use hashi_types::committee::Committee as HashiCommittee; + +/// Enclave's config & state +pub struct Enclave { + /// Immutable config (set once during init) + pub config: EnclaveConfig, + /// Mutable state + pub state: EnclaveState, + /// Initialization scratchpad + pub scratchpad: Scratchpad, +} + +/// Configuration set during initialization (immutable after set) +pub struct EnclaveConfig { + /// Ephemeral keypair (set on boot) + eph_keys: EphemeralKeyPairs, + /// S3 client & config (set in operator_init) + s3_logger: OnceLock, + /// Enclave BTC private key (set in provisioner_init) + pub(crate) enclave_btc_keypair: OnceLock, + /// BTC network: mainnet, testnet, regtest (set in operator_init) + btc_network: OnceLock, + /// Hashi BTC public key used to derive child keys (set in provisioner_init) + pub(crate) hashi_btc_master_pubkey: OnceLock, + /// Withdraw related config's (set in provisioner_init) + withdrawal_config: OnceLock, +} + +/// Mutable state that changes during operation. +/// Note: State is initialized during provisioner_init. +pub struct EnclaveState { + /// Current Hashi committee. + committee: RwLock>>, + /// Rate limiter. Set once during provisioner_init. + /// Uses `Arc` so the guard can be held across `.await`. + rate_limiter: OnceLock>>, +} + +/// Scratchpad used only during initialization. +/// Note: We don't clear it post-init because it does not have a lot of data. +#[derive(Default)] +pub struct Scratchpad { + /// The received shares + /// TODO: Investigate if it can be moved to std::sync::Mutex + pub shares: tokio::sync::Mutex>, + /// The share commitments + pub share_commitments: OnceLock, + /// Hash of the state in ProvisionerInitRequest + pub state_hash: OnceLock<[u8; 32]>, + /// Set once operator_init has successfully written all logs to S3. + /// This prevents heartbeats from being emitted before operator_init logs. + pub operator_init_logging_complete: OnceLock<()>, + /// Set once the provisioner init flow has successfully logged EnclaveFullyInitialized. + /// This prevents withdrawals from starting before provisioner_init logs. + pub provisioner_init_logging_complete: OnceLock<()>, +} + +pub struct EphemeralKeyPairs { + pub signing_keys: GuardianSignKeyPair, + pub encryption_keys: GuardianEncKeyPair, +} + +impl EnclaveConfig { + pub fn new(signing_keys: GuardianSignKeyPair, encryption_keys: GuardianEncKeyPair) -> Self { + EnclaveConfig { + eph_keys: EphemeralKeyPairs { + signing_keys, + encryption_keys, + }, + s3_logger: OnceLock::new(), + enclave_btc_keypair: OnceLock::new(), + btc_network: OnceLock::new(), + hashi_btc_master_pubkey: OnceLock::new(), + withdrawal_config: OnceLock::new(), + } + } + + // ======================================================================== + // Bitcoin Configuration + // ======================================================================== + + pub fn bitcoin_network(&self) -> GuardianResult { + self.btc_network + .get() + .copied() + .ok_or(InvalidInputs("Network is uninitialized".into())) + } + + pub fn set_bitcoin_network(&self, network: Network) -> GuardianResult<()> { + self.btc_network + .set(network) + .map_err(|_| InvalidInputs("Network is already initialized".into())) + } + + pub fn set_btc_keypair(&self, keypair: Keypair) -> GuardianResult<()> { + self.enclave_btc_keypair + .set(keypair) + .map_err(|_| InvalidInputs("Bitcoin key already set".into())) + } + + pub fn set_hashi_btc_pk(&self, pk: BitcoinPubkey) -> GuardianResult<()> { + self.hashi_btc_master_pubkey + .set(pk) + .map_err(|_| InvalidInputs("Hashi BTC key is already set".into())) + } + + /// Sign a BTC tx. Returns an Err if enclave btc keypair or hashi btc pk is not set. + pub fn btc_sign(&self, tx_utxos: &TxUTXOs) -> GuardianResult<(Txid, Vec)> { + let enclave_keypair = self + .enclave_btc_keypair + .get() + .ok_or(InvalidInputs("Bitcoin key is not initialized".into()))?; + let hashi_btc_pk = self + .hashi_btc_master_pubkey + .get() + .ok_or(InvalidInputs("Hashi BTC public key not set".into()))?; + + let enclave_btc_pk = enclave_keypair.x_only_public_key().0; + let (messages, txid) = tx_utxos.signing_messages_and_txid(&enclave_btc_pk, hashi_btc_pk); + Ok((txid, sign_btc_tx(&messages, enclave_keypair))) + } + + // ======================================================================== + // Withdrawal Configuration + // ======================================================================== + + pub fn withdrawal_config(&self) -> GuardianResult<&WithdrawalConfig> { + self.withdrawal_config + .get() + .ok_or(InvalidInputs("WithdrawalConfig is not initialized".into())) + } + + pub fn set_withdrawal_config(&self, config: WithdrawalConfig) -> GuardianResult<()> { + self.withdrawal_config + .set(config) + .map_err(|_| InvalidInputs("WithdrawalConfig already set".into())) + } + + pub fn committee_threshold(&self) -> GuardianResult { + Ok(self.withdrawal_config()?.committee_threshold) + } + + // ======================================================================== + // S3 Logger + // ======================================================================== + + pub fn s3_logger(&self) -> GuardianResult<&S3Logger> { + self.s3_logger + .get() + .ok_or(InvalidInputs("S3 logger is not initialized".into())) + } + + pub fn set_s3_logger(&self, logger: S3Logger) -> GuardianResult<()> { + self.s3_logger + .set(logger) + .map_err(|_| InvalidInputs("S3 logger already set".into())) + } + + // ======================================================================== + // Initialization Status + // ======================================================================== + + /// Check if operator_init configuration is complete (S3 logger and network) + pub fn is_operator_init_complete(&self) -> bool { + self.s3_logger.get().is_some() && self.btc_network.get().is_some() + } + + /// Check if any operator_init configuration has been set + pub fn is_operator_init_partially_complete(&self) -> bool { + self.s3_logger.get().is_some() || self.btc_network.get().is_some() + } + + /// Check if provisioner_init configuration is complete (BTC keys and withdrawal config) + pub fn is_provisioner_init_complete(&self) -> bool { + self.enclave_btc_keypair.get().is_some() + && self.hashi_btc_master_pubkey.get().is_some() + && self.withdrawal_config.get().is_some() + } + + /// Check if any provisioner_init configuration has been set + pub fn is_provisioner_init_partially_complete(&self) -> bool { + self.enclave_btc_keypair.get().is_some() + || self.hashi_btc_master_pubkey.get().is_some() + || self.withdrawal_config.get().is_some() + } +} + +impl EnclaveState { + pub fn init(&self, incoming_state: ProvisionerInitState) -> GuardianResult<()> { + let rate_limiter = incoming_state.build_rate_limiter()?; + let (committee, _, _, _) = incoming_state.into_parts(); + + self.set_committee(committee)?; + self.set_rate_limiter(rate_limiter)?; + Ok(()) + } + + // ======================================================================== + // Initialization Status + // ======================================================================== + + fn status_check_inner(&self) -> (bool, bool) { + let committee_init = self + .committee + .read() + .expect("rwlock read should not fail") + .is_some(); + + let limiter_init = self.rate_limiter.get().is_some(); + + (committee_init, limiter_init) + } + + /// Check if state init is complete + pub fn is_provisioner_init_complete(&self) -> bool { + let (committee_init, limiter_init) = self.status_check_inner(); + committee_init && limiter_init + } + + /// Check if any state has been set + pub fn is_provisioner_init_partially_complete(&self) -> bool { + let (committee_init, limiter_init) = self.status_check_inner(); + committee_init || limiter_init + } + + // ======================================================================== + // Committee Management + // ======================================================================== + + /// Get the current committee. + pub fn get_committee(&self) -> GuardianResult> { + let guard = self + .committee + .read() + .expect("rwlock should never throw an error"); + guard + .as_ref() + .cloned() + .ok_or_else(|| InvalidInputs("committee not initialized".into())) + } + + /// Set committee. Called only from init(ProvisionerInitState) + fn set_committee(&self, committee: HashiCommittee) -> GuardianResult<()> { + info!("Setting committee for epoch {}.", committee.epoch()); + + let mut guard = self + .committee + .write() + .expect("rwlock should never throw an error"); + if guard.is_some() { + return Err(InvalidInputs("committee already initialized".into())); + } + *guard = Some(Arc::new(committee)); + Ok(()) + } + + // ======================================================================== + // Rate Limiter Management + // ======================================================================== + + fn set_rate_limiter(&self, limiter: RateLimiter) -> GuardianResult<()> { + info!("Setting rate limiter."); + + self.rate_limiter + .set(Arc::new(tokio::sync::Mutex::new(limiter))) + .map_err(|_| InvalidInputs("rate_limiter already initialized".into())) + } + + /// Acquire exclusive access to the limiter, consume tokens, and return a guard. + /// The guard holds the mutex lock — no other withdrawal can start until it is + /// committed or dropped (which reverts). + /// Timeout for acquiring the limiter lock. If a withdrawal is in progress and + /// takes longer than this, we bail rather than queue up requests indefinitely. + const LIMITER_LOCK_TIMEOUT: Duration = Duration::from_secs(10); + + pub async fn consume_from_limiter( + &self, + seq: u64, + timestamp: u64, + amount_sats: u64, + ) -> GuardianResult { + let rate_limiter = self + .rate_limiter + .get() + .ok_or_else(|| InvalidInputs("rate_limiter not initialized".into()))?; + let mut guard = tokio::time::timeout( + Self::LIMITER_LOCK_TIMEOUT, + rate_limiter.clone().lock_owned(), + ) + .await + .map_err(|_| InvalidInputs("timed out waiting for rate limiter lock".into()))?; + guard.consume(seq, timestamp, amount_sats)?; + Ok(LimiterGuard::new(guard)) + } +} + +impl Enclave { + // ======================================================================== + // Construction & Initialization Status + // ======================================================================== + + pub fn new(signing_keys: GuardianSignKeyPair, encryption_keys: GuardianEncKeyPair) -> Self { + Enclave { + config: EnclaveConfig::new(signing_keys, encryption_keys), + state: EnclaveState { + committee: RwLock::new(None), + rate_limiter: OnceLock::new(), + }, + scratchpad: Scratchpad::default(), + } + } + + pub fn is_provisioner_init_complete(&self) -> bool { + self.config.is_provisioner_init_complete() + && self.state.is_provisioner_init_complete() + && self + .scratchpad + .provisioner_init_logging_complete + .get() + .is_some() + } + + pub fn is_provisioner_init_partially_complete(&self) -> bool { + self.config.is_provisioner_init_partially_complete() + || self.state.is_provisioner_init_partially_complete() + } + + pub fn is_operator_init_complete(&self) -> bool { + self.config.is_operator_init_complete() + && self.scratchpad.share_commitments.get().is_some() + && self + .scratchpad + .operator_init_logging_complete + .get() + .is_some() + } + + pub fn is_operator_init_partially_complete(&self) -> bool { + self.config.is_operator_init_partially_complete() + || self.scratchpad.share_commitments.get().is_some() + } + + pub fn is_fully_initialized(&self) -> bool { + self.is_provisioner_init_complete() && self.is_operator_init_complete() + } + + // ======================================================================== + // Ephemeral Keypairs (Encryption & Signing) + // ======================================================================== + + /// Get the enclave's encryption secret key + pub fn encryption_secret_key(&self) -> &EncSecKey { + self.config.eph_keys.encryption_keys.secret_key() + } + + /// Get the enclave's encryption public key + pub fn encryption_public_key(&self) -> &EncPubKey { + self.config.eph_keys.encryption_keys.public_key() + } + + /// Get the enclave's verification key + pub fn signing_pubkey(&self) -> GuardianPubKey { + self.config.eph_keys.signing_keys.verification_key() + } + + pub fn sign(&self, data: T) -> GuardianSigned { + let kp = &self.config.eph_keys.signing_keys; + let timestamp = now_timestamp_ms(); + GuardianSigned::new(data, kp, timestamp) + } + + // ======================================================================== + // Enclave Info + // ======================================================================== + + pub fn info(&self) -> GuardianInfo { + GuardianInfo { + share_commitments: self.share_commitments().ok().cloned(), + bucket_info: self + .config + .s3_logger() + .ok() + .map(|l| l.bucket_info().clone()), + encryption_pubkey: self.encryption_public_key().to_bytes().to_vec(), + // TODO: Change it + server_version: "v1".to_string(), + } + } + + // ======================================================================== + // S3 Logging + // ======================================================================== + + /// A unique session ID for the current enclave session. + pub fn s3_session_id(&self) -> String { + session_id_from_signing_pubkey(&self.signing_pubkey()) + } + + async fn write_log(&self, message: LogMessage) -> GuardianResult<()> { + let log = LogRecord::new( + self.s3_session_id(), + message, + &self.config.eph_keys.signing_keys, + ); + + self.config.s3_logger()?.write_log_record(log).await + } + + pub async fn log_init(&self, msg: InitLogMessage) -> GuardianResult<()> { + self.write_log(LogMessage::Init(Box::new(msg))).await + } + + pub async fn log_withdraw(&self, msg: WithdrawalLogMessage) -> GuardianResult<()> { + self.write_log(LogMessage::Withdrawal(Box::new(msg))).await + } + + pub async fn log_heartbeat(&self, seq: u64) -> GuardianResult<()> { + self.write_log(LogMessage::Heartbeat { seq }).await + } + + // ======================================================================== + // Scratchpad (Initialization-only data) + // ======================================================================== + + pub fn decrypted_shares(&self) -> &tokio::sync::Mutex> { + &self.scratchpad.shares + } + + pub fn share_commitments(&self) -> GuardianResult<&ShareCommitments> { + self.scratchpad + .share_commitments + .get() + .ok_or(InvalidInputs("Share commitments not set".into())) + } + + pub fn set_share_commitments(&self, commitments: ShareCommitments) -> GuardianResult<()> { + self.scratchpad + .share_commitments + .set(commitments) + .map_err(|_| InvalidInputs("Share commitments already set".into())) + } + + pub fn state_hash(&self) -> Option<&[u8; 32]> { + self.scratchpad.state_hash.get() + } + + pub fn set_state_hash(&self, hash: [u8; 32]) -> GuardianResult<()> { + self.scratchpad + .state_hash + .set(hash) + .map_err(|_| InvalidInputs("State hash already set".into())) + } +} diff --git a/crates/hashi-guardian/src/getters.rs b/crates/hashi-guardian/src/getters.rs index bb29c38a6..65320a6ad 100644 --- a/crates/hashi-guardian/src/getters.rs +++ b/crates/hashi-guardian/src/getters.rs @@ -9,7 +9,7 @@ use tracing::info; // Only needed in enclave builds (for NSM hardware interaction). // The `non-enclave-dev` feature and `cfg(test)` both route to the stub below. #[cfg(not(any(test, feature = "non-enclave-dev")))] -use crate::GuardianError; +use hashi_types::guardian::GuardianError; #[cfg(not(any(test, feature = "non-enclave-dev")))] use nsm_api::api::Request as NsmRequest; #[cfg(not(any(test, feature = "non-enclave-dev")))] diff --git a/crates/hashi-guardian/src/lib.rs b/crates/hashi-guardian/src/lib.rs index b8bc9af95..11fc4d1b9 100644 --- a/crates/hashi-guardian/src/lib.rs +++ b/crates/hashi-guardian/src/lib.rs @@ -8,4 +8,28 @@ pub const HEARTBEAT_INTERVAL: Duration = Duration::from_mins(1); pub const HEARTBEAT_RETRY_INTERVAL: Duration = Duration::from_secs(10); pub const MAX_HEARTBEAT_FAILURES_INTERVAL: Duration = Duration::from_mins(5); +pub mod enclave; +pub mod getters; +pub mod heartbeat; +pub mod init; +pub mod rpc; pub mod s3_logger; // used by the monitor +pub mod setup; +pub mod withdraw; + +#[cfg(any(test, feature = "test-utils"))] +pub mod test_utils; + +pub use enclave::Enclave; +pub use s3_logger::S3Logger; + +#[cfg(any(test, feature = "test-utils"))] +pub use test_utils::create_fully_initialized_enclave; +#[cfg(any(test, feature = "test-utils"))] +pub use test_utils::create_operator_initialized_enclave; +#[cfg(any(test, feature = "test-utils"))] +pub use test_utils::mock_logger; +#[cfg(any(test, feature = "test-utils"))] +pub use test_utils::FullyInitializedArgs; +#[cfg(any(test, feature = "test-utils"))] +pub use test_utils::OperatorInitTestArgs; diff --git a/crates/hashi-guardian/src/main.rs b/crates/hashi-guardian/src/main.rs index 24a88b61b..c130f4319 100644 --- a/crates/hashi-guardian/src/main.rs +++ b/crates/hashi-guardian/src/main.rs @@ -2,102 +2,20 @@ // SPDX-License-Identifier: Apache-2.0 use anyhow::Result; -use bitcoin::secp256k1::Keypair; -use bitcoin::Network; -use bitcoin::Txid; +use hashi_guardian::heartbeat::HeartbeatWriter; +use hashi_guardian::rpc::GuardianGrpc; +use hashi_guardian::Enclave; use hashi_guardian::HEARTBEAT_INTERVAL; use hashi_guardian::HEARTBEAT_RETRY_INTERVAL; use hashi_guardian::MAX_HEARTBEAT_FAILURES_INTERVAL; -use hashi_types::guardian::bitcoin_utils::sign_btc_tx; -use hashi_types::guardian::bitcoin_utils::TxUTXOs; -use hashi_types::guardian::crypto::Share; -use hashi_types::guardian::GuardianError::InvalidInputs; -use hashi_types::guardian::*; -use hpke::Serializable; -use serde::Serialize; +use hashi_types::guardian::GuardianEncKeyPair; +use hashi_types::guardian::GuardianSignKeyPair; +use hashi_types::proto::guardian_service_server::GuardianServiceServer; use std::sync::Arc; -use std::sync::OnceLock; -use std::sync::RwLock; -use std::time::Duration; use tonic::transport::Server; use tonic_health::server::health_reporter; use tracing::info; -mod getters; -mod heartbeat; -mod init; -mod rpc; -mod s3_logger; -mod setup; -mod withdraw; - -use crate::heartbeat::HeartbeatWriter; -use crate::rpc::GuardianGrpc; -use crate::s3_logger::S3Logger; -use crate::withdraw::LimiterGuard; -use hashi_types::committee::Committee as HashiCommittee; -use hashi_types::proto::guardian_service_server::GuardianServiceServer; - -/// Enclave's config & state -pub struct Enclave { - /// Immutable config (set once during init) - pub config: EnclaveConfig, - /// Mutable state - pub state: EnclaveState, - /// Initialization scratchpad - pub scratchpad: Scratchpad, -} - -/// Configuration set during initialization (immutable after set) -pub struct EnclaveConfig { - /// Ephemeral keypair (set on boot) - eph_keys: EphemeralKeyPairs, - /// S3 client & config (set in operator_init) - s3_logger: OnceLock, - /// Enclave BTC private key (set in provisioner_init) - enclave_btc_keypair: OnceLock, - /// BTC network: mainnet, testnet, regtest (set in operator_init) - btc_network: OnceLock, - /// Hashi BTC public key used to derive child keys (set in provisioner_init) - hashi_btc_master_pubkey: OnceLock, - /// Withdraw related config's (set in provisioner_init) - withdrawal_config: OnceLock, -} - -/// Mutable state that changes during operation. -/// Note: State is initialized during provisioner_init. -pub struct EnclaveState { - /// Current Hashi committee. - committee: RwLock>>, - /// Rate limiter. Set once during provisioner_init. - /// Uses `Arc` so the guard can be held across `.await`. - rate_limiter: OnceLock>>, -} - -/// Scratchpad used only during initialization. -/// Note: We don't clear it post-init because it does not have a lot of data. -#[derive(Default)] -pub struct Scratchpad { - /// The received shares - /// TODO: Investigate if it can be moved to std::sync::Mutex - pub shares: tokio::sync::Mutex>, - /// The share commitments - pub share_commitments: OnceLock, - /// Hash of the state in ProvisionerInitRequest - pub state_hash: OnceLock<[u8; 32]>, - /// Set once operator_init has successfully written all logs to S3. - /// This prevents heartbeats from being emitted before operator_init logs. - pub operator_init_logging_complete: OnceLock<()>, - /// Set once the provisioner init flow has successfully logged EnclaveFullyInitialized. - /// This prevents withdrawals from starting before provisioner_init logs. - pub provisioner_init_logging_complete: OnceLock<()>, -} - -pub struct EphemeralKeyPairs { - pub signing_keys: GuardianSignKeyPair, - pub encryption_keys: GuardianEncKeyPair, -} - /// Enclave initialization. /// SETUP_MODE=true: only get_attestation, operator_init and setup_new_key are enabled. /// SETUP_MODE=false: all endpoints except setup_new_key are enabled. @@ -155,501 +73,3 @@ async fn main() -> Result<()> { } } } - -impl EnclaveConfig { - pub fn new(signing_keys: GuardianSignKeyPair, encryption_keys: GuardianEncKeyPair) -> Self { - EnclaveConfig { - eph_keys: EphemeralKeyPairs { - signing_keys, - encryption_keys, - }, - s3_logger: OnceLock::new(), - enclave_btc_keypair: OnceLock::new(), - btc_network: OnceLock::new(), - hashi_btc_master_pubkey: OnceLock::new(), - withdrawal_config: OnceLock::new(), - } - } - - // ======================================================================== - // Bitcoin Configuration - // ======================================================================== - - pub fn bitcoin_network(&self) -> GuardianResult { - self.btc_network - .get() - .copied() - .ok_or(InvalidInputs("Network is uninitialized".into())) - } - - pub fn set_bitcoin_network(&self, network: Network) -> GuardianResult<()> { - self.btc_network - .set(network) - .map_err(|_| InvalidInputs("Network is already initialized".into())) - } - - pub fn set_btc_keypair(&self, keypair: Keypair) -> GuardianResult<()> { - self.enclave_btc_keypair - .set(keypair) - .map_err(|_| InvalidInputs("Bitcoin key already set".into())) - } - - pub fn set_hashi_btc_pk(&self, pk: BitcoinPubkey) -> GuardianResult<()> { - self.hashi_btc_master_pubkey - .set(pk) - .map_err(|_| InvalidInputs("Hashi BTC key is already set".into())) - } - - /// Sign a BTC tx. Returns an Err if enclave btc keypair or hashi btc pk is not set. - pub fn btc_sign(&self, tx_utxos: &TxUTXOs) -> GuardianResult<(Txid, Vec)> { - let enclave_keypair = self - .enclave_btc_keypair - .get() - .ok_or(InvalidInputs("Bitcoin key is not initialized".into()))?; - let hashi_btc_pk = self - .hashi_btc_master_pubkey - .get() - .ok_or(InvalidInputs("Hashi BTC public key not set".into()))?; - - let enclave_btc_pk = enclave_keypair.x_only_public_key().0; - let (messages, txid) = tx_utxos.signing_messages_and_txid(&enclave_btc_pk, hashi_btc_pk); - Ok((txid, sign_btc_tx(&messages, enclave_keypair))) - } - - // ======================================================================== - // Withdrawal Configuration - // ======================================================================== - - pub fn withdrawal_config(&self) -> GuardianResult<&WithdrawalConfig> { - self.withdrawal_config - .get() - .ok_or(InvalidInputs("WithdrawalConfig is not initialized".into())) - } - - pub fn set_withdrawal_config(&self, config: WithdrawalConfig) -> GuardianResult<()> { - self.withdrawal_config - .set(config) - .map_err(|_| InvalidInputs("WithdrawalConfig already set".into())) - } - - pub fn committee_threshold(&self) -> GuardianResult { - Ok(self.withdrawal_config()?.committee_threshold) - } - - // ======================================================================== - // S3 Logger - // ======================================================================== - - pub fn s3_logger(&self) -> GuardianResult<&S3Logger> { - self.s3_logger - .get() - .ok_or(InvalidInputs("S3 logger is not initialized".into())) - } - - pub fn set_s3_logger(&self, logger: S3Logger) -> GuardianResult<()> { - self.s3_logger - .set(logger) - .map_err(|_| InvalidInputs("S3 logger already set".into())) - } - - // ======================================================================== - // Initialization Status - // ======================================================================== - - /// Check if operator_init configuration is complete (S3 logger and network) - pub fn is_operator_init_complete(&self) -> bool { - self.s3_logger.get().is_some() && self.btc_network.get().is_some() - } - - /// Check if any operator_init configuration has been set - pub fn is_operator_init_partially_complete(&self) -> bool { - self.s3_logger.get().is_some() || self.btc_network.get().is_some() - } - - /// Check if provisioner_init configuration is complete (BTC keys and withdrawal config) - pub fn is_provisioner_init_complete(&self) -> bool { - self.enclave_btc_keypair.get().is_some() - && self.hashi_btc_master_pubkey.get().is_some() - && self.withdrawal_config.get().is_some() - } - - /// Check if any provisioner_init configuration has been set - pub fn is_provisioner_init_partially_complete(&self) -> bool { - self.enclave_btc_keypair.get().is_some() - || self.hashi_btc_master_pubkey.get().is_some() - || self.withdrawal_config.get().is_some() - } -} - -impl EnclaveState { - pub fn init(&self, incoming_state: ProvisionerInitState) -> GuardianResult<()> { - let rate_limiter = incoming_state.build_rate_limiter()?; - let (committee, _, _, _) = incoming_state.into_parts(); - - self.set_committee(committee)?; - self.set_rate_limiter(rate_limiter)?; - Ok(()) - } - - // ======================================================================== - // Initialization Status - // ======================================================================== - - fn status_check_inner(&self) -> (bool, bool) { - let committee_init = self - .committee - .read() - .expect("rwlock read should not fail") - .is_some(); - - let limiter_init = self.rate_limiter.get().is_some(); - - (committee_init, limiter_init) - } - - /// Check if state init is complete - pub fn is_provisioner_init_complete(&self) -> bool { - let (committee_init, limiter_init) = self.status_check_inner(); - committee_init && limiter_init - } - - /// Check if any state has been set - pub fn is_provisioner_init_partially_complete(&self) -> bool { - let (committee_init, limiter_init) = self.status_check_inner(); - committee_init || limiter_init - } - - // ======================================================================== - // Committee Management - // ======================================================================== - - /// Get the current committee. - pub fn get_committee(&self) -> GuardianResult> { - let guard = self - .committee - .read() - .expect("rwlock should never throw an error"); - guard - .as_ref() - .cloned() - .ok_or_else(|| InvalidInputs("committee not initialized".into())) - } - - /// Set committee. Called only from init(ProvisionerInitState) - fn set_committee(&self, committee: HashiCommittee) -> GuardianResult<()> { - info!("Setting committee for epoch {}.", committee.epoch()); - - let mut guard = self - .committee - .write() - .expect("rwlock should never throw an error"); - if guard.is_some() { - return Err(InvalidInputs("committee already initialized".into())); - } - *guard = Some(Arc::new(committee)); - Ok(()) - } - - // ======================================================================== - // Rate Limiter Management - // ======================================================================== - - fn set_rate_limiter(&self, limiter: RateLimiter) -> GuardianResult<()> { - info!("Setting rate limiter."); - - self.rate_limiter - .set(Arc::new(tokio::sync::Mutex::new(limiter))) - .map_err(|_| InvalidInputs("rate_limiter already initialized".into())) - } - - /// Acquire exclusive access to the limiter, consume tokens, and return a guard. - /// The guard holds the mutex lock — no other withdrawal can start until it is - /// committed or dropped (which reverts). - /// Timeout for acquiring the limiter lock. If a withdrawal is in progress and - /// takes longer than this, we bail rather than queue up requests indefinitely. - const LIMITER_LOCK_TIMEOUT: Duration = Duration::from_secs(10); - - pub async fn consume_from_limiter( - &self, - seq: u64, - timestamp: u64, - amount_sats: u64, - ) -> GuardianResult { - let rate_limiter = self - .rate_limiter - .get() - .ok_or_else(|| InvalidInputs("rate_limiter not initialized".into()))?; - let mut guard = tokio::time::timeout( - Self::LIMITER_LOCK_TIMEOUT, - rate_limiter.clone().lock_owned(), - ) - .await - .map_err(|_| InvalidInputs("timed out waiting for rate limiter lock".into()))?; - guard.consume(seq, timestamp, amount_sats)?; - Ok(LimiterGuard::new(guard)) - } -} - -impl Enclave { - // ======================================================================== - // Construction & Initialization Status - // ======================================================================== - - pub fn new(signing_keys: GuardianSignKeyPair, encryption_keys: GuardianEncKeyPair) -> Self { - Enclave { - config: EnclaveConfig::new(signing_keys, encryption_keys), - state: EnclaveState { - committee: RwLock::new(None), - rate_limiter: OnceLock::new(), - }, - scratchpad: Scratchpad::default(), - } - } - - pub fn is_provisioner_init_complete(&self) -> bool { - self.config.is_provisioner_init_complete() - && self.state.is_provisioner_init_complete() - && self - .scratchpad - .provisioner_init_logging_complete - .get() - .is_some() - } - - pub fn is_provisioner_init_partially_complete(&self) -> bool { - self.config.is_provisioner_init_partially_complete() - || self.state.is_provisioner_init_partially_complete() - } - - pub fn is_operator_init_complete(&self) -> bool { - self.config.is_operator_init_complete() - && self.scratchpad.share_commitments.get().is_some() - && self - .scratchpad - .operator_init_logging_complete - .get() - .is_some() - } - - pub fn is_operator_init_partially_complete(&self) -> bool { - self.config.is_operator_init_partially_complete() - || self.scratchpad.share_commitments.get().is_some() - } - - pub fn is_fully_initialized(&self) -> bool { - self.is_provisioner_init_complete() && self.is_operator_init_complete() - } - - // ======================================================================== - // Ephemeral Keypairs (Encryption & Signing) - // ======================================================================== - - /// Get the enclave's encryption secret key - pub fn encryption_secret_key(&self) -> &EncSecKey { - self.config.eph_keys.encryption_keys.secret_key() - } - - /// Get the enclave's encryption public key - pub fn encryption_public_key(&self) -> &EncPubKey { - self.config.eph_keys.encryption_keys.public_key() - } - - /// Get the enclave's verification key - pub fn signing_pubkey(&self) -> GuardianPubKey { - self.config.eph_keys.signing_keys.verification_key() - } - - pub fn sign(&self, data: T) -> GuardianSigned { - let kp = &self.config.eph_keys.signing_keys; - let timestamp = now_timestamp_ms(); - GuardianSigned::new(data, kp, timestamp) - } - - // ======================================================================== - // Enclave Info - // ======================================================================== - - pub fn info(&self) -> GuardianInfo { - GuardianInfo { - share_commitments: self.share_commitments().ok().cloned(), - bucket_info: self - .config - .s3_logger() - .ok() - .map(|l| l.bucket_info().clone()), - encryption_pubkey: self.encryption_public_key().to_bytes().to_vec(), - // TODO: Change it - server_version: "v1".to_string(), - } - } - - // ======================================================================== - // S3 Logging - // ======================================================================== - - /// A unique session ID for the current enclave session. - pub fn s3_session_id(&self) -> String { - session_id_from_signing_pubkey(&self.signing_pubkey()) - } - - async fn write_log(&self, message: LogMessage) -> GuardianResult<()> { - let log = LogRecord::new( - self.s3_session_id(), - message, - &self.config.eph_keys.signing_keys, - ); - - self.config.s3_logger()?.write_log_record(log).await - } - - pub async fn log_init(&self, msg: InitLogMessage) -> GuardianResult<()> { - self.write_log(LogMessage::Init(Box::new(msg))).await - } - - pub async fn log_withdraw(&self, msg: WithdrawalLogMessage) -> GuardianResult<()> { - self.write_log(LogMessage::Withdrawal(Box::new(msg))).await - } - - pub async fn log_heartbeat(&self, seq: u64) -> GuardianResult<()> { - self.write_log(LogMessage::Heartbeat { seq }).await - } - - // ======================================================================== - // Scratchpad (Initialization-only data) - // ======================================================================== - - pub fn decrypted_shares(&self) -> &tokio::sync::Mutex> { - &self.scratchpad.shares - } - - pub fn share_commitments(&self) -> GuardianResult<&ShareCommitments> { - self.scratchpad - .share_commitments - .get() - .ok_or(InvalidInputs("Share commitments not set".into())) - } - - pub fn set_share_commitments(&self, commitments: ShareCommitments) -> GuardianResult<()> { - self.scratchpad - .share_commitments - .set(commitments) - .map_err(|_| InvalidInputs("Share commitments already set".into())) - } - - pub fn state_hash(&self) -> Option<&[u8; 32]> { - self.scratchpad.state_hash.get() - } - - pub fn set_state_hash(&self, hash: [u8; 32]) -> GuardianResult<()> { - self.scratchpad - .state_hash - .set(hash) - .map_err(|_| InvalidInputs("State hash already set".into())) - } -} - -// --------------------------------- -// Tests and related utilities -// --------------------------------- - -// Mock S3 logger for use in APIs calls post operator_init, e.g., provisioner_init, withdrawals. -#[cfg(test)] -pub fn mock_logger() -> S3Logger { - use aws_sdk_s3::operation::put_object::PutObjectOutput; - use aws_sdk_s3::Client; - use aws_smithy_mocks::mock; - use aws_smithy_mocks::mock_client; - use aws_smithy_mocks::RuleMode; - use hashi_types::guardian::S3Config; - - // For unit tests we only need PutObject to succeed, because `sign_and_log()` calls `S3Logger::write()`. - // The `then_output` helper creates a "simple" rule that repeats indefinitely. - let put_ok = mock!(Client::put_object).then_output(|| PutObjectOutput::builder().build()); - - let client = mock_client!(aws_sdk_s3, RuleMode::MatchAny, &[&put_ok]); - - let config = S3Config::mock_for_testing(); - - S3Logger::from_client_for_tests(config, client) -} - -#[cfg(test)] -pub struct OperatorInitTestArgs { - pub network: Network, - pub commitments: ShareCommitments, - pub s3_logger: S3Logger, -} - -#[cfg(test)] -impl Default for OperatorInitTestArgs { - fn default() -> Self { - let commitments = (1..=NUM_OF_SHARES) - .map(|id| ShareCommitment { - id: std::num::NonZeroU16::new(id as u16).unwrap(), - digest: vec![], - }) - .collect(); - - Self { - network: Network::Regtest, - commitments: ShareCommitments::new(commitments).unwrap(), - s3_logger: mock_logger(), - } - } -} - -#[cfg(test)] -impl OperatorInitTestArgs { - pub fn with_network(mut self, network: Network) -> Self { - self.network = network; - self - } - - pub fn with_commitments(mut self, commitments: ShareCommitments) -> Self { - self.commitments = commitments; - self - } - - pub fn with_s3_logger(mut self, s3_logger: S3Logger) -> Self { - self.s3_logger = s3_logger; - self - } -} - -#[cfg(test)] -impl Enclave { - pub fn create_with_random_keys() -> Arc { - let signing_keys = GuardianSignKeyPair::new(rand::thread_rng()); - let encryption_keys = GuardianEncKeyPair::random(&mut rand::thread_rng()); - Arc::new(Enclave::new(signing_keys, encryption_keys)) - } - - // Create an enclave post operator_init() but pre provisioner_init(). - pub async fn create_operator_initialized() -> Arc { - Self::create_operator_initialized_with(OperatorInitTestArgs::default()).await - } - - pub async fn create_operator_initialized_with(args: OperatorInitTestArgs) -> Arc { - let enclave = Self::create_with_random_keys(); - - // Initialize S3 logger - enclave.config.set_s3_logger(args.s3_logger).unwrap(); - - // Set bitcoin network - enclave.config.set_bitcoin_network(args.network).unwrap(); - - // Set share commitments - enclave.set_share_commitments(args.commitments).unwrap(); - - // In tests, treat "operator initialized" as including the operator-init identity logs. - enclave - .scratchpad - .operator_init_logging_complete - .set(()) - .expect("operator_init_logging_complete should only be set once"); - - assert!(enclave.is_operator_init_complete() && !enclave.is_provisioner_init_complete()); - - enclave - } -} diff --git a/crates/hashi-guardian/src/test_utils.rs b/crates/hashi-guardian/src/test_utils.rs new file mode 100644 index 000000000..82a2e6ed2 --- /dev/null +++ b/crates/hashi-guardian/src/test_utils.rs @@ -0,0 +1,161 @@ +// Copyright (c) Mysten Labs, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//! Helpers for constructing enclaves at various init stages. + +use crate::enclave::Enclave; +use crate::s3_logger::S3Logger; +use bitcoin::secp256k1::Keypair; +use bitcoin::secp256k1::Secp256k1; +use bitcoin::secp256k1::SecretKey; +use bitcoin::Network; +use hashi_types::guardian::*; +use rand::RngCore; +use std::sync::Arc; + +/// Mock S3 logger that returns success for every PutObject call. +pub fn mock_logger() -> S3Logger { + use aws_sdk_s3::operation::put_object::PutObjectOutput; + use aws_sdk_s3::Client; + use aws_smithy_mocks::mock; + use aws_smithy_mocks::mock_client; + use aws_smithy_mocks::RuleMode; + use hashi_types::guardian::S3Config; + + let put_ok = mock!(Client::put_object).then_output(|| PutObjectOutput::builder().build()); + let client = mock_client!(aws_sdk_s3, RuleMode::MatchAny, &[&put_ok]); + S3Logger::from_client_for_tests(S3Config::mock_for_testing(), client) +} + +pub struct OperatorInitTestArgs { + pub network: Network, + pub commitments: ShareCommitments, + pub s3_logger: S3Logger, +} + +impl Default for OperatorInitTestArgs { + fn default() -> Self { + let commitments = (1..=NUM_OF_SHARES) + .map(|id| ShareCommitment { + id: std::num::NonZeroU16::new(id as u16).unwrap(), + digest: vec![], + }) + .collect(); + + Self { + network: Network::Regtest, + commitments: ShareCommitments::new(commitments).unwrap(), + s3_logger: mock_logger(), + } + } +} + +impl OperatorInitTestArgs { + pub fn with_network(mut self, network: Network) -> Self { + self.network = network; + self + } + + pub fn with_commitments(mut self, commitments: ShareCommitments) -> Self { + self.commitments = commitments; + self + } + + pub fn with_s3_logger(mut self, s3_logger: S3Logger) -> Self { + self.s3_logger = s3_logger; + self + } +} + +impl Enclave { + pub fn create_with_random_keys() -> Arc { + let signing_keys = GuardianSignKeyPair::new(rand::thread_rng()); + let encryption_keys = GuardianEncKeyPair::random(&mut rand::thread_rng()); + Arc::new(Enclave::new(signing_keys, encryption_keys)) + } + + /// Create an enclave post operator_init() but pre provisioner_init(). + pub async fn create_operator_initialized() -> Arc { + Self::create_operator_initialized_with(OperatorInitTestArgs::default()).await + } + + pub async fn create_operator_initialized_with(args: OperatorInitTestArgs) -> Arc { + let enclave = Self::create_with_random_keys(); + enclave.config.set_s3_logger(args.s3_logger).unwrap(); + enclave.config.set_bitcoin_network(args.network).unwrap(); + enclave.set_share_commitments(args.commitments).unwrap(); + enclave + .scratchpad + .operator_init_logging_complete + .set(()) + .expect("operator_init_logging_complete should only be set once"); + + assert!(enclave.is_operator_init_complete() && !enclave.is_provisioner_init_complete()); + enclave + } +} + +pub async fn create_operator_initialized_enclave(args: OperatorInitTestArgs) -> Arc { + Enclave::create_operator_initialized_with(args).await +} + +pub struct FullyInitializedArgs { + pub network: Network, + pub committee: HashiCommittee, + pub master_pubkey: BitcoinPubkey, + pub withdrawal_config: WithdrawalConfig, + pub limiter_state: LimiterState, +} + +/// Drive an enclave to fully-initialized state, skipping the real share +/// encryption round-trip. Uses a freshly-generated BTC keypair. +pub async fn create_fully_initialized_enclave(args: FullyInitializedArgs) -> Arc { + let FullyInitializedArgs { + network, + committee, + master_pubkey, + withdrawal_config, + limiter_state, + } = args; + + let enclave = + create_operator_initialized_enclave(OperatorInitTestArgs::default().with_network(network)) + .await; + + let secp = Secp256k1::new(); + let mut sk_bytes = [0u8; 32]; + rand::thread_rng().fill_bytes(&mut sk_bytes); + let enclave_btc_keypair = Keypair::from_secret_key( + &secp, + &SecretKey::from_slice(&sk_bytes).expect("random bytes form a valid secp256k1 key"), + ); + enclave + .config + .set_btc_keypair(enclave_btc_keypair) + .expect("fresh enclave should not have a btc keypair set"); + enclave + .config + .set_hashi_btc_pk(master_pubkey) + .expect("fresh enclave should not have a master pubkey set"); + enclave + .config + .set_withdrawal_config(withdrawal_config) + .expect("fresh enclave should not have a withdrawal config set"); + + let init_state = + ProvisionerInitState::new(committee, withdrawal_config, limiter_state, master_pubkey) + .expect("valid ProvisionerInitState"); + enclave + .state + .init(init_state) + .expect("provisioner state init should succeed on a fresh enclave"); + + enclave + .scratchpad + .provisioner_init_logging_complete + .set(()) + .expect("provisioner_init_logging_complete should only be set once"); + + assert!(enclave.is_fully_initialized()); + enclave +} diff --git a/crates/hashi/src/config.rs b/crates/hashi/src/config.rs index 9ce3bf84f..13a3926f7 100644 --- a/crates/hashi/src/config.rs +++ b/crates/hashi/src/config.rs @@ -134,6 +134,11 @@ pub struct Config { #[serde(skip_serializing_if = "Option::is_none")] pub screener_endpoint: Option, + /// URL of the `hashi-guardian` gRPC endpoint. When not set, the guardian + /// integration is bypassed. + #[serde(skip_serializing_if = "Option::is_none")] + pub guardian_endpoint: Option, + /// Maximum gRPC decoding message size in bytes. /// /// Defaults to 16 MiB if not specified. Tonic's built-in default is 4 MiB, @@ -353,6 +358,10 @@ impl Config { self.screener_endpoint.as_deref() } + pub fn guardian_endpoint(&self) -> Option<&str> { + self.guardian_endpoint.as_deref() + } + pub fn grpc_max_decoding_message_size(&self) -> usize { self.grpc_max_decoding_message_size .unwrap_or(16 * 1024 * 1024) diff --git a/crates/hashi/src/grpc/guardian_client.rs b/crates/hashi/src/grpc/guardian_client.rs new file mode 100644 index 000000000..6f55bd2dd --- /dev/null +++ b/crates/hashi/src/grpc/guardian_client.rs @@ -0,0 +1,40 @@ +// Copyright (c) Mysten Labs, Inc. +// SPDX-License-Identifier: Apache-2.0 + +use std::time::Duration; + +use hashi_types::proto::guardian_service_client::GuardianServiceClient; +use tonic::transport::Channel; +use tonic::transport::Endpoint; + +type BoxError = Box; + +/// Lazy gRPC channel to a `hashi-guardian`. +#[derive(Clone, Debug)] +pub struct GuardianClient { + endpoint: String, + channel: Channel, +} + +impl GuardianClient { + pub fn new(endpoint: &str) -> Result { + let channel = Endpoint::from_shared(endpoint.to_string()) + .map_err(Into::::into) + .map_err(tonic::Status::from_error)? + .connect_timeout(Duration::from_secs(5)) + .http2_keep_alive_interval(Duration::from_secs(5)) + .connect_lazy(); + Ok(Self { + endpoint: endpoint.to_string(), + channel, + }) + } + + pub fn endpoint(&self) -> &str { + &self.endpoint + } + + pub fn guardian_service_client(&self) -> GuardianServiceClient { + GuardianServiceClient::new(self.channel.clone()) + } +} diff --git a/crates/hashi/src/grpc/mod.rs b/crates/hashi/src/grpc/mod.rs index cd9d9e2de..370c5f2a1 100644 --- a/crates/hashi/src/grpc/mod.rs +++ b/crates/hashi/src/grpc/mod.rs @@ -15,6 +15,7 @@ pub use client::Client; pub use client::MPC_PROTOCOL_METADATA_KEY; pub mod bridge_service; +pub mod guardian_client; pub mod metrics_layer; pub mod screener_client; diff --git a/crates/hashi/src/lib.rs b/crates/hashi/src/lib.rs index 1452b16b9..454a5da60 100644 --- a/crates/hashi/src/lib.rs +++ b/crates/hashi/src/lib.rs @@ -49,6 +49,7 @@ pub struct Hashi { mpc_handle: OnceLock, btc_monitor: OnceLock, screener_client: OnceLock>, + guardian_client: OnceLock>, /// Reconfig completion signatures by epoch. reconfig_signatures: RwLock>>, } @@ -70,6 +71,7 @@ impl Hashi { mpc_handle: OnceLock::new(), btc_monitor: OnceLock::new(), screener_client: OnceLock::new(), + guardian_client: OnceLock::new(), reconfig_signatures: RwLock::new(HashMap::new()), })) } @@ -94,6 +96,7 @@ impl Hashi { mpc_handle: OnceLock::new(), btc_monitor: OnceLock::new(), screener_client: OnceLock::new(), + guardian_client: OnceLock::new(), reconfig_signatures: RwLock::new(HashMap::new()), })) } @@ -185,6 +188,10 @@ impl Hashi { self.screener_client.get().and_then(|opt| opt.as_ref()) } + pub fn guardian_client(&self) -> Option<&grpc::guardian_client::GuardianClient> { + self.guardian_client.get().and_then(|opt| opt.as_ref()) + } + async fn initialize_onchain_state(&self) -> anyhow::Result { let (onchain_state, service) = onchain::OnchainState::new( self.config.sui_rpc.as_deref().unwrap(), @@ -386,6 +393,30 @@ impl Hashi { .set(screener) .map_err(|_| anyhow!("Screener client already initialized"))?; + let guardian = if let Some(endpoint) = self.config.guardian_endpoint() { + match grpc::guardian_client::GuardianClient::new(endpoint) { + Ok(client) => { + tracing::info!("Guardian client configured for {}", client.endpoint()); + Some(client) + } + Err(e) => { + tracing::warn!( + "Failed to configure guardian client for {}: {}", + endpoint, + e + ); + None + } + } + } else { + tracing::info!("No guardian endpoint configured; guardian integration disabled"); + None + }; + + self.guardian_client + .set(guardian) + .map_err(|_| anyhow!("Guardian client already initialized"))?; + // Verify Sui RPC is on the expected chain before loading any state. self.verify_sui_chain_id().await?; From fc429c28a9f0f097b8c82b203ab7627c45c2d7dd Mon Sep 17 00:00:00 2001 From: Siddharth Sharma Date: Sun, 26 Apr 2026 19:20:39 -0700 Subject: [PATCH 2/8] [guardian] extract finalize_enclave helper into test_utils `GuardianHarness::finalize` and `create_fully_initialized_enclave` were running the same "drive an operator-init'd enclave to fully-initialized state" sequence (set BTC keys, build init state, state.init, mark logged) inline. Lift it into `hashi_guardian::test_utils::finalize_enclave` and call it from both. The harness's `finalize` shrinks to a single helper call plus the `is_fully_initialized` post-condition. No visibility or API changes to the `Enclave` types. --- crates/e2e-tests/src/guardian_harness.rs | 47 +++----------- crates/hashi-guardian/src/test_utils.rs | 82 ++++++++++++++---------- 2 files changed, 56 insertions(+), 73 deletions(-) diff --git a/crates/e2e-tests/src/guardian_harness.rs b/crates/e2e-tests/src/guardian_harness.rs index 656d4a07d..c160725e8 100644 --- a/crates/e2e-tests/src/guardian_harness.rs +++ b/crates/e2e-tests/src/guardian_harness.rs @@ -16,7 +16,6 @@ use hashi_guardian::rpc::GuardianGrpc; use hashi_types::committee::Committee as HashiCommittee; use hashi_types::guardian::BitcoinPubkey; use hashi_types::guardian::LimiterState; -use hashi_types::guardian::ProvisionerInitState; use hashi_types::guardian::WithdrawalConfig; use hashi_types::proto::guardian_service_server::GuardianServiceServer; use std::net::SocketAddr; @@ -87,44 +86,14 @@ impl GuardianHarness { withdrawal_config: WithdrawalConfig, limiter_state: LimiterState, ) -> Result<()> { - use bitcoin::secp256k1::Keypair; - use bitcoin::secp256k1::Secp256k1; - use bitcoin::secp256k1::SecretKey; - use rand::RngCore; - - let secp = Secp256k1::new(); - let mut sk_bytes = [0u8; 32]; - rand::thread_rng().fill_bytes(&mut sk_bytes); - let enclave_btc_keypair = Keypair::from_secret_key( - &secp, - &SecretKey::from_slice(&sk_bytes).expect("random bytes form a valid secp256k1 key"), - ); - self.enclave - .config - .set_btc_keypair(enclave_btc_keypair) - .context("set enclave btc keypair")?; - self.enclave - .config - .set_hashi_btc_pk(master_pubkey) - .context("set hashi btc master pubkey")?; - self.enclave - .config - .set_withdrawal_config(withdrawal_config) - .context("set withdrawal config")?; - - let init_state = - ProvisionerInitState::new(committee, withdrawal_config, limiter_state, master_pubkey) - .context("valid ProvisionerInitState")?; - self.enclave - .state - .init(init_state) - .context("init enclave state")?; - - self.enclave - .scratchpad - .provisioner_init_logging_complete - .set(()) - .map_err(|_| anyhow::anyhow!("provisioner_init already finalized"))?; + hashi_guardian::test_utils::finalize_enclave( + &self.enclave, + committee, + master_pubkey, + withdrawal_config, + limiter_state, + ) + .map_err(|e| anyhow::anyhow!("finalize guardian enclave: {e:?}"))?; anyhow::ensure!( self.enclave.is_fully_initialized(), diff --git a/crates/hashi-guardian/src/test_utils.rs b/crates/hashi-guardian/src/test_utils.rs index 82a2e6ed2..e95cc63c4 100644 --- a/crates/hashi-guardian/src/test_utils.rs +++ b/crates/hashi-guardian/src/test_utils.rs @@ -107,21 +107,19 @@ pub struct FullyInitializedArgs { pub limiter_state: LimiterState, } -/// Drive an enclave to fully-initialized state, skipping the real share -/// encryption round-trip. Uses a freshly-generated BTC keypair. -pub async fn create_fully_initialized_enclave(args: FullyInitializedArgs) -> Arc { - let FullyInitializedArgs { - network, - committee, - master_pubkey, - withdrawal_config, - limiter_state, - } = args; - - let enclave = - create_operator_initialized_enclave(OperatorInitTestArgs::default().with_network(network)) - .await; - +/// Drive an already operator-initialized enclave to fully-initialized state, +/// skipping the real share-encryption round-trip. Generates a fresh random +/// BTC keypair internally. +/// +/// Used by [`create_fully_initialized_enclave`] and by the `e2e-tests` +/// `GuardianHarness` when DKG output becomes available on chain. +pub fn finalize_enclave( + enclave: &Arc, + committee: HashiCommittee, + master_pubkey: BitcoinPubkey, + withdrawal_config: WithdrawalConfig, + limiter_state: LimiterState, +) -> GuardianResult<()> { let secp = Secp256k1::new(); let mut sk_bytes = [0u8; 32]; rand::thread_rng().fill_bytes(&mut sk_bytes); @@ -129,32 +127,48 @@ pub async fn create_fully_initialized_enclave(args: FullyInitializedArgs) -> Arc &secp, &SecretKey::from_slice(&sk_bytes).expect("random bytes form a valid secp256k1 key"), ); - enclave - .config - .set_btc_keypair(enclave_btc_keypair) - .expect("fresh enclave should not have a btc keypair set"); - enclave - .config - .set_hashi_btc_pk(master_pubkey) - .expect("fresh enclave should not have a master pubkey set"); - enclave - .config - .set_withdrawal_config(withdrawal_config) - .expect("fresh enclave should not have a withdrawal config set"); + enclave.config.set_btc_keypair(enclave_btc_keypair)?; + enclave.config.set_hashi_btc_pk(master_pubkey)?; + enclave.config.set_withdrawal_config(withdrawal_config)?; let init_state = - ProvisionerInitState::new(committee, withdrawal_config, limiter_state, master_pubkey) - .expect("valid ProvisionerInitState"); - enclave - .state - .init(init_state) - .expect("provisioner state init should succeed on a fresh enclave"); + ProvisionerInitState::new(committee, withdrawal_config, limiter_state, master_pubkey)?; + enclave.state.init(init_state)?; enclave .scratchpad .provisioner_init_logging_complete .set(()) - .expect("provisioner_init_logging_complete should only be set once"); + .map_err(|_| { + GuardianError::InvalidInputs("provisioner_init_logging_complete already set".into()) + })?; + Ok(()) +} + +/// Construct an operator-initialized enclave and drive it to fully-initialized +/// state in one shot. Convenience for unit tests that don't need the +/// intermediate "operator-init only" stage. +pub async fn create_fully_initialized_enclave(args: FullyInitializedArgs) -> Arc { + let FullyInitializedArgs { + network, + committee, + master_pubkey, + withdrawal_config, + limiter_state, + } = args; + + let enclave = + create_operator_initialized_enclave(OperatorInitTestArgs::default().with_network(network)) + .await; + + finalize_enclave( + &enclave, + committee, + master_pubkey, + withdrawal_config, + limiter_state, + ) + .expect("finalize_enclave should succeed on a fresh enclave"); assert!(enclave.is_fully_initialized()); enclave From 44fa42cef819fb7d1fedf73b83ef0f48aa4198ef Mon Sep 17 00:00:00 2001 From: Siddharth Sharma Date: Sun, 26 Apr 2026 20:23:24 -0700 Subject: [PATCH 3/8] [guardian] drop dead network getter on GuardianHarness The `network` field was set in `start()` but never read internally; the getter was marked `#[allow(dead_code)]`. The enclave already owns the network, and `pub` already opts out of dead-code lint, so the attribute on a `pub` API was a smell. --- crates/e2e-tests/src/guardian_harness.rs | 7 ------- 1 file changed, 7 deletions(-) diff --git a/crates/e2e-tests/src/guardian_harness.rs b/crates/e2e-tests/src/guardian_harness.rs index c160725e8..c0d6c0032 100644 --- a/crates/e2e-tests/src/guardian_harness.rs +++ b/crates/e2e-tests/src/guardian_harness.rs @@ -30,7 +30,6 @@ use tonic::transport::Server; pub struct GuardianHarness { enclave: Arc, endpoint: String, - network: Network, shutdown_tx: Option>, server_handle: Option>, } @@ -71,7 +70,6 @@ impl GuardianHarness { Ok(Self { enclave, endpoint, - network, shutdown_tx: Some(shutdown_tx), server_handle: Some(server_handle), }) @@ -109,11 +107,6 @@ impl GuardianHarness { pub fn enclave(&self) -> &Arc { &self.enclave } - - #[allow(dead_code)] - pub fn network(&self) -> Network { - self.network - } } impl Drop for GuardianHarness { From 265b39e406c088ae777c66952fc16eb9ed2afaeb Mon Sep 17 00:00:00 2001 From: Siddharth Sharma Date: Sun, 26 Apr 2026 20:26:03 -0700 Subject: [PATCH 4/8] [guardian] parameterize bitcoin withdrawal e2e tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extract `run_bitcoin_withdrawal_e2e(with_guardian: bool)` as the shared body for the deposit → withdrawal → confirm flow. Both `test_bitcoin_withdrawal_e2e_flow` and `test_bitcoin_withdrawal_with_guardian_e2e_flow` now reduce to one-line wrappers that pass `false` / `true`. Forces the two flows to stay in lockstep — any future change to the deposit/withdrawal path is automatically exercised under both configurations. The guardian variant gains the post-deposit balance assertion and the withdrawal_request_id / withdrawal_txid log lines that were previously only in the no-guardian variant; behavior is otherwise unchanged. `setup_test_networks` is no longer called from the bitcoin-withdrawal test (its setup is now inlined into the parameterized helper) but remains in use by the other 3 e2e tests. --- crates/e2e-tests/src/e2e_flow.rs | 161 +++++++++++-------------------- 1 file changed, 58 insertions(+), 103 deletions(-) diff --git a/crates/e2e-tests/src/e2e_flow.rs b/crates/e2e-tests/src/e2e_flow.rs index 7ebe650be..52150f10c 100644 --- a/crates/e2e-tests/src/e2e_flow.rs +++ b/crates/e2e-tests/src/e2e_flow.rs @@ -255,10 +255,65 @@ mod tests { #[tokio::test] async fn test_bitcoin_withdrawal_e2e_flow() -> Result<()> { + run_bitcoin_withdrawal_e2e(false).await + } + + /// Same flow as `test_bitcoin_withdrawal_e2e_flow`, but with an + /// in-process guardian in the loop. + #[tokio::test] + async fn test_bitcoin_withdrawal_with_guardian_e2e_flow() -> Result<()> { + run_bitcoin_withdrawal_e2e(true).await + } + + /// Shared body for the bitcoin-withdrawal e2e tests, parameterized on + /// whether to spin up an in-process guardian alongside the bridge. + /// Keeping a single body forces the two flows to stay in lockstep. + async fn run_bitcoin_withdrawal_e2e(with_guardian: bool) -> Result<()> { init_test_logging(); - info!("=== Starting Bitcoin Withdrawal E2E Test ==="); + let label = if with_guardian { + " (with guardian)" + } else { + "" + }; + info!("=== Starting Bitcoin Withdrawal E2E Test{label} ==="); - let mut networks = setup_test_networks().await?; + let mut builder = TestNetworksBuilder::new().with_nodes(4); + if with_guardian { + builder = builder.with_guardian(); + } + let mut networks = builder.build().await?; + + info!("Test networks initialized"); + info!(" - Sui RPC: {}", networks.sui_network.rpc_url); + info!(" - Bitcoin RPC: {}", networks.bitcoin_node.rpc_url()); + info!(" - Hashi nodes: {}", networks.hashi_network.nodes().len()); + + info!("Waiting for MPC key to be ready..."); + networks.hashi_network.nodes()[0] + .wait_for_mpc_key(Duration::from_secs(60)) + .await?; + info!("MPC key ready"); + + if with_guardian { + for node in networks.hashi_network.nodes() { + assert!( + node.hashi().config.guardian_endpoint().is_some(), + "guardian endpoint should be configured on every test node" + ); + assert!( + node.hashi().guardian_client().is_some(), + "guardian client should be populated after Hashi::start()" + ); + } + let harness = networks + .guardian_harness + .as_ref() + .expect("harness present when .with_guardian() is set"); + assert!( + harness.enclave().is_fully_initialized(), + "guardian harness should have reached fully-initialized state" + ); + } let deposit_amount_sats = 100_000u64; let hbtc_recipient = create_deposit_and_wait(&mut networks, deposit_amount_sats).await?; @@ -335,107 +390,7 @@ mod tests { ) .await?; - info!("=== Bitcoin Withdrawal E2E Test Passed ==="); - Ok(()) - } - - /// Same flow as `test_bitcoin_withdrawal_e2e_flow`, but with an - /// in-process guardian in the loop. - #[tokio::test] - async fn test_bitcoin_withdrawal_with_guardian_e2e_flow() -> Result<()> { - init_test_logging(); - info!("=== Starting Bitcoin Withdrawal E2E Test (with guardian) ==="); - - let mut networks = TestNetworksBuilder::new() - .with_nodes(4) - .with_guardian() - .build() - .await?; - - info!("Waiting for MPC key to be ready..."); - networks.hashi_network.nodes()[0] - .wait_for_mpc_key(Duration::from_secs(60)) - .await?; - info!("MPC key ready"); - - for node in networks.hashi_network.nodes() { - assert!( - node.hashi().config.guardian_endpoint().is_some(), - "guardian endpoint should be configured on every test node" - ); - assert!( - node.hashi().guardian_client().is_some(), - "guardian client should be populated after Hashi::start()" - ); - } - let harness = networks - .guardian_harness - .as_ref() - .expect("harness present when .with_guardian() is set"); - assert!( - harness.enclave().is_fully_initialized(), - "guardian harness should have reached fully-initialized state" - ); - - let deposit_amount_sats = 100_000u64; - let hbtc_recipient = create_deposit_and_wait(&mut networks, deposit_amount_sats).await?; - - let hashi = networks.hashi_network.nodes()[0].hashi().clone(); - let user_key = networks.sui_network.user_keys.first().unwrap(); - let withdrawal_amount_sats = 30_000u64; - let btc_destination = networks.bitcoin_node.get_new_address()?; - let destination_bytes = extract_witness_program(&btc_destination)?; - info!( - "Requesting withdrawal of {} sats to {}", - withdrawal_amount_sats, btc_destination - ); - - let mut withdrawal_executor = - SuiTxExecutor::from_config(&hashi.config, hashi.onchain_state())? - .with_signer(user_key.clone()); - withdrawal_executor - .execute_create_withdrawal_request(withdrawal_amount_sats, destination_bytes) - .await?; - - let miner = BackgroundMiner::start(&networks.bitcoin_node); - - let confirmed_event = wait_for_withdrawal_confirmation( - &mut networks.sui_network.client, - Duration::from_secs(60), - ) - .await?; - info!("Withdrawal confirmed on Sui"); - - drop(miner); - - let hbtc_balance_after = get_hbtc_balance( - &mut networks.sui_network.client, - networks.hashi_network.ids().package_id, - hbtc_recipient, - ) - .await?; - let expected_remaining = deposit_amount_sats - withdrawal_amount_sats; - assert_eq!( - hbtc_balance_after, expected_remaining, - "Expected remaining hBTC after withdrawal" - ); - - let withdrawal_txid: Txid = confirmed_event.txid.into(); - let max_output = Amount::from_sat(withdrawal_amount_sats); - let min_output = Amount::from_sat( - withdrawal_amount_sats.saturating_sub(hashi.onchain_state().worst_case_network_fee()), - ); - wait_for_withdrawal_tx_success( - &networks.bitcoin_node, - &withdrawal_txid, - &btc_destination, - max_output, - min_output, - Duration::from_secs(30), - ) - .await?; - - info!("=== Bitcoin Withdrawal E2E Test (with guardian) Passed ==="); + info!("=== Bitcoin Withdrawal E2E Test{label} Passed ==="); Ok(()) } From 4d8d06aa5731f3fd094e9a8f5ccdb1b8ff8014c4 Mon Sep 17 00:00:00 2001 From: Siddharth Sharma Date: Mon, 27 Apr 2026 14:14:40 -0700 Subject: [PATCH 5/8] [guardian] use workspace tokio-stream --- Cargo.toml | 1 + crates/e2e-tests/Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 710c6496a..800ed3b8e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,7 @@ axum = "0.8.4" tower = "0.5.3" reqwest = { version = "0.12", default-features = false, features = ["http2", "json", "rustls-tls"] } tokio = { version = "1.46.1", features = ["full"] } +tokio-stream = "0.1" tokio-util = { version = "0.7", features = ["rt"] } prometheus = "0.14" tracing = "0.1.41" diff --git a/crates/e2e-tests/Cargo.toml b/crates/e2e-tests/Cargo.toml index e3c945147..444c448db 100644 --- a/crates/e2e-tests/Cargo.toml +++ b/crates/e2e-tests/Cargo.toml @@ -36,7 +36,7 @@ clap.workspace = true colored.workspace = true rand.workspace = true nix = { version = "0.26.4", features = ["signal"] } -tokio-stream = "0.1" +tokio-stream.workspace = true tonic.workspace = true tracing-subscriber.workspace = true From 4c8fd40a1570bcbf897a43746c864654096bf18a Mon Sep 17 00:00:00 2001 From: Siddharth Sharma Date: Mon, 27 Apr 2026 14:14:45 -0700 Subject: [PATCH 6/8] [guardian] trim docs and dedupe e2e setup Reuse setup_test_networks via a shared `_with(builder)` helper so the guardian withdrawal flow doesn't re-implement the network-init logging. Trim doc comments on guardian_harness, finalize_enclave, and create_fully_initialized_enclave. --- crates/e2e-tests/src/e2e_flow.rs | 39 ++++++------------------ crates/e2e-tests/src/guardian_harness.rs | 17 +++++------ crates/hashi-guardian/src/test_utils.rs | 12 ++------ 3 files changed, 19 insertions(+), 49 deletions(-) diff --git a/crates/e2e-tests/src/e2e_flow.rs b/crates/e2e-tests/src/e2e_flow.rs index 52150f10c..b2bc8f179 100644 --- a/crates/e2e-tests/src/e2e_flow.rs +++ b/crates/e2e-tests/src/e2e_flow.rs @@ -33,8 +33,12 @@ mod tests { use crate::test_helpers::txid_to_address; async fn setup_test_networks() -> Result { + setup_test_networks_with(TestNetworksBuilder::new().with_nodes(4)).await + } + + async fn setup_test_networks_with(builder: TestNetworksBuilder) -> Result { info!("Setting up test networks..."); - let networks = TestNetworksBuilder::new().with_nodes(4).build().await?; + let networks = builder.build().await?; info!("Test networks initialized"); info!(" - Sui RPC: {}", networks.sui_network.rpc_url); @@ -258,16 +262,11 @@ mod tests { run_bitcoin_withdrawal_e2e(false).await } - /// Same flow as `test_bitcoin_withdrawal_e2e_flow`, but with an - /// in-process guardian in the loop. #[tokio::test] async fn test_bitcoin_withdrawal_with_guardian_e2e_flow() -> Result<()> { run_bitcoin_withdrawal_e2e(true).await } - /// Shared body for the bitcoin-withdrawal e2e tests, parameterized on - /// whether to spin up an in-process guardian alongside the bridge. - /// Keeping a single body forces the two flows to stay in lockstep. async fn run_bitcoin_withdrawal_e2e(with_guardian: bool) -> Result<()> { init_test_logging(); let label = if with_guardian { @@ -281,38 +280,18 @@ mod tests { if with_guardian { builder = builder.with_guardian(); } - let mut networks = builder.build().await?; - - info!("Test networks initialized"); - info!(" - Sui RPC: {}", networks.sui_network.rpc_url); - info!(" - Bitcoin RPC: {}", networks.bitcoin_node.rpc_url()); - info!(" - Hashi nodes: {}", networks.hashi_network.nodes().len()); - - info!("Waiting for MPC key to be ready..."); - networks.hashi_network.nodes()[0] - .wait_for_mpc_key(Duration::from_secs(60)) - .await?; - info!("MPC key ready"); + let mut networks = setup_test_networks_with(builder).await?; if with_guardian { for node in networks.hashi_network.nodes() { - assert!( - node.hashi().config.guardian_endpoint().is_some(), - "guardian endpoint should be configured on every test node" - ); - assert!( - node.hashi().guardian_client().is_some(), - "guardian client should be populated after Hashi::start()" - ); + assert!(node.hashi().config.guardian_endpoint().is_some()); + assert!(node.hashi().guardian_client().is_some()); } let harness = networks .guardian_harness .as_ref() .expect("harness present when .with_guardian() is set"); - assert!( - harness.enclave().is_fully_initialized(), - "guardian harness should have reached fully-initialized state" - ); + assert!(harness.enclave().is_fully_initialized()); } let deposit_amount_sats = 100_000u64; diff --git a/crates/e2e-tests/src/guardian_harness.rs b/crates/e2e-tests/src/guardian_harness.rs index c0d6c0032..1220cdfcd 100644 --- a/crates/e2e-tests/src/guardian_harness.rs +++ b/crates/e2e-tests/src/guardian_harness.rs @@ -1,10 +1,9 @@ // Copyright (c) Mysten Labs, Inc. // SPDX-License-Identifier: Apache-2.0 -//! In-process `hashi-guardian` for integration tests. Started in two -//! stages because hashi-side DKG is a prerequisite for provisioner-init: -//! [`GuardianHarness::start`] serves gRPC immediately; [`GuardianHarness::finalize`] -//! completes provisioner-init once DKG output is on chain. +//! In-process `hashi-guardian` for integration tests. Two stages: +//! [`GuardianHarness::start`] serves gRPC; [`GuardianHarness::finalize`] +//! runs provisioner-init once hashi DKG output is on chain. use anyhow::Context; use anyhow::Result; @@ -25,8 +24,7 @@ use tokio::sync::oneshot; use tokio::task::JoinHandle; use tonic::transport::Server; -/// In-process guardian reachable over gRPC on a local TCP socket. Drop -/// shuts the server down. +/// In-process guardian reachable over gRPC on a local TCP socket. pub struct GuardianHarness { enclave: Arc, endpoint: String, @@ -35,8 +33,8 @@ pub struct GuardianHarness { } impl GuardianHarness { - /// Start an operator-init'd guardian. Withdrawal RPCs stay gated - /// on [`Self::finalize`] completing provisioner-init. + /// Start an operator-init'd guardian. Withdrawal RPCs stay gated until + /// [`Self::finalize`] completes provisioner-init. pub async fn start(network: Network) -> Result { let enclave = create_operator_initialized_enclave( OperatorInitTestArgs::default().with_network(network), @@ -75,8 +73,7 @@ impl GuardianHarness { }) } - /// Complete provisioner-init with the committee + BTC master pubkey - /// from hashi's DKG. + /// Complete provisioner-init using the committee and master pubkey from hashi DKG. pub async fn finalize( &self, committee: HashiCommittee, diff --git a/crates/hashi-guardian/src/test_utils.rs b/crates/hashi-guardian/src/test_utils.rs index e95cc63c4..b0487f79a 100644 --- a/crates/hashi-guardian/src/test_utils.rs +++ b/crates/hashi-guardian/src/test_utils.rs @@ -107,12 +107,8 @@ pub struct FullyInitializedArgs { pub limiter_state: LimiterState, } -/// Drive an already operator-initialized enclave to fully-initialized state, -/// skipping the real share-encryption round-trip. Generates a fresh random -/// BTC keypair internally. -/// -/// Used by [`create_fully_initialized_enclave`] and by the `e2e-tests` -/// `GuardianHarness` when DKG output becomes available on chain. +/// Drive an operator-initialized enclave to fully-initialized state without +/// running the share-encryption round-trip. Generates a fresh BTC keypair. pub fn finalize_enclave( enclave: &Arc, committee: HashiCommittee, @@ -145,9 +141,7 @@ pub fn finalize_enclave( Ok(()) } -/// Construct an operator-initialized enclave and drive it to fully-initialized -/// state in one shot. Convenience for unit tests that don't need the -/// intermediate "operator-init only" stage. +/// Operator-init + finalize in one shot. pub async fn create_fully_initialized_enclave(args: FullyInitializedArgs) -> Arc { let FullyInitializedArgs { network, From c64923a59269e4a85ce90a541129d9eef5497eba Mon Sep 17 00:00:00 2001 From: Siddharth Sharma Date: Mon, 27 Apr 2026 14:50:39 -0700 Subject: [PATCH 7/8] [guardian] take TestNetworksBuilder param in setup_test_networks Collapse the wrapper + helper pair into one function that always takes a builder. The 3 existing call sites pass `TestNetworksBuilder::new() .with_nodes(4)` explicitly, which removes the wrapper-of-a-wrapper and keeps any future setup customization at the call site. --- crates/e2e-tests/src/e2e_flow.rs | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/crates/e2e-tests/src/e2e_flow.rs b/crates/e2e-tests/src/e2e_flow.rs index b2bc8f179..9a658692e 100644 --- a/crates/e2e-tests/src/e2e_flow.rs +++ b/crates/e2e-tests/src/e2e_flow.rs @@ -32,11 +32,7 @@ mod tests { use crate::test_helpers::lookup_vout; use crate::test_helpers::txid_to_address; - async fn setup_test_networks() -> Result { - setup_test_networks_with(TestNetworksBuilder::new().with_nodes(4)).await - } - - async fn setup_test_networks_with(builder: TestNetworksBuilder) -> Result { + async fn setup_test_networks(builder: TestNetworksBuilder) -> Result { info!("Setting up test networks..."); let networks = builder.build().await?; @@ -240,7 +236,7 @@ mod tests { init_test_logging(); info!("=== Starting Bitcoin Deposit E2E Test ==="); - let mut networks = setup_test_networks().await?; + let mut networks = setup_test_networks(TestNetworksBuilder::new().with_nodes(4)).await?; let amount_sats = 31337u64; let hbtc_recipient = create_deposit_and_wait(&mut networks, amount_sats).await?; @@ -280,7 +276,7 @@ mod tests { if with_guardian { builder = builder.with_guardian(); } - let mut networks = setup_test_networks_with(builder).await?; + let mut networks = setup_test_networks(builder).await?; if with_guardian { for node in networks.hashi_network.nodes() { @@ -416,7 +412,7 @@ mod tests { #[tokio::test] async fn test_presigning_recovery_within_batch() -> Result<()> { init_test_logging(); - let mut networks = setup_test_networks().await?; + let mut networks = setup_test_networks(TestNetworksBuilder::new().with_nodes(4)).await?; let deposit_amount_sats = 100_000u64; let withdrawal_amount_sats = 30_000u64; let user_key = networks.sui_network.user_keys.first().unwrap().clone(); @@ -677,7 +673,7 @@ mod tests { init_test_logging(); info!("=== Starting Unconfirmed Change UTXO Chaining Test ==="); - let mut networks = setup_test_networks().await?; + let mut networks = setup_test_networks(TestNetworksBuilder::new().with_nodes(4)).await?; // Deposit enough that after withdrawal 1 there is substantial change. let deposit_amount_sats = 200_000u64; From dffff49c02c365908e6743e18d792316e46651b5 Mon Sep 17 00:00:00 2001 From: Siddharth Sharma Date: Tue, 28 Apr 2026 09:30:03 -0700 Subject: [PATCH 8/8] [guardian] add is_*_set() predicates and re-privatize btc key fields --- crates/hashi-guardian/src/enclave.rs | 20 ++++++++++++++------ crates/hashi-guardian/src/init.rs | 10 +++++----- 2 files changed, 19 insertions(+), 11 deletions(-) diff --git a/crates/hashi-guardian/src/enclave.rs b/crates/hashi-guardian/src/enclave.rs index 7ad2a661e..f44f2d916 100644 --- a/crates/hashi-guardian/src/enclave.rs +++ b/crates/hashi-guardian/src/enclave.rs @@ -44,11 +44,11 @@ pub struct EnclaveConfig { /// S3 client & config (set in operator_init) s3_logger: OnceLock, /// Enclave BTC private key (set in provisioner_init) - pub(crate) enclave_btc_keypair: OnceLock, + enclave_btc_keypair: OnceLock, /// BTC network: mainnet, testnet, regtest (set in operator_init) btc_network: OnceLock, /// Hashi BTC public key used to derive child keys (set in provisioner_init) - pub(crate) hashi_btc_master_pubkey: OnceLock, + hashi_btc_master_pubkey: OnceLock, /// Withdraw related config's (set in provisioner_init) withdrawal_config: OnceLock, } @@ -187,6 +187,14 @@ impl EnclaveConfig { // Initialization Status // ======================================================================== + pub fn is_enclave_btc_keypair_set(&self) -> bool { + self.enclave_btc_keypair.get().is_some() + } + + pub fn is_hashi_btc_master_pubkey_set(&self) -> bool { + self.hashi_btc_master_pubkey.get().is_some() + } + /// Check if operator_init configuration is complete (S3 logger and network) pub fn is_operator_init_complete(&self) -> bool { self.s3_logger.get().is_some() && self.btc_network.get().is_some() @@ -199,15 +207,15 @@ impl EnclaveConfig { /// Check if provisioner_init configuration is complete (BTC keys and withdrawal config) pub fn is_provisioner_init_complete(&self) -> bool { - self.enclave_btc_keypair.get().is_some() - && self.hashi_btc_master_pubkey.get().is_some() + self.is_enclave_btc_keypair_set() + && self.is_hashi_btc_master_pubkey_set() && self.withdrawal_config.get().is_some() } /// Check if any provisioner_init configuration has been set pub fn is_provisioner_init_partially_complete(&self) -> bool { - self.enclave_btc_keypair.get().is_some() - || self.hashi_btc_master_pubkey.get().is_some() + self.is_enclave_btc_keypair_set() + || self.is_hashi_btc_master_pubkey_set() || self.withdrawal_config.get().is_some() } } diff --git a/crates/hashi-guardian/src/init.rs b/crates/hashi-guardian/src/init.rs index d8b2ebc40..be711f05a 100644 --- a/crates/hashi-guardian/src/init.rs +++ b/crates/hashi-guardian/src/init.rs @@ -284,11 +284,11 @@ mod tests { i ); assert!( - enclave.config.enclave_btc_keypair.get().is_some(), + enclave.config.is_enclave_btc_keypair_set(), "Bitcoin key should be set after threshold" ); assert!( - enclave.config.hashi_btc_master_pubkey.get().is_some(), + enclave.config.is_hashi_btc_master_pubkey_set(), "Hashi BTC key should be set after threshold" ); } else if i >= THRESHOLD { @@ -300,18 +300,18 @@ mod tests { result ); assert!( - enclave.config.enclave_btc_keypair.get().is_some(), + enclave.config.is_enclave_btc_keypair_set(), "Bitcoin key should still be set" ); } else { // Before threshold, call should succeed assert!(result.is_ok(), "Init should succeed before threshold"); assert!( - enclave.config.enclave_btc_keypair.get().is_none(), + !enclave.config.is_enclave_btc_keypair_set(), "Bitcoin key should not be set before threshold" ); assert!( - enclave.config.hashi_btc_master_pubkey.get().is_none(), + !enclave.config.is_hashi_btc_master_pubkey_set(), "Hashi BTC key should not be set before threshold" ); }