From 6a650e79b584ac1fd0f7e01dcf373830578ab496 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Esteban=20Dimitroff=20H=C3=B3di?=
 <esteban.dimitroff@lambdaclass.com>
Date: Wed, 27 Mar 2024 13:59:52 -0300
Subject: [PATCH] Checking weights sum for overflow

---
 .../roles/src/validator/messages/consensus.rs | 10 +++++++-
 node/libs/roles/src/validator/tests.rs        | 23 +++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/node/libs/roles/src/validator/messages/consensus.rs b/node/libs/roles/src/validator/messages/consensus.rs
index 183bb4b2..ac30b40e 100644
--- a/node/libs/roles/src/validator/messages/consensus.rs
+++ b/node/libs/roles/src/validator/messages/consensus.rs
@@ -77,6 +77,7 @@ impl Default for Fork {
 pub struct ValidatorCommittee {
     vec: Vec<WeightedValidator>,
     indexes: BTreeMap<validator::PublicKey, usize>,
+    total_weight: u64,
 }
 
 impl ValidatorCommittee {
@@ -84,11 +85,17 @@ impl ValidatorCommittee {
     pub fn new(validators: impl IntoIterator<Item = WeightedValidator>) -> anyhow::Result<Self> {
         let mut set = BTreeSet::new();
         let mut weighted_validators = BTreeMap::new();
+        let mut total_weight: u64 = 0;
         for validator in validators {
             anyhow::ensure!(
                 set.insert(validator.key.clone()),
                 "Duplicate validator in ValidatorCommittee"
             );
+            if let Some(result) = total_weight.checked_add(validator.weight) {
+                total_weight = result
+            } else {
+                anyhow::bail!("Sum of weights overflows in ValidatorCommittee")
+            }
             weighted_validators.insert(validator.key.clone(), validator);
         }
         anyhow::ensure!(
@@ -106,6 +113,7 @@ impl ValidatorCommittee {
                 .enumerate()
                 .map(|(i, pk)| (pk, i))
                 .collect(),
+            total_weight,
         })
     }
 
@@ -168,7 +176,7 @@ impl ValidatorCommittee {
 
     /// Sum of all validators' weight in the committee
     pub fn total_weight(&self) -> u64 {
-        self.vec.iter().map(|v| v.weight).sum()
+        self.total_weight
     }
 }
 
diff --git a/node/libs/roles/src/validator/tests.rs b/node/libs/roles/src/validator/tests.rs
index e1c3ce00..07e9dcf3 100644
--- a/node/libs/roles/src/validator/tests.rs
+++ b/node/libs/roles/src/validator/tests.rs
@@ -1,5 +1,6 @@
 use super::*;
 use crate::validator::testonly::Setup;
+use anyhow::Error;
 use assert_matches::assert_matches;
 use rand::{seq::SliceRandom, Rng};
 use std::vec;
@@ -313,3 +314,25 @@ fn test_validator_committee_weights() {
         );
     }
 }
+
+#[test]
+fn test_validator_weights_sanity() {
+    let ctx = ctx::test_root(&ctx::RealClock);
+    let rng = &mut ctx.rng();
+
+    let setup = Setup::new(rng, 6);
+    // Validators weights
+    let weight = u64::MAX / 5;
+    let weights = [weight, weight, weight, weight, weight, weight];
+    let validators: Vec<WeightedValidator> = weights
+        .iter()
+        .enumerate()
+        .map(|(i, w)| WeightedValidator {
+            key: setup.keys[i].public(),
+            weight: *w,
+        })
+        .collect();
+
+    // Creation should overflow
+    assert_matches!(ValidatorCommittee::new(validators), Err(Error { .. }));
+}