From 7fbcf02bb38382f674e36bf45c76ab5b41bc8fcc Mon Sep 17 00:00:00 2001
From: Jude Nelson <judecn@gmail.com>
Date: Thu, 8 Sep 2022 15:02:39 -0400
Subject: [PATCH] refactor: track key registration with txid, not
 BurnchainSigner

---
 testnet/stacks-node/src/node.rs | 64 +++++++++++++++++++++++++--------
 1 file changed, 49 insertions(+), 15 deletions(-)

diff --git a/testnet/stacks-node/src/node.rs b/testnet/stacks-node/src/node.rs
index 45dae4db8..714f4839a 100644
--- a/testnet/stacks-node/src/node.rs
+++ b/testnet/stacks-node/src/node.rs
@@ -30,7 +30,7 @@ use stacks::net::{
     Error as NetError, PeerAddress,
 };
 use stacks::types::chainstate::TrieHash;
-use stacks::types::chainstate::{BlockHeaderHash, BurnchainHeaderHash, StacksAddress, VRFSeed};
+use stacks::types::chainstate::{BlockHeaderHash, BurnchainHeaderHash, VRFSeed};
 use stacks::util::get_epoch_time_secs;
 use stacks::util::hash::Sha256Sum;
 use stacks::util::secp256k1::Secp256k1PrivateKey;
@@ -98,6 +98,8 @@ pub struct Node {
     event_dispatcher: EventDispatcher,
     nonce: u64,
     attachments_tx: SyncSender<HashSet<AttachmentInstance>>,
+    leader_key_registers: HashSet<Txid>,
+    block_commits: HashSet<Txid>,
 }
 
 pub fn get_account_lockups(
@@ -375,6 +377,8 @@ impl Node {
             nonce: 0,
             event_dispatcher,
             attachments_tx,
+            leader_key_registers: HashSet::new(),
+            block_commits: HashSet::new(),
         }
     }
 
@@ -418,6 +422,8 @@ impl Node {
             nonce: 0,
             event_dispatcher,
             attachments_tx,
+            leader_key_registers: HashSet::new(),
+            block_commits: HashSet::new(),
         };
 
         node.spawn_peer_server(attachments_rx);
@@ -580,9 +586,27 @@ impl Node {
             .keychain
             .rotate_vrf_keypair(burnchain_tip.block_snapshot.block_height);
         let consensus_hash = burnchain_tip.block_snapshot.consensus_hash;
+
+        let burnchain = Burnchain::regtest(&self.config.get_burn_db_path());
+        let sortdb = SortitionDB::open(
+            &self.config.get_burn_db_file_path(),
+            true,
+            burnchain.pox_constants.clone(),
+        )
+        .expect("Error while opening sortition db");
+
+        let cur_epoch =
+            SortitionDB::get_stacks_epoch(sortdb.conn(), burnchain_tip.block_snapshot.block_height)
+                .expect("FATAL: failed to read sortition DB")
+                .expect("FATAL: no epoch defined");
+
         let key_reg_op = self.generate_leader_key_register_op(vrf_pk, &consensus_hash);
         let mut op_signer = self.keychain.generate_op_signer();
-        burnchain_controller.submit_operation(key_reg_op, &mut op_signer, 1);
+        let key_txid = burnchain_controller
+            .submit_operation(cur_epoch.epoch_id, key_reg_op, &mut op_signer, 1)
+            .expect("FATAL: failed to submit leader key register operation");
+
+        self.leader_key_registers.insert(key_txid);
     }
 
     /// Process an state coming from the burnchain, by extracting the validated KeyRegisterOp
@@ -600,7 +624,7 @@ impl Node {
         for op in ops.iter() {
             match op {
                 BlockstackOperationType::LeaderKeyRegister(ref op) => {
-                    if op.address == self.keychain.get_address(is_mainnet) {
+                    if self.leader_key_registers.contains(&op.txid) {
                         // Registered key has been mined
                         new_key = Some(RegisteredKey {
                             vrf_public_key: op.public_key.clone(),
@@ -612,15 +636,12 @@ impl Node {
                 BlockstackOperationType::LeaderBlockCommit(ref op) => {
                     if op.txid == burnchain_tip.block_snapshot.winning_block_txid {
                         last_sortitioned_block = Some(burnchain_tip.clone());
-                        if op.apparent_sender == self.keychain.get_burnchain_signer() {
+                        if self.block_commits.contains(&op.txid) {
                             won_sortition = true;
                         }
                     }
                 }
-                BlockstackOperationType::PreStx(_)
-                | BlockstackOperationType::StackStx(_)
-                | BlockstackOperationType::TransferStx(_)
-                | BlockstackOperationType::UserBurnSupport(_) => {
+                _ => {
                     // no-op, ops are not supported / produced at this point.
                 }
             }
@@ -764,8 +785,27 @@ impl Node {
                 VRFSeed::from_proof(&vrf_proof),
             );
 
+            let burnchain = Burnchain::regtest(&self.config.get_burn_db_path());
+            let sortdb = SortitionDB::open(
+                &self.config.get_burn_db_file_path(),
+                true,
+                burnchain.pox_constants.clone(),
+            )
+            .expect("Error while opening sortition db");
+
+            let cur_epoch = SortitionDB::get_stacks_epoch(
+                sortdb.conn(),
+                burnchain_tip.block_snapshot.block_height,
+            )
+            .expect("FATAL: failed to read sortition DB")
+            .expect("FATAL: no epoch defined");
+
             let mut op_signer = self.keychain.generate_op_signer();
-            burnchain_controller.submit_operation(op, &mut op_signer, 1);
+            let txid = burnchain_controller
+                .submit_operation(cur_epoch.epoch_id, op, &mut op_signer, 1)
+                .expect("FATAL: failed to submit block-commit");
+
+            self.block_commits.insert(txid);
         }
     }
 
@@ -961,11 +1001,6 @@ impl Node {
         attachments_instances
     }
 
-    /// Returns the Stacks address of the node
-    pub fn get_address(&self) -> StacksAddress {
-        self.keychain.get_address(self.config.is_mainnet())
-    }
-
     /// Constructs and returns a LeaderKeyRegisterOp out of the provided params
     fn generate_leader_key_register_op(
         &mut self,
@@ -975,7 +1010,6 @@ impl Node {
         BlockstackOperationType::LeaderKeyRegister(LeaderKeyRegisterOp {
             public_key: vrf_public_key,
             memo: vec![],
-            address: self.keychain.get_address(self.config.is_mainnet()),
             consensus_hash: consensus_hash.clone(),
             vtxindex: 0,
             txid: Txid([0u8; 32]),