From 7197a15c78a55c0e4be3f45d830a059372ff1a85 Mon Sep 17 00:00:00 2001 From: Zhou Fang Date: Wed, 15 Apr 2026 14:11:45 -0700 Subject: [PATCH] fix!: cache rejected `send_messages` responses --- crates/hashi/src/mpc/mpc_except_signing.rs | 22 ++++++++-------- .../hashi/src/mpc/mpc_except_signing_tests.rs | 26 +++++++++---------- crates/hashi/src/mpc/types.rs | 2 +- 3 files changed, 25 insertions(+), 25 deletions(-) diff --git a/crates/hashi/src/mpc/mpc_except_signing.rs b/crates/hashi/src/mpc/mpc_except_signing.rs index c8377c2e7..d505da49f 100644 --- a/crates/hashi/src/mpc/mpc_except_signing.rs +++ b/crates/hashi/src/mpc/mpc_except_signing.rs @@ -101,7 +101,7 @@ pub struct MpcManager { pub dkg_messages: HashMap, pub rotation_messages: HashMap, pub nonce_messages: HashMap, - pub message_responses: HashMap, + pub message_responses: HashMap>, pub complaints_to_process: HashMap, pub complaint_responses: HashMap<(Address, ProtocolTypeIndicator), ComplaintResponses>, pub public_messages_store: Box, @@ -262,18 +262,18 @@ impl MpcManager { reason: "Dealer sent different messages".to_string(), }); } - if let Some(response) = self.message_responses.get(&sender) { - return Ok(response.clone()); + if let Some(cached) = self.message_responses.get(&sender) { + return cached.clone(); } tracing::info!( "handle_send_messages_request: existing message from {sender:?} but no \ cached response (e.g. post-restart), re-processing" ); } - let signature = match &request.messages { + let result = match &request.messages { Messages::Dkg(msg) => { self.store_dkg_message(self.mpc_config.epoch, sender, msg)?; - self.try_sign_dkg_message(sender, &request.messages)? + self.try_sign_dkg_message(sender, &request.messages) } Messages::Rotation(msgs) => { let previous = self @@ -281,16 +281,16 @@ impl MpcManager { .clone() .ok_or_else(|| MpcError::NotReady("Rotation not started".into()))?; self.store_rotation_messages(self.mpc_config.epoch, sender, msgs)?; - self.try_sign_rotation_messages(&previous, sender, &request.messages)? + self.try_sign_rotation_messages(&previous, sender, &request.messages) } Messages::NonceGeneration(nonce) => { self.store_nonce_message(self.mpc_config.epoch, sender, nonce)?; - self.try_sign_nonce_message(sender, &request.messages)? + self.try_sign_nonce_message(sender, &request.messages) } - }; - let response = SendMessagesResponse { signature }; - self.message_responses.insert(sender, response.clone()); - Ok(response) + } + .map(|signature| SendMessagesResponse { signature }); + self.message_responses.insert(sender, result.clone()); + result } pub fn handle_retrieve_messages_request( diff --git a/crates/hashi/src/mpc/mpc_except_signing_tests.rs b/crates/hashi/src/mpc/mpc_except_signing_tests.rs index a32aaf573..3c65b2c2c 100644 --- a/crates/hashi/src/mpc/mpc_except_signing_tests.rs +++ b/crates/hashi/src/mpc/mpc_except_signing_tests.rs @@ -4294,7 +4294,7 @@ async fn test_handle_send_messages_request_equivocation() { } #[tokio::test] -async fn test_handle_send_messages_request_invalid_shares_no_panic_on_retry() { +async fn test_handle_send_messages_request_invalid_shares_cached_on_retry() { // Second RPC call with invalid shares should not panic. let mut rng = rand::thread_rng(); @@ -4328,20 +4328,16 @@ async fn test_handle_send_messages_request_invalid_shares_no_panic_on_retry() { _ => panic!("Expected InvalidMessage error"), } - // Second call: same message — re-processes and returns the same - // "Invalid shares" error. The point of this test is that the second - // call must NOT panic; the specific error message is allowed to be - // either the original Complaint reason or a short-circuit reason. + // Second call: same message — returns the cached error immediately + // without re-processing. let result2 = receiver_manager.handle_send_messages_request(dealer_addr, &request); - assert!(result2.is_err(), "Second call should also return error"); + assert!(result2.is_err(), "Second call should return cached error"); match result2.unwrap_err() { MpcError::InvalidMessage { sender, reason } => { assert_eq!(sender, dealer_addr); assert!( - reason.contains("Invalid shares") - || reason.contains("previously received but no valid response"), - "Second call should return either the original Complaint error or a \ - 'previously received' short-circuit, got: {reason}" + reason.contains("Invalid shares"), + "Should return the cached original error, got: {reason}" ); } _ => panic!("Expected InvalidMessage error"), @@ -4353,12 +4349,16 @@ async fn test_handle_send_messages_request_invalid_shares_no_panic_on_retry() { "Message should be stored even if invalid" ); - // Verify no response was cached (since we returned error) + // Verify the error was cached so subsequent retries don't re-process. assert!( - !receiver_manager + receiver_manager .message_responses .contains_key(&dealer_addr), - "Response should not be cached for invalid shares" + "Error should be cached for invalid shares" + ); + assert!( + receiver_manager.message_responses[&dealer_addr].is_err(), + "Cached response should be an error" ); // Verify receiver can still serve the message via RetrieveMessagesRequest diff --git a/crates/hashi/src/mpc/types.rs b/crates/hashi/src/mpc/types.rs index 7f9be3036..2f6ea29b6 100644 --- a/crates/hashi/src/mpc/types.rs +++ b/crates/hashi/src/mpc/types.rs @@ -391,7 +391,7 @@ impl CertificateV1 { pub type MpcResult = Result; -#[derive(Debug, thiserror::Error)] +#[derive(Clone, Debug, thiserror::Error)] pub enum MpcError { #[error("Invalid configuration: {0}")] InvalidConfig(String),