From 651ab96cfe8ac55e89aafdd3e6bc13087f2802c5 Mon Sep 17 00:00:00 2001 From: cylewitruk Date: Sun, 9 Feb 2025 00:07:28 +0100 Subject: [PATCH] more validation, test improvements and a little cleanup --- signer/src/dkg/testing.rs | 30 ++++--- signer/src/dkg/verification.rs | 132 ++++++++++++++++++++++--------- signer/src/transaction_signer.rs | 78 ++++++++++++------ 3 files changed, 164 insertions(+), 76 deletions(-) diff --git a/signer/src/dkg/testing.rs b/signer/src/dkg/testing.rs index 06d08130d..a5db0ec8f 100644 --- a/signer/src/dkg/testing.rs +++ b/signer/src/dkg/testing.rs @@ -14,7 +14,7 @@ use wsts::{ }; use crate::{ - keys::{PrivateKey, PublicKey}, + keys::PublicKey, wsts_state_machine::{FrostCoordinator, WstsCoordinator as _}, }; @@ -23,7 +23,6 @@ use super::{verification::StateMachine, wsts::WstsNetMessageType}; pub struct TestSetup { pub state_machine: StateMachine, pub signers: Vec>, - pub senders: Vec, #[allow(dead_code)] pub aggregate_key: XOnlyPublicKey, } @@ -39,20 +38,26 @@ impl TestSetup { let aggregate_key = pubkey_xonly(); let coordinator: FrostCoordinator = coordinators.into_iter().next().unwrap().into(); - let state_machine = StateMachine::new(coordinator, aggregate_key, Duration::from_secs(60)); + let state_machine = StateMachine::new(coordinator, aggregate_key, Duration::from_secs(60)) + .expect("failed to create new dkg verification state machine"); Self { state_machine, signers, - senders: (0..num_parties).map(|_| pubkey()).collect(), aggregate_key, } } -} -pub fn pubkey() -> PublicKey { - let keypair = secp256k1::Keypair::new_global(&mut OsRng); - PublicKey::from_private_key(&PrivateKey::from(keypair.secret_key())) + pub fn sender(&self, index: usize) -> PublicKey { + self.state_machine + .coordinator + .get_config() + .signer_public_keys + .iter() + .map(|(_, point)| crate::keys::PublicKey::try_from(point)) + .collect::, _>>() + .expect("failed to convert public keys")[index] + } } pub fn pubkey_xonly() -> secp256k1::XOnlyPublicKey { @@ -60,13 +65,6 @@ pub fn pubkey_xonly() -> secp256k1::XOnlyPublicKey { keypair.x_only_public_key().0 } -pub fn keypair() -> (PrivateKey, PublicKey) { - let keypair = secp256k1::Keypair::new_global(&mut OsRng); - let private_key = PrivateKey::from(keypair.secret_key()); - let public_key = PublicKey::from_private_key(&private_key); - (private_key, public_key) -} - pub fn nonce_request(dkg_id: u64, sign_id: u64, sign_iter_id: u64) -> Message { Message::NonceRequest(NonceRequest { dkg_id, @@ -99,7 +97,7 @@ pub fn signature_share_request( dkg_id, sign_id, sign_iter_id, - message: vec![], + message: vec![0; 5], signature_type: SignatureType::Taproot(None), nonce_responses, }) diff --git a/signer/src/dkg/verification.rs b/signer/src/dkg/verification.rs index 51c43bb8e..ff5371b42 100644 --- a/signer/src/dkg/verification.rs +++ b/signer/src/dkg/verification.rs @@ -56,6 +56,27 @@ pub enum Error { /// to be in a specific subset of WSTS coordinator states. #[error("the FROST coordinator is in an invalid state: {0:?}")] InvalidCoordinatorState(wsts::state_machine::coordinator::State), + + /// The sender is not part of the signing set. + #[error("the sender is not part of the signing set: {0}")] + UnknownSender(PublicKey), + + /// One or more of the public keys in the signing set are invalid. This error + /// can be returned when converting the public keys from the coordinator's + /// [`p256k1::point::Point`] types to [`PublicKey`] types. + #[error("one or more of the public keys in the signing set are invalid")] + InvalidPublicKeys, + + /// The public key of the signer does not match the public key of the sender. + #[error("signer public key mismatch: signer_id: {signer_id}, wsts: {wsts}, sender: {sender}")] + SignerPublicKeyMismatch { + /// The signer ID according to the FROST coordinator. + signer_id: u32, + /// The public key expected by the FROST coordinator. + wsts: Box, + /// The public key of the sender. + sender: Box, + }, } /// Represents the state of a DKG verification. @@ -150,20 +171,24 @@ pub struct StateMachine { impl StateMachine { /// Creates a new [`StateMachine`] with the given [`FrostCoordinator`], /// aggregate key, and timeout. - pub fn new(coordinator: FrostCoordinator, aggregate_key: X, timeout: Duration) -> Self + pub fn new( + coordinator: FrostCoordinator, + aggregate_key: X, + timeout: Duration, + ) -> Result where X: Into, { let aggregate_key = aggregate_key.into(); - Self { + Ok(Self { aggregate_key, coordinator, wsts_messages: HashMap::new(), created_at: Instant::now(), timeout, state: State::Idle, - } + }) } /// Processes a WSTS message, updating the internal state of the @@ -191,6 +216,9 @@ impl StateMachine { // Assert that we're in a valid state to process messages. self.assert_valid_state()?; + // Validate that the sender is a valid signer in the signing set. + self.assert_is_known_sender(sender)?; + // Enqueue the message to be processed. If the constraints for // enqueueing the message are not met, an error will be returned and we // don't make any further updates or process anything. @@ -212,6 +240,50 @@ impl StateMachine { Ok(()) } + /// Validates that the sender of a message is a valid signer in the signing + /// set. + pub fn validate_sender(&self, signer_id: u32, sender: PublicKey) -> Result<(), Error> { + let config = self.coordinator.get_config(); + let wsts: PublicKey = config + .signer_public_keys + .get(&signer_id) + .ok_or(Error::UnknownSender(sender))? + .try_into() + .map_err(|_| Error::InvalidPublicKeys)?; + + if wsts != sender { + return Err(Error::SignerPublicKeyMismatch { + signer_id, + wsts: wsts.into(), + sender: sender.into(), + }); + } + + Ok(()) + } + + /// Checks whether or not the given public key is a known sender in the + /// signing set. Returns `true` if the sender is known, `false` otherwise. + /// Returns an error if the public key is invalid. + fn assert_is_known_sender(&self, sender: PublicKey) -> Result<(), Error> { + let is_known = self + .coordinator + .get_config() + .signer_public_keys + .values() + .any(|key| { + PublicKey::try_from(key) + .map(|pub_key| pub_key == sender) + .unwrap_or(false) + }); + + if !is_known { + return Err(Error::UnknownSender(sender)); + } + + Ok(()) + } + /// Resets the [`StateMachine`] to its initial state, clearing all messages, /// setting its creation time to the current time, its state to /// [`State::Idle`] and also calling [`Coordinator::reset`] on the @@ -361,18 +433,12 @@ impl StateMachine { #[cfg(test)] mod tests { - use std::time::Duration; - use wsts::net::Message; - use crate::{ - dkg::testing::*, - testing::IterTestExt, - wsts_state_machine::{FrostCoordinator, WstsCoordinator}, - }; + use crate::{dkg::testing::*, testing::IterTestExt}; use super::{ - Error, State, StateMachine, WstsNetMessageType, WstsNetMessageType::NonceRequest, + State, WstsNetMessageType, WstsNetMessageType::NonceRequest, WstsNetMessageType::NonceResponse, WstsNetMessageType::SignatureShareResponse, }; @@ -400,10 +466,11 @@ mod tests { #[test] fn test_reset() { let signers = TestSetup::setup(5); + let sender1 = signers.sender(0); let mut state_machine = signers.state_machine; state_machine - .process_message(pubkey(), nonce_request(1, 1, 1)) + .process_message(sender1, nonce_request(1, 1, 1)) .expect("should be able to enqueue message"); assert_message_counts!(state_machine, @@ -442,18 +509,10 @@ mod tests { #[test] fn test_enqueue_message() { - let (signer_privkey, signer_pubkey) = keypair(); - let sender1 = pubkey(); - let sender2 = pubkey(); - assert_ne!(&sender1, &sender2); - - let aggregate_key = pubkey_xonly(); - - let mut state_machine = StateMachine::new( - FrostCoordinator::new([sender1, signer_pubkey], 1, signer_privkey), - aggregate_key, - Duration::from_secs(60), - ); + let setup = TestSetup::setup(2); + let sender1 = setup.sender(0); + let sender2 = setup.sender(1); + let mut state_machine = setup.state_machine; let dkg_id = 0; let sign_id = 0; @@ -490,11 +549,11 @@ mod tests { #[test] fn test_out_of_order_messages() { let mut setup = TestSetup::setup(2); + let sender1 = setup.sender(0); + let sender2 = setup.sender(1); let mut state_machine = setup.state_machine; let mut signer1 = setup.signers.pop().unwrap(); let mut signer2 = setup.signers.pop().unwrap(); - let sender1 = setup.senders[0]; - let sender2 = setup.senders[1]; let nonce_request = nonce_request(1, 1, 1); let nonce_response1 = signer1.process(&nonce_request).unwrap().single(); @@ -542,11 +601,11 @@ mod tests { #[test] fn test_nonce_phase_with_in_order_messages() { let mut setup = TestSetup::setup(2); + let sender1 = setup.sender(0); + let sender2 = setup.sender(1); let mut state_machine = setup.state_machine; let mut signer1 = setup.signers.pop().unwrap(); let mut signer2 = setup.signers.pop().unwrap(); - let sender1 = setup.senders[0]; - let sender2 = setup.senders[1]; assert_eq!(state_machine.signer_count(), 2); @@ -596,11 +655,11 @@ mod tests { #[test] fn test_dkg_verification_state_machine() { let mut setup = TestSetup::setup(2); + let sender1 = setup.sender(0); + let sender2 = setup.sender(1); let mut state_machine = setup.state_machine; let mut signer1 = setup.signers.pop().unwrap(); let mut signer2 = setup.signers.pop().unwrap(); - let sender1 = setup.senders[0]; - let sender2 = setup.senders[1]; let nonce_request = nonce_request(1, 1, 1); @@ -645,7 +704,7 @@ mod tests { 1, 1, 1, - vec![nonce_response1.clone(), nonce_response2.clone()], + vec![nonce_response2.clone(), nonce_response1.clone()], ); // Process the signature share request @@ -683,12 +742,11 @@ mod tests { // Process the second signature share response -- this should result // in the FROST coordinator transitioning into an end-state and thus // also the state machine. - let result = state_machine.process_message(sender2, sig_share_response2); + state_machine + .process_message(sender2, sig_share_response2) + .expect("should be able to process message"); - // TODO: This currently fails with a `BadPartySigs` error, so - // something's probably off with the setup of the signers. But what - // we're really testing is the state machine and that we have an end - // state here (either success or failure). - assert!(matches!(result, Err(Error::SigningFailure(_)))); + // ... which should be SUCCESS! + assert_state!(state_machine, State::Success(_)); } } diff --git a/signer/src/transaction_signer.rs b/signer/src/transaction_signer.rs index 31a086b47..05c45ea0f 100644 --- a/signer/src/transaction_signer.rs +++ b/signer/src/transaction_signer.rs @@ -600,6 +600,7 @@ where &state_machine_id, msg.id, msg_public_key, + None, &msg.inner, &chain_tip.block_hash, ) @@ -623,6 +624,7 @@ where &state_machine_id, msg.id, msg_public_key, + None, &msg.inner, &chain_tip.block_hash, ) @@ -635,11 +637,11 @@ where tracing::debug!("processing message"); let state_machine_id = StateMachineId::Dkg(*chain_tip); - self.validate_sender(&state_machine_id, request.signer_id, &msg_public_key)?; self.relay_message( &state_machine_id, msg.id, msg_public_key, + Some(request.signer_id), &msg.inner, &chain_tip.block_hash, ) @@ -652,11 +654,11 @@ where tracing::debug!("processing message"); let state_machine_id = StateMachineId::Dkg(*chain_tip); - self.validate_sender(&state_machine_id, request.signer_id, &msg_public_key)?; self.relay_message( &state_machine_id, msg.id, msg_public_key, + Some(request.signer_id), &msg.inner, &chain_tip.block_hash, ) @@ -679,6 +681,7 @@ where &state_machine_id, msg.id, msg_public_key, + None, &msg.inner, &chain_tip.block_hash, ) @@ -774,7 +777,7 @@ where ) .await?; - let (state_machine_id, _) = + let state_machine_id = self.ensure_dkg_verification_state_machine(new_key).await?; let tap_sighash = @@ -803,6 +806,7 @@ where &state_machine_id, msg.id, msg_public_key, + None, &msg.inner, &chain_tip.block_hash, ) @@ -869,7 +873,7 @@ where "🔐 responding to signature-share-request for DKG verification signing" ); - let (state_machine_id, _) = + let state_machine_id = self.ensure_dkg_verification_state_machine(new_key).await?; let tap_sighash = @@ -889,6 +893,7 @@ where &state_machine_id, msg.id, msg_public_key, + None, &msg.inner, &chain_tip.block_hash, ) @@ -921,10 +926,7 @@ where ) .await?; - let (state_machine_id, _) = - self.ensure_dkg_verification_state_machine(new_key).await?; - - self.validate_sender(&state_machine_id, request.signer_id, &msg_public_key)?; + let state_machine_id = self.ensure_dkg_verification_state_machine(new_key).await?; let tap_sighash = UnsignedMockTransaction::new(new_key.into()).compute_sighash()?; if tap_sighash.as_byte_array() != request.message.as_slice() { @@ -932,8 +934,13 @@ where return Err(Error::InvalidSigningOperation); } - self.handle_dkg_verification_message(state_machine_id, msg_public_key, &msg.inner) - .await?; + self.handle_dkg_verification_message( + state_machine_id, + msg_public_key, + Some(request.signer_id), + &msg.inner, + ) + .await?; } WstsNetMessage::SignatureShareResponse(request) => { span.record(WSTS_DKG_ID, request.dkg_id); @@ -956,13 +963,15 @@ where ) .await?; - let (state_machine_id, _) = - self.ensure_dkg_verification_state_machine(new_key).await?; - - self.validate_sender(&state_machine_id, request.signer_id, &msg_public_key)?; + let state_machine_id = self.ensure_dkg_verification_state_machine(new_key).await?; - self.handle_dkg_verification_message(state_machine_id, msg_public_key, &msg.inner) - .await?; + self.handle_dkg_verification_message( + state_machine_id, + msg_public_key, + Some(request.signer_id), + &msg.inner, + ) + .await?; } } @@ -1145,7 +1154,8 @@ where .await?; let state_machine = - dkg::verification::StateMachine::new(coordinator, aggregate_key, timeout); + dkg::verification::StateMachine::new(coordinator, aggregate_key, timeout) + .map_err(Error::DkgVerification)?; Ok(state_machine) } @@ -1160,7 +1170,7 @@ where async fn ensure_dkg_verification_state_machine( &mut self, aggregate_key: PublicKeyXOnly, - ) -> Result<(StateMachineId, &mut dkg::verification::StateMachine), Error> { + ) -> Result { let state_machine_id = StateMachineId::RotateKey(aggregate_key); if !self @@ -1210,7 +1220,7 @@ where dkg::verification::State::Idle | dkg::verification::State::Signing => {} } - Ok((state_machine_id, state_machine)) + Ok(state_machine_id) } #[tracing::instrument(skip_all)] @@ -1218,6 +1228,7 @@ where &mut self, state_machine_id: StateMachineId, sender: PublicKey, + signer_id: Option, msg: &WstsNetMessage, ) -> Result<(), Error> { // We should only be handling messages for the DKG verification state @@ -1238,6 +1249,14 @@ where return Err(Error::MissingStateMachine(state_machine_id)); }; + // Validate that the sender is a valid member of the signing set and + // has the correct id according to the state machine/coordinator. + if let Some(signer_id) = signer_id { + state_machine + .validate_sender(signer_id, sender) + .map_err(Error::DkgVerification)?; + } + tracing::trace!(?msg, "🔐 processing FROST coordinator message"); // Process the message in the DKG verification state machine. @@ -1294,6 +1313,7 @@ where state_machine_id: &StateMachineId, wsts_id: WstsMessageId, sender: PublicKey, + signer_id: Option, msg: &WstsNetMessage, bitcoin_chain_tip: &model::BitcoinBlockHash, ) -> Result<(), Error> { @@ -1306,15 +1326,22 @@ where } }; + // Validate that the sender is a valid member of the signing set and + // has the correct id according to the signer state machine. + if let Some(signer_id) = signer_id { + self.validate_sender(state_machine_id, signer_id, &sender)?; + } + // Check and store if this is a DKG verification-related message. let is_dkg_verification = matches!(state_machine_id, StateMachineId::RotateKey(_)); // If this is a DKG verification then we need to process the message in // the frost coordinator as well to be able to properly follow the // signing round (which is otherwise handled by the signer state - // machine). + // machine). We pass `None` as the `signer_id` because we have just + // validated the sender above. if is_dkg_verification { - self.handle_dkg_verification_message(*state_machine_id, sender, msg) + self.handle_dkg_verification_message(*state_machine_id, sender, None, msg) .await?; } @@ -1333,8 +1360,13 @@ where // in the FROST state machine as well for it to properly follow // the signing round. if is_dkg_verification { - self.handle_dkg_verification_message(*state_machine_id, sender, outbound_message) - .await?; + self.handle_dkg_verification_message( + *state_machine_id, + sender, + signer_id, + outbound_message, + ) + .await?; } }