From ec2002f3f409e535786b361b8c416f89313f6ede Mon Sep 17 00:00:00 2001
From: Jiloc <francescoleacche@gmail.com>
Date: Thu, 14 Nov 2024 00:12:47 +0000
Subject: [PATCH] chore: mock RotateKeysTransactions with consistent addresses
 (#833)

* mock RotateKeysTransactions with consistent addresses

* randomize test aggregate_key
---
 signer/src/testing/dummy.rs             | 19 +++++++++--
 signer/src/testing/wsts.rs              | 44 ++++++++++++++++---------
 signer/tests/integration/rotate_keys.rs |  2 +-
 3 files changed, 45 insertions(+), 20 deletions(-)

diff --git a/signer/src/testing/dummy.rs b/signer/src/testing/dummy.rs
index 0025bc2af..e4b1f4034 100644
--- a/signer/src/testing/dummy.rs
+++ b/signer/src/testing/dummy.rs
@@ -13,6 +13,7 @@ use bitcoin::TxOut;
 use bitvec::array::BitArray;
 use blockstack_lib::burnchains::Txid as StacksTxid;
 use blockstack_lib::chainstate::{nakamoto, stacks};
+use clarity::util::secp256k1::Secp256k1PublicKey;
 use fake::Fake;
 use rand::seq::IteratorRandom as _;
 use rand::Rng;
@@ -20,6 +21,8 @@ use sbtc::deposits::DepositScriptInputs;
 use sbtc::deposits::ReclaimScriptInputs;
 use secp256k1::ecdsa::RecoverableSignature;
 use secp256k1::SECP256K1;
+use stacks_common::address::AddressHashMode;
+use stacks_common::address::C32_ADDRESS_VERSION_TESTNET_MULTISIG;
 use stacks_common::types::chainstate::StacksAddress;
 
 use crate::keys::PrivateKey;
@@ -362,14 +365,24 @@ impl fake::Dummy<SignerSetConfig> for RotateKeysTransaction {
         let signer_set: Vec<PublicKey> = std::iter::repeat_with(|| fake::Faker.fake_with_rng(rng))
             .take(config.num_keys as usize)
             .collect();
-        let aggregate_key = PublicKey::combine_keys(signer_set.iter()).unwrap();
+
         let address = StacksPrincipal::from(clarity::vm::types::PrincipalData::from(
-            StacksAddress::p2pkh(false, &aggregate_key.into()),
+            StacksAddress::from_public_keys(
+                C32_ADDRESS_VERSION_TESTNET_MULTISIG,
+                &AddressHashMode::SerializeP2SH,
+                config.signatures_required as usize,
+                &signer_set
+                    .iter()
+                    .map(Secp256k1PublicKey::from)
+                    .collect::<Vec<_>>(),
+            )
+            .expect("failed to create StacksAddress"),
         ));
+
         RotateKeysTransaction {
             txid: fake::Faker.fake_with_rng(rng),
             address,
-            aggregate_key,
+            aggregate_key: fake::Faker.fake_with_rng(rng),
             signer_set,
             signatures_required: config.signatures_required,
         }
diff --git a/signer/src/testing/wsts.rs b/signer/src/testing/wsts.rs
index 10a7f7381..74f0ef719 100644
--- a/signer/src/testing/wsts.rs
+++ b/signer/src/testing/wsts.rs
@@ -1,30 +1,33 @@
 //! Test utilities for running a wsts signer and coordinator.
 
-use std::collections::{BTreeMap, BTreeSet};
+use std::collections::BTreeMap;
+use std::collections::BTreeSet;
 use std::time::Duration;
 
+use clarity::util::secp256k1::Secp256k1PublicKey;
+use clarity::vm::types::PrincipalData;
+use fake::Fake;
+use stacks_common::address::AddressHashMode;
+use stacks_common::address::C32_ADDRESS_VERSION_TESTNET_MULTISIG;
+use stacks_common::types::chainstate::StacksAddress;
+use wsts::net::SignatureType;
+use wsts::state_machine::coordinator;
+use wsts::state_machine::coordinator::frost;
+use wsts::state_machine::coordinator::Coordinator as _;
+use wsts::state_machine::StateMachine as _;
+
+use crate::ecdsa::SignEcdsa as _;
 use crate::keys::PrivateKey;
 use crate::keys::PublicKey;
 use crate::message;
 use crate::network;
+use crate::network::MessageTransfer as _;
 use crate::storage;
 use crate::storage::model;
 use crate::storage::model::EncryptedDkgShares;
 use crate::storage::model::StacksPrincipal;
 use crate::wsts_state_machine;
 
-use fake::Fake;
-use wsts::state_machine::coordinator;
-use wsts::state_machine::coordinator::frost;
-
-use crate::ecdsa::SignEcdsa as _;
-use crate::network::MessageTransfer as _;
-
-use stacks_common::types::chainstate::StacksAddress;
-use wsts::net::SignatureType;
-use wsts::state_machine::coordinator::Coordinator as _;
-use wsts::state_machine::StateMachine as _;
-
 /// Signer info
 #[derive(Debug, Clone)]
 pub struct SignerInfo {
@@ -547,10 +550,19 @@ impl SignerSet {
             tx_type: model::TransactionType::RotateKeys,
             block_hash: stacks_chain_tip.block_hash.to_bytes(),
         };
-        let address = StacksPrincipal::from(clarity::vm::types::PrincipalData::from(
-            StacksAddress::p2pkh(false, &shares.aggregate_key.into()),
+        let address = StacksPrincipal::from(PrincipalData::from(
+            StacksAddress::from_public_keys(
+                C32_ADDRESS_VERSION_TESTNET_MULTISIG,
+                &AddressHashMode::SerializeP2SH,
+                self.signers.len(),
+                &self
+                    .signer_keys()
+                    .iter()
+                    .map(Secp256k1PublicKey::from)
+                    .collect::<Vec<_>>(),
+            )
+            .expect("failed to create StacksAddress"),
         ));
-
         let rotate_keys_tx = model::RotateKeysTransaction {
             aggregate_key: shares.aggregate_key,
             address,
diff --git a/signer/tests/integration/rotate_keys.rs b/signer/tests/integration/rotate_keys.rs
index d3b684241..2046844d6 100644
--- a/signer/tests/integration/rotate_keys.rs
+++ b/signer/tests/integration/rotate_keys.rs
@@ -131,7 +131,7 @@ impl TestRotateKeySetup {
 
         let aggregate_key: PublicKey = self.aggregate_key();
         let address = StacksPrincipal::from(clarity::vm::types::PrincipalData::from(
-            StacksAddress::p2pkh(false, &aggregate_key.into()),
+            self.wallet.address().clone(),
         ));
         let rotate_key_tx = RotateKeysTransaction {
             address: address,