From 9644bf1d72597eb2f3f9baa4498d998b5a2866eb Mon Sep 17 00:00:00 2001
From: Moshe Shababo <17073733+moshababo@users.noreply.github.com>
Date: Tue, 5 Dec 2023 14:20:09 +0200
Subject: [PATCH] Validate signer to be in validator set (BFT-384) (#45)

Add validation for messages sent by replica to have their signer be in
the validator set.
---
 node/actors/bft/src/leader/replica_commit.rs  | 14 +++
 node/actors/bft/src/leader/replica_prepare.rs | 14 +++
 node/actors/bft/src/leader/tests.rs           | 89 +++++++++++--------
 node/actors/bft/src/testonly/ut_harness.rs    |  2 +
 4 files changed, 82 insertions(+), 37 deletions(-)

diff --git a/node/actors/bft/src/leader/replica_commit.rs b/node/actors/bft/src/leader/replica_commit.rs
index 2f038969..c6263409 100644
--- a/node/actors/bft/src/leader/replica_commit.rs
+++ b/node/actors/bft/src/leader/replica_commit.rs
@@ -16,6 +16,12 @@ pub(crate) enum Error {
         /// Local version.
         local_version: ProtocolVersion,
     },
+    /// Message signer isn't part of the validator set.
+    #[error("Message signer isn't part of the validator set (signer: {signer:?})")]
+    NonValidatorSigner {
+        /// Signer of the message.
+        signer: validator::PublicKey,
+    },
     /// Unexpected proposal.
     #[error("unexpected proposal")]
     UnexpectedProposal,
@@ -77,6 +83,14 @@ impl StateMachine {
             });
         }
 
