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/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 7d744d360..444c448db 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.workspace = true +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..9a658692e 100644 --- a/crates/e2e-tests/src/e2e_flow.rs +++ b/crates/e2e-tests/src/e2e_flow.rs @@ -32,9 +32,9 @@ mod tests { use crate::test_helpers::lookup_vout; use crate::test_helpers::txid_to_address; - async fn setup_test_networks() -> Result { + async fn setup_test_networks(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); @@ -236,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?; @@ -255,10 +255,40 @@ mod tests { #[tokio::test] async fn test_bitcoin_withdrawal_e2e_flow() -> Result<()> { + run_bitcoin_withdrawal_e2e(false).await + } + + #[tokio::test] + async fn test_bitcoin_withdrawal_with_guardian_e2e_flow() -> Result<()> { + run_bitcoin_withdrawal_e2e(true).await + } + + 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 builder = TestNetworksBuilder::new().with_nodes(4); + if with_guardian { + builder = builder.with_guardian(); + } + let mut networks = setup_test_networks(builder).await?; - let mut networks = setup_test_networks().await?; + if with_guardian { + for node in networks.hashi_network.nodes() { + 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()); + } let deposit_amount_sats = 100_000u64; let hbtc_recipient = create_deposit_and_wait(&mut networks, deposit_amount_sats).await?; @@ -335,7 +365,7 @@ mod tests { ) .await?; - info!("=== Bitcoin Withdrawal E2E Test Passed ==="); + info!("=== Bitcoin Withdrawal E2E Test{label} Passed ==="); Ok(()) } @@ -382,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(); @@ -643,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; diff --git a/crates/e2e-tests/src/guardian_harness.rs b/crates/e2e-tests/src/guardian_harness.rs new file mode 100644 index 000000000..1220cdfcd --- /dev/null +++ b/crates/e2e-tests/src/guardian_harness.rs @@ -0,0 +1,136 @@ +// Copyright (c) Mysten Labs, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//! 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; +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::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. +pub struct GuardianHarness { + enclave: Arc, + endpoint: String, + shutdown_tx: Option>, + server_handle: Option>, +} + +impl GuardianHarness { + /// 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), + ) + .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, + shutdown_tx: Some(shutdown_tx), + server_handle: Some(server_handle), + }) + } + + /// Complete provisioner-init using the committee and master pubkey from hashi DKG. + pub async fn finalize( + &self, + committee: HashiCommittee, + master_pubkey: BitcoinPubkey, + withdrawal_config: WithdrawalConfig, + limiter_state: LimiterState, + ) -> Result<()> { + 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(), + "guardian did not reach fully-initialized state" + ); + Ok(()) + } + + pub fn endpoint(&self) -> &str { + &self.endpoint + } + + pub fn enclave(&self) -> &Arc { + &self.enclave + } +} + +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..f44f2d916 --- /dev/null +++ b/crates/hashi-guardian/src/enclave.rs @@ -0,0 +1,488 @@ +// 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) + 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, +} + +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 + // ======================================================================== + + 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() + } + + /// 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.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.is_enclave_btc_keypair_set() + || self.is_hashi_btc_master_pubkey_set() + || 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/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" ); } 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..b0487f79a --- /dev/null +++ b/crates/hashi-guardian/src/test_utils.rs @@ -0,0 +1,169 @@ +// 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 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, + 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); + 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)?; + 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)?; + enclave.state.init(init_state)?; + + enclave + .scratchpad + .provisioner_init_logging_complete + .set(()) + .map_err(|_| { + GuardianError::InvalidInputs("provisioner_init_logging_complete already set".into()) + })?; + Ok(()) +} + +/// Operator-init + finalize in one shot. +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 +} 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?;