Mid term double stake protection

massa-consensus-worker/src/state/mod.rs

@@ -90,6 +90,9 @@ pub struct ConsensusState {
     pub wishlist: PreHashMap<BlockId, Option<SecuredHeader>>,
     /// previous blockclique notified to Execution
     pub prev_blockclique: PreHashMap<BlockId, Slot>,
+    /// Blocks for each slots (used for multi-stake). Blocks
+    /// should be saved in this map when we receive the header or directly full block
+    pub nonfinal_active_blocks_per_slot: HashMap<Slot, PreHashSet<BlockId>>,
 }
 
 impl ConsensusState {

massa-consensus-worker/src/state/process.rs

@@ -20,7 +20,7 @@ use massa_models::{
 use massa_signature::PublicKey;
 use massa_storage::Storage;
 use massa_time::MassaTime;
-use tracing::log::{debug, info};
+use tracing::log::{debug, info, warn};
 
 use super::ConsensusState;
 
@@ -112,6 +112,22 @@ impl ConsensusState {
                         block_id
                     )));
                 };
+                // Verify that we haven't already received 2 blocks for this slot
+                let entry = self
+                    .nonfinal_active_blocks_per_slot
+                    .entry(header.content.slot)
+                    .or_default();
+                if !entry.contains(&header.id) {
+                    if entry.len() > 1 && !self.wishlist.contains_key(&block_id) {
+                        warn!(
+                            "received more than 2 blocks for slot {}",
+                            header.content.slot
+                        );
+                        return Ok(BTreeSet::new());
+                    } else {
+                        entry.insert(block_id);
+                    }
+                }
                 self.check_block_header_and_store(block_id, header, current_slot)?;
                 return Ok(BTreeSet::new());
             }
@@ -140,7 +156,22 @@ impl ConsensusState {
                     .get(&block_id)
                     .cloned()
                     .expect("incoming block not found in storage");
-
+                // Verify that we haven't already received 2 blocks for this slot
+                let entry = self
+                    .nonfinal_active_blocks_per_slot
+                    .entry(stored_block.content.header.content.slot)
+                    .or_default();
+                if !entry.contains(&stored_block.id) {
+                    if entry.len() > 1 && !self.wishlist.contains_key(&block_id) {
+                        warn!(
+                            "received more than 2 blocks for slot {}",
+                            stored_block.content.header.content.slot
+                        );
+                        return Ok(BTreeSet::default());
+                    } else {
+                        entry.insert(block_id);
+                    }
+                }
                 match self.check_block_and_store(
                     block_id,
                     slot,
@@ -149,7 +180,7 @@ impl ConsensusState {
                     current_slot,
                 )? {
                     Some(block_infos) => block_infos,
-                    None => return Ok(BTreeSet::default()),
+                    None => return Ok(BTreeSet::new()),
                 }
             }
 

massa-consensus-worker/src/state/prune.rs

@@ -135,7 +135,7 @@ impl ConsensusState {
         });
     }
 
-    // Keep only a certain (`config.max_future_processing_blocks`) number of blocks that are discarded
+    // Keep only a certain (`config.max_discarded_blocks`) number of blocks that are discarded
     // to avoid too big memory consumption
     fn prune_discarded(&mut self) -> Result<(), ConsensusError> {
         if self.discarded_index.len() <= self.config.max_discarded_blocks {
@@ -332,6 +332,21 @@ impl ConsensusState {
         Ok(())
     }
 
+    fn prune_nonfinal_blocks_per_slot(&mut self) {
+        let keys = self
+            .nonfinal_active_blocks_per_slot
+            .keys()
+            .cloned()
+            .collect::<Vec<_>>();
+        for (thread, (_, period)) in self.latest_final_blocks_periods.iter().enumerate() {
+            for key in &keys {
+                if key.thread == thread as u8 && key.period <= *period {
+                    self.nonfinal_active_blocks_per_slot.remove(&key);
+                }
+            }
+        }
+    }
+
     pub fn prune(&mut self) -> Result<(), ConsensusError> {
         let before = self.max_cliques.len();
         // Step 1: discard final blocks that are not useful to the graph anymore and return them
@@ -346,6 +361,9 @@ impl ConsensusState {
         // Step 4: prune discarded
         self.prune_discarded()?;
 
+        // Step 5: prune nonfinal blocks per slot
+        self.prune_nonfinal_blocks_per_slot();
+
         let after = self.max_cliques.len();
         if before != after {
             debug!(

massa-consensus-worker/src/worker/mod.rs

@@ -96,6 +96,7 @@ pub fn start_consensus_worker(
             config.stats_timespan,
         ),
         prev_blockclique: Default::default(),
+        nonfinal_active_blocks_per_slot: Default::default(),
     }));
 
     let shared_state_cloned = shared_state.clone();