+        // Check that the message signer is in the validator set.
+        consensus
+            .validator_set
+            .index(author)
+            .ok_or(Error::NonValidatorSigner {
+                signer: author.clone(),
+            })?;
+
         // If the message is from the "past", we discard it.
         if (message.view, validator::Phase::Commit) < (self.view, self.phase) {
             return Err(Error::Old {
diff --git a/node/actors/bft/src/leader/replica_prepare.rs b/node/actors/bft/src/leader/replica_prepare.rs
index 3576892e..2d06f7ab 100644
--- a/node/actors/bft/src/leader/replica_prepare.rs
+++ b/node/actors/bft/src/leader/replica_prepare.rs
@@ -18,6 +18,12 @@ pub(crate) enum Error {
         /// Local version.
         local_version: ProtocolVersion,
     },
+    /// Message signer isn't part of the validator set.
+    #[error("Message signer isn't part of the validator set (signer: {signer:?})")]
+    NonValidatorSigner {
+        /// Signer of the message.
+        signer: validator::PublicKey,
+    },
     /// Past view or phase.
     #[error("past view/phase (current view: {current_view:?}, current phase: {current_phase:?})")]
     Old {
@@ -89,6 +95,14 @@ impl StateMachine {
             });
         }
 
+        // Check that the message signer is in the validator set.
+        consensus
+            .validator_set
+            .index(author)
+            .ok_or(Error::NonValidatorSigner {
+                signer: author.clone(),
+            })?;
+
         // If the message is from the "past", we discard it.
         if (message.view, validator::Phase::Prepare) < (self.view, self.phase) {
             return Err(Error::Old {
diff --git a/node/actors/bft/src/leader/tests.rs b/node/actors/bft/src/leader/tests.rs
index f510ced0..cedc362f 100644
--- a/node/actors/bft/src/leader/tests.rs
+++ b/node/actors/bft/src/leader/tests.rs
@@ -56,24 +56,6 @@ async fn replica_prepare_sanity_yield_leader_prepare() {
     );
 }
 
-#[tokio::test]
-async fn replica_prepare_incompatible_protocol_version() {
-    let mut util = UTHarness::new_one().await;
-
-    let incompatible_protocol_version = util.incompatible_protocol_version();
-    let replica_prepare = util.new_current_replica_prepare(|msg| {
-        msg.protocol_version = incompatible_protocol_version;
-    });
-    let res = util.dispatch_replica_prepare_one(replica_prepare);
-    assert_matches!(
-        res,
-        Err(ReplicaPrepareError::IncompatibleProtocolVersion { message_version, local_version }) => {
-            assert_eq!(message_version, incompatible_protocol_version);
-            assert_eq!(local_version, util.protocol_version());
-        }
-    )
-}
-
 #[tokio::test]
 async fn replica_prepare_sanity_yield_leader_prepare_reproposal() {
     let mut util = UTHarness::new_many().await;
@@ -119,6 +101,41 @@ async fn replica_prepare_sanity_yield_leader_prepare_reproposal() {
     );
 }
 
+#[tokio::test]
+async fn replica_prepare_incompatible_protocol_version() {
+    let mut util = UTHarness::new_one().await;
+
+    let incompatible_protocol_version = util.incompatible_protocol_version();
+    let replica_prepare = util.new_current_replica_prepare(|msg| {
+        msg.protocol_version = incompatible_protocol_version;
+    });
+    let res = util.dispatch_replica_prepare_one(replica_prepare);
+    assert_matches!(
+        res,
+        Err(ReplicaPrepareError::IncompatibleProtocolVersion { message_version, local_version }) => {
+            assert_eq!(message_version, incompatible_protocol_version);
+            assert_eq!(local_version, util.protocol_version());
+        }
+    )
+}
+
+#[tokio::test]
+async fn replica_prepare_non_validator_signer() {
+    let mut util = UTHarness::new_one().await;
+
+    let replica_prepare = util.new_current_replica_prepare(|_| {}).cast().unwrap().msg;
+    let non_validator_key: validator::SecretKey = util.rng().gen();
+    let signed = non_validator_key.sign_msg(ConsensusMsg::ReplicaPrepare(replica_prepare));
+
+    let res = util.dispatch_replica_prepare_one(signed);
+    assert_matches!(
+        res,
+        Err(ReplicaPrepareError::NonValidatorSigner { signer }) => {
+            assert_eq!(signer, non_validator_key.public());
+        }
+    );
+}
+
 #[tokio::test]
 async fn replica_prepare_old_view() {
     let mut util = UTHarness::new_one().await;
@@ -272,25 +289,6 @@ async fn replica_prepare_high_qc_of_future_view() {
     );
 }
 
-#[ignore = "fails/unsupported"]
-#[tokio::test]
-async fn replica_prepare_non_validator_signer() {
-    let mut util = UTHarness::new_with(2).await;
-
-    let view = ViewNumber(2);
-    util.set_view(view);
-    assert_eq!(util.view_leader(view), util.key_at(0).public());
-
-    let replica_prepare = util.new_current_replica_prepare(|_| {});
-    let _ = util.dispatch_replica_prepare_one(replica_prepare.clone());
-
-    let non_validator: validator::SecretKey = util.rng().gen();
-    let replica_prepare = non_validator.sign_msg(replica_prepare.msg);
-    util.dispatch_replica_prepare_one(replica_prepare).unwrap();
-    // PANICS:
-    // "Couldn't create justification from valid replica messages!: Message signer isn't in the validator set"
-}
-
 #[tokio::test]
 async fn replica_commit_sanity() {
     let mut util = UTHarness::new_many().await;
@@ -356,6 +354,23 @@ async fn replica_commit_incompatible_protocol_version() {
     )
 }
 
+#[tokio::test]
+async fn replica_commit_non_validator_signer() {
+    let mut util = UTHarness::new_one().await;
+
+    let replica_commit = util.new_current_replica_commit(|_| {}).cast().unwrap().msg;
+    let non_validator_key: validator::SecretKey = util.rng().gen();
+    let signed = non_validator_key.sign_msg(ConsensusMsg::ReplicaCommit(replica_commit));
+
+    let res = util.dispatch_replica_commit_one(signed);
+    assert_matches!(
+        res,
+        Err(ReplicaCommitError::NonValidatorSigner { signer }) => {
+            assert_eq!(signer, non_validator_key.public());
+        }
+    );
+}
+
 #[tokio::test]
 async fn replica_commit_old() {
     let mut util = UTHarness::new_one().await;
diff --git a/node/actors/bft/src/testonly/ut_harness.rs b/node/actors/bft/src/testonly/ut_harness.rs
index d364ba5f..c620f567 100644
--- a/node/actors/bft/src/testonly/ut_harness.rs
+++ b/node/actors/bft/src/testonly/ut_harness.rs
@@ -323,6 +323,7 @@ impl UTHarness {
             .unwrap()
     }
 
+    #[allow(clippy::result_large_err)]
     pub(crate) fn dispatch_replica_commit_one(
         &mut self,
         msg: Signed<ConsensusMsg>,
@@ -334,6 +335,7 @@ impl UTHarness {
         )
     }
 
+    #[allow(clippy::result_large_err)]
     pub(crate) fn dispatch_replica_commit_many(
         &mut self,
         messages: Vec<ReplicaCommit>,