From dffecbc3b4ab188dfde4e278307dd911d968b459 Mon Sep 17 00:00:00 2001
From: Wen <113942165+wen-coding@users.noreply.github.com>
Date: Sat, 15 Jun 2024 10:50:32 -0700
Subject: [PATCH] wen_restart: Ignore Gossip messages from my own pubkey.
 (#1678)

---
 wen-restart/src/heaviest_fork_aggregate.rs    | 41 +++++++++++++++++
 .../src/last_voted_fork_slots_aggregate.rs    | 45 +++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/wen-restart/src/heaviest_fork_aggregate.rs b/wen-restart/src/heaviest_fork_aggregate.rs
index ce1c19283876ce..0b43b800d18573 100644
--- a/wen-restart/src/heaviest_fork_aggregate.rs
+++ b/wen-restart/src/heaviest_fork_aggregate.rs
@@ -118,6 +118,9 @@ impl HeaviestForkAggregate {
             );
             return None;
         }
+        if from == &self.my_pubkey {
+            return None;
+        }
         if received_heaviest_fork.shred_version != self.my_shred_version {
             warn!(
                 "Gossip should not accept RestartLastVotedFork with different shred version {} from {:?}",
@@ -445,6 +448,23 @@ mod tests {
                 .total_active_stake_seen_supermajority(),
             1500
         );
+
+        // test that message from my pubkey is ignored.
+        assert_eq!(
+            test_state
+                .heaviest_fork_aggregate
+                .aggregate(RestartHeaviestFork {
+                    from: test_state.validator_voting_keypairs[MY_INDEX]
+                        .node_keypair
+                        .pubkey(),
+                    wallclock: timestamp(),
+                    last_slot: test_state.heaviest_slot,
+                    last_slot_hash: test_state.heaviest_hash,
+                    observed_stake: 100,
+                    shred_version: SHRED_VERSION,
+                },),
+            None,
+        );
     }
 
     #[test]
@@ -535,6 +555,27 @@ mod tests {
         );
         // percentage doesn't change since the previous aggregate is ignored.
         assert_eq!(test_state.heaviest_fork_aggregate.total_active_stake(), 200);
+
+        // Record from my pubkey should be ignored.
+        assert_eq!(
+            test_state
+                .heaviest_fork_aggregate
+                .aggregate_from_record(
+                    &test_state.validator_voting_keypairs[MY_INDEX]
+                        .node_keypair
+                        .pubkey()
+                        .to_string(),
+                    &HeaviestForkRecord {
+                        wallclock: timestamp(),
+                        slot: test_state.heaviest_slot,
+                        bankhash: test_state.heaviest_hash.to_string(),
+                        shred_version: SHRED_VERSION as u32,
+                        total_active_stake: 100,
+                    }
+                )
+                .unwrap(),
+            None,
+        );
     }
 
     #[test]
diff --git a/wen-restart/src/last_voted_fork_slots_aggregate.rs b/wen-restart/src/last_voted_fork_slots_aggregate.rs
index b042ba80d9bb51..67cd7c1c77ea87 100644
--- a/wen-restart/src/last_voted_fork_slots_aggregate.rs
+++ b/wen-restart/src/last_voted_fork_slots_aggregate.rs
@@ -20,6 +20,7 @@ pub(crate) struct LastVotedForkSlotsAggregate {
     slots_stake_map: HashMap<Slot, u64>,
     active_peers: HashSet<Pubkey>,
     slots_to_repair: HashSet<Slot>,
+    my_pubkey: Pubkey,
 }
 
 #[derive(Clone, Debug, PartialEq)]
@@ -53,6 +54,7 @@ impl LastVotedForkSlotsAggregate {
             slots_stake_map,
             active_peers,
             slots_to_repair: HashSet::new(),
+            my_pubkey: *my_pubkey,
         }
     }
 
@@ -70,6 +72,9 @@ impl LastVotedForkSlotsAggregate {
         record: &LastVotedForkSlotsRecord,
     ) -> Result<Option<LastVotedForkSlotsRecord>> {
         let from = Pubkey::from_str(key_string)?;
+        if from == self.my_pubkey {
+            return Ok(None);
+        }
         let last_voted_hash = Hash::from_str(&record.last_vote_bankhash)?;
         let converted_record = RestartLastVotedForkSlots::new(
             from,
@@ -88,6 +93,9 @@ impl LastVotedForkSlotsAggregate {
         let total_stake = self.epoch_stakes.total_stake();
         let threshold_stake = (total_stake as f64 * self.repair_threshold) as u64;
         let from = &new_slots.from;
+        if from == &self.my_pubkey {
+            return None;
+        }
         let sender_stake = Self::validator_stake(&self.epoch_stakes, from);
         if sender_stake == 0 {
             warn!(
@@ -354,6 +362,23 @@ mod tests {
             Vec::from_iter(test_state.slots_aggregate.slots_to_repair_iter().cloned());
         actual_slots.sort();
         assert_eq!(actual_slots, vec![root_slot + 1]);
+
+        // test that message from my pubkey is ignored.
+        assert_eq!(
+            test_state.slots_aggregate.aggregate(
+                RestartLastVotedForkSlots::new(
+                    test_state.validator_voting_keypairs[MY_INDEX]
+                        .node_keypair
+                        .pubkey(),
+                    timestamp(),
+                    &[root_slot + 1, root_slot + 4, root_slot + 5],
+                    Hash::default(),
+                    SHRED_VERSION,
+                )
+                .unwrap(),
+            ),
+            None,
+        );
     }
 
     #[test]
@@ -446,6 +471,26 @@ mod tests {
         );
         // percentage doesn't change since the previous aggregate is ignored.
         assert_eq!(test_state.slots_aggregate.active_percent(), 20.0);
+
+        // Record from my pubkey should be ignored.
+        assert_eq!(
+            test_state
+                .slots_aggregate
+                .aggregate_from_record(
+                    &test_state.validator_voting_keypairs[MY_INDEX]
+                        .node_keypair
+                        .pubkey()
+                        .to_string(),
+                    &LastVotedForkSlotsRecord {
+                        wallclock: timestamp(),
+                        last_voted_fork_slots: vec![root_slot + 10, root_slot + 300],
+                        last_vote_bankhash: Hash::new_unique().to_string(),
+                        shred_version: SHRED_VERSION as u32,
+                    }
+                )
+                .unwrap(),
+            None,
+        );
     }
 
     #[test]