[Guardian] Improve S3 log write robustness

crates/hashi-guardian/src/enclave.rs

@@ -23,6 +23,7 @@ use std::sync::Arc;
 use std::sync::OnceLock;
 use std::sync::RwLock;
 use std::time::Duration;
+use tokio::sync::OwnedMutexGuard;
 use tracing::info;
 
 use crate::s3_client::GuardianS3Client;
@@ -283,19 +284,19 @@ impl EnclaveState {
             .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);
 
+    /// Acquire exclusive access to the limiter, consume tokens, and return a guard.
+    /// The guard is held through signing and durable logging so no other withdrawal
+    /// can start until this one is durably logged or the enclave aborts.
     pub async fn consume_from_limiter(
         &self,
         seq: u64,
         timestamp: u64,
         amount_sats: u64,
-    ) -> GuardianResult<LimiterGuard> {
+    ) -> GuardianResult<OwnedMutexGuard<RateLimiter>> {
         let rate_limiter = self
             .rate_limiter
             .get()
@@ -307,50 +308,29 @@ impl EnclaveState {
         .await
         .map_err(|_| InvalidInputs("timed out waiting for rate limiter lock".into()))?;
         guard.consume(seq, timestamp, amount_sats)?;
-        Ok(LimiterGuard::new(guard))
+        Ok(guard)
     }
 
     pub async fn limiter_state(&self) -> Option<LimiterState> {
-        let limiter = self.rate_limiter.get()?;
-        Some(*limiter.lock().await.state())
+        Some(*self.lock_limiter_for_read().await?.state())
     }
 
     pub async fn limiter_config(&self) -> Option<LimiterConfig> {
-        let limiter = self.rate_limiter.get()?;
-        Some(*limiter.lock().await.config())
-    }
-}
-
-/// RAII guard that holds the limiter mutex via an owned guard, returned by
-/// [`EnclaveState::consume_from_limiter`]. Reverts on drop unless committed.
-pub struct LimiterGuard {
-    guard: tokio::sync::OwnedMutexGuard<RateLimiter>,
-    committed: bool,
-}
-
-impl LimiterGuard {
-    pub(crate) fn new(guard: tokio::sync::OwnedMutexGuard<RateLimiter>) -> Self {
-        Self {
-            guard,
-            committed: false,
-        }
+        Some(*self.lock_limiter_for_read().await?.config())
     }
 
-    /// Mark this withdrawal as successful. Prevents revert on drop.
-    pub fn commit(mut self) {
-        self.committed = true;
-    }
-
-    pub fn state(&self) -> &LimiterState {
-        self.guard.state()
-    }
-}
-
-impl Drop for LimiterGuard {
-    fn drop(&mut self) {
-        if !self.committed {
-            self.guard.revert();
-        }
+    /// Lock the limiter for a read-only status query, bounded by
+    /// `LIMITER_LOCK_TIMEOUT` so an in-flight durable write cannot stall `info`
+    /// for minutes. Returns None if uninitialized or the lock is still held at
+    /// the deadline.
+    async fn lock_limiter_for_read(&self) -> Option<OwnedMutexGuard<RateLimiter>> {
+        let rate_limiter = self.rate_limiter.get()?;
+        tokio::time::timeout(
+            Self::LIMITER_LOCK_TIMEOUT,
+            rate_limiter.clone().lock_owned(),
+        )
+        .await
+        .ok()
     }
 }
 
@@ -541,29 +521,48 @@ impl Enclave {
         self.config.s3_logger()?.write_log_record(log).await
     }
 
+    async fn write_log_or_abort(&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_or_abort(log)
+            .await
+    }
+
+    /// Only init skips grace-period retries, providing quick fail-stop for basic
+    /// S3 write/access issues. The incomplete enclave cannot serve and will restart,
+    /// so S3 being ahead after a lost acknowledgement is acceptable.
     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
+        self.write_log_or_abort(LogMessage::Withdrawal(Box::new(msg)))
+            .await
     }
 
     pub async fn log_committee_update(&self, msg: CommitteeUpdateLogMessage) -> GuardianResult<()> {
-        self.write_log(LogMessage::CommitteeUpdate(Box::new(msg)))
+        self.write_log_or_abort(LogMessage::CommitteeUpdate(Box::new(msg)))
             .await
     }
 
     pub async fn log_genesis(&self, msg: GenesisLogMessage) -> GuardianResult<()> {
-        self.write_log(LogMessage::Genesis(Box::new(msg))).await
+        self.write_log_or_abort(LogMessage::Genesis(Box::new(msg)))
+            .await
     }
 
     pub async fn log_heartbeat(&self, seq: u64) -> GuardianResult<()> {
-        self.write_log(LogMessage::Heartbeat { seq }).await
+        self.write_log_or_abort(LogMessage::Heartbeat { seq }).await
     }
 
     pub async fn log_ceremony(&self, state: CeremonyLogMessage) -> GuardianResult<()> {
-        self.write_log(LogMessage::Ceremony(Box::new(state))).await
+        self.write_log_or_abort(LogMessage::Ceremony(Box::new(state)))
+            .await
     }
 
     /// Persist the current encrypted KP share state to `kp-shares/` for recovery.
@@ -575,7 +574,7 @@ impl Enclave {
         cert_seq: u64,
         encrypted_shares: KPEncryptedShares,
     ) -> GuardianResult<()> {
-        self.write_log(LogMessage::KpShareState(Box::new(KpShareState::new(
+        self.write_log_or_abort(LogMessage::KpShareState(Box::new(KpShareState::new(
             sharing_seq,
             cert_seq,
             encrypted_shares,

crates/hashi-guardian/src/lib.rs

@@ -3,10 +3,21 @@
 
 use std::time::Duration;
 
-// TODO: Leave as consts or make them configurable?
+/// Interval between successful heartbeat writes.
 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);
+/// Maximum write failure interval before the enclave aborts.
+pub const MAX_S3_WRITE_FAILURE_INTERVAL: Duration = Duration::from_mins(5);
+/// Maximum acceptable heartbeat age for the activating standby session.
+pub const LIVE_SESSION_MAX_AGE: Duration = Duration::from_mins(3);
+/// Silence required before another session is considered inactive.
+pub const OTHER_SESSION_QUIET_PERIOD: Duration = Duration::from_mins(10);
+
+// An enclave that cannot write its next heartbeat must abort before another
+// session is allowed to treat it as quiet.
+const _: () = assert!(
+    HEARTBEAT_INTERVAL.as_secs() + MAX_S3_WRITE_FAILURE_INTERVAL.as_secs()
+        < OTHER_SESSION_QUIET_PERIOD.as_secs()
+);
 
 pub mod attestation;
 pub mod ceremony_mode;

crates/hashi-guardian/src/main.rs

@@ -6,8 +6,6 @@ use hashi_guardian::rpc::GuardianGrpc;
 use hashi_guardian::withdraw_mode::heartbeat::HeartbeatWriter;
 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::EnclaveMode;
 use hashi_types::guardian::GuardianEncKeyPair;
 use hashi_types::guardian::GuardianSignKeyPair;
@@ -68,31 +66,21 @@ async fn main() -> Result<()> {
         .set_serving::<GuardianServiceServer<GuardianGrpc>>()
         .await;
 
-    let server_future = Server::builder()
-        .add_service(health_service)
-        .add_service(GuardianServiceServer::new(svc))
-        .serve(addr);
-
     // Don't emit heartbeats in ceremony mode: their primary function is
     // to allow KPs to detect old sessions that might still be running
     // in order to bypass limiter. Not a concern for ceremony mode.
-    if ceremony_mode {
-        return server_future
-            .await
-            .map_err(|e| anyhow::anyhow!("Server error: {}", e));
+    if !ceremony_mode {
+        drop(tokio::spawn(
+            HeartbeatWriter::new(enclave).run(HEARTBEAT_INTERVAL),
+        ));
     }
 
-    let heartbeat_future = HeartbeatWriter::new(enclave, MAX_HEARTBEAT_FAILURES_INTERVAL)
-        .run(HEARTBEAT_INTERVAL, HEARTBEAT_RETRY_INTERVAL);
-
-    tokio::select! {
-        res = server_future => {
-            res.map_err(|e| anyhow::anyhow!("Server error: {}", e))
-        }
-        res = heartbeat_future => {
-            panic!("Heartbeat failed: {:?}", res)
-        }
-    }
+    Server::builder()
+        .add_service(health_service)
+        .add_service(GuardianServiceServer::new(svc))
+        .serve(addr)
+        .await
+        .map_err(|e| anyhow::anyhow!("Server error: {}", e))
 }
 
 /// Make any panic abort the process instead of unwinding to the tokio task

crates/hashi-guardian/src/rpc.rs

@@ -9,7 +9,7 @@ use crate::operator_init;
 use crate::withdraw_mode::committee_update;
 use crate::withdraw_mode::genesis;
 use crate::withdraw_mode::provisioner_init;
-use crate::withdraw_mode::standard;
+use crate::withdraw_mode::standard_withdrawal;
 use crate::Enclave;
 use hashi_types::guardian::proto_conversions;
 use hashi_types::guardian::proto_conversions::pb_to_signed_committee_transition;
@@ -193,9 +193,10 @@ impl proto::guardian_service_server::GuardianService for GuardianGrpc {
                 .map_err(to_status)?;
 
         // core withdraw call
-        let response = standard::standard_withdrawal(self.enclave.clone(), validated_req)
-            .await
-            .map_err(to_status)?;
+        let response =
+            standard_withdrawal::standard_withdrawal(self.enclave.clone(), validated_req)
+                .await
+                .map_err(to_status)?;
 
         // domain to proto
         let resp_pb = proto_conversions::standard_withdrawal_response_signed_to_pb(response);