Adds a restrictions ID check into handshake protocols

Cargo.toml

@@ -46,7 +46,7 @@ default-features = false
 
 [workspace.dependencies.snarkvm]
 git = "https://github.com/AleoNet/snarkVM.git"
-rev = "74a437897"
+rev = "06bba06"
 #version = "=0.16.18"
 features = [ "circuit", "console", "rocks" ]
 

node/bft/events/src/challenge_response.rs

@@ -16,6 +16,7 @@ use super::*;
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct ChallengeResponse<N: Network> {
+    pub restrictions_id: Field<N>,
     pub signature: Data<Signature<N>>,
     pub nonce: u64,
 }
@@ -30,6 +31,7 @@ impl<N: Network> EventTrait for ChallengeResponse<N> {
 
 impl<N: Network> ToBytes for ChallengeResponse<N> {
     fn write_le<W: Write>(&self, mut writer: W) -> IoResult<()> {
+        self.restrictions_id.write_le(&mut writer)?;
         self.signature.write_le(&mut writer)?;
         self.nonce.write_le(&mut writer)?;
         Ok(())
@@ -38,10 +40,11 @@ impl<N: Network> ToBytes for ChallengeResponse<N> {
 
 impl<N: Network> FromBytes for ChallengeResponse<N> {
     fn read_le<R: Read>(mut reader: R) -> IoResult<Self> {
+        let restrictions_id = Field::read_le(&mut reader)?;
         let signature = Data::read_le(&mut reader)?;
         let nonce = u64::read_le(&mut reader)?;
 
-        Ok(Self { signature, nonce })
+        Ok(Self { restrictions_id, signature, nonce })
     }
 }
 
@@ -51,7 +54,7 @@ pub mod prop_tests {
     use snarkvm::{
         console::prelude::{FromBytes, ToBytes},
         ledger::narwhal::Data,
-        prelude::{PrivateKey, Signature},
+        prelude::{Field, PrivateKey, Signature},
         utilities::rand::{TestRng, Uniform},
     };
 
@@ -61,6 +64,10 @@ pub mod prop_tests {
 
     type CurrentNetwork = snarkvm::prelude::MainnetV0;
 
+    pub fn any_restrictions_id() -> Field<CurrentNetwork> {
+        Uniform::rand(&mut TestRng::default())
+    }
+
     pub fn any_signature() -> BoxedStrategy<Signature<CurrentNetwork>> {
         (0..64)
             .prop_map(|message_size| {
@@ -74,7 +81,11 @@ pub mod prop_tests {
 
     pub fn any_challenge_response() -> BoxedStrategy<ChallengeResponse<CurrentNetwork>> {
         (any_signature(), any::<u64>())
-            .prop_map(|(sig, nonce)| ChallengeResponse { signature: Data::Object(sig), nonce })
+            .prop_map(|(sig, nonce)| ChallengeResponse {
+                restrictions_id: any_restrictions_id(),
+                signature: Data::Object(sig),
+                nonce,
+            })
             .boxed()
     }
 

node/bft/ledger-service/src/ledger.rs

@@ -82,6 +82,11 @@ impl<N: Network, C: ConsensusStorage<N>> LedgerService<N> for CoreLedgerService<
         self.ledger.latest_block()
     }
 
+    /// Returns the latest restrictions ID in the ledger.
+    fn latest_restrictions_id(&self) -> Field<N> {
+        self.ledger.vm().restrictions().restrictions_id()
+    }
+
     /// Returns the latest cached leader and its associated round.
     fn latest_leader(&self) -> Option<(u64, Address<N>)> {
         *self.latest_leader.read()

node/bft/ledger-service/src/mock.rs

@@ -20,7 +20,7 @@ use snarkvm::{
         narwhal::{BatchCertificate, Data, Subdag, Transmission, TransmissionID},
         puzzle::{Solution, SolutionID},
     },
-    prelude::{bail, ensure, Address, Field, Network, Result},
+    prelude::{bail, ensure, Address, Field, Network, Result, Zero},
 };
 
 use indexmap::IndexMap;
@@ -68,6 +68,11 @@ impl<N: Network> LedgerService<N> for MockLedgerService<N> {
         unreachable!("MockLedgerService does not support latest_block")
     }
 
+    /// Returns the latest restrictions ID in the ledger.
+    fn latest_restrictions_id(&self) -> Field<N> {
+        Field::zero()
+    }
+
     /// Returns the latest cached leader and its associated round.
     fn latest_leader(&self) -> Option<(u64, Address<N>)> {
         None

node/bft/ledger-service/src/prover.rs

@@ -20,7 +20,7 @@ use snarkvm::{
         narwhal::{BatchCertificate, Data, Subdag, Transmission, TransmissionID},
         puzzle::{Solution, SolutionID},
     },
-    prelude::{bail, Address, Field, Network, Result},
+    prelude::{bail, Address, Field, Network, Result, Zero},
 };
 
 use indexmap::IndexMap;
@@ -56,6 +56,11 @@ impl<N: Network> LedgerService<N> for ProverLedgerService<N> {
         unreachable!("Latest block does not exist in prover")
     }
 
+    /// Returns the latest restrictions ID in the ledger.
+    fn latest_restrictions_id(&self) -> Field<N> {
+        Field::zero()
+    }
+
     /// Returns the latest cached leader and its associated round.
     fn latest_leader(&self) -> Option<(u64, Address<N>)> {
         unreachable!("Latest leader does not exist in prover");

node/bft/ledger-service/src/traits.rs

@@ -36,6 +36,9 @@ pub trait LedgerService<N: Network>: Debug + Send + Sync {
     /// Returns the latest block in the ledger.
     fn latest_block(&self) -> Block<N>;
 
+    /// Returns the latest restrictions ID in the ledger.
+    fn latest_restrictions_id(&self) -> Field<N>;
+
     /// Returns the latest cached leader and its associated round.
     fn latest_leader(&self) -> Option<(u64, Address<N>)>;
 

node/bft/ledger-service/src/translucent.rs

@@ -67,6 +67,11 @@ impl<N: Network, C: ConsensusStorage<N>> LedgerService<N> for TranslucentLedgerS
         self.inner.latest_block()
     }
 
+    /// Returns the latest restrictions ID in the ledger.
+    fn latest_restrictions_id(&self) -> Field<N> {
+        self.inner.latest_restrictions_id()
+    }
+
     /// Returns the latest cached leader and its associated round.
     fn latest_leader(&self) -> Option<(u64, Address<N>)> {
         self.inner.latest_leader()

node/bft/src/gateway.rs

@@ -56,7 +56,7 @@ use snarkvm::{
         committee::Committee,
         narwhal::{BatchHeader, Data},
     },
-    prelude::Address,
+    prelude::{Address, Field},
 };
 
 use colored::Colorize;
@@ -1115,11 +1115,14 @@ impl<N: Network> Handshake for Gateway<N> {
             Some(peer_addr)
         };
 
+        // Retrieve the restrictions ID.
+        let restrictions_id = self.ledger.latest_restrictions_id();
+
         // Perform the handshake; we pass on a mutable reference to peer_ip in case the process is broken at any point in time.
         let handshake_result = if peer_side == ConnectionSide::Responder {
-            self.handshake_inner_initiator(peer_addr, peer_ip, stream).await
+            self.handshake_inner_initiator(peer_addr, peer_ip, restrictions_id, stream).await
         } else {
-            self.handshake_inner_responder(peer_addr, &mut peer_ip, stream).await
+            self.handshake_inner_responder(peer_addr, &mut peer_ip, restrictions_id, stream).await
         };
 
         // Remove the address from the collection of connecting peers (if the handshake got to the point where it's known).
@@ -1183,6 +1186,7 @@ impl<N: Network> Gateway<N> {
         &'a self,
         peer_addr: SocketAddr,
         peer_ip: Option<SocketAddr>,
+        restrictions_id: Field<N>,
         stream: &'a mut TcpStream,
     ) -> io::Result<(SocketAddr, Framed<&mut TcpStream, EventCodec<N>>)> {
         // This value is immediately guaranteed to be present, so it can be unwrapped.
@@ -1210,8 +1214,9 @@ impl<N: Network> Gateway<N> {
         let peer_request = expect_event!(Event::ChallengeRequest, framed, peer_addr);
 
         // Verify the challenge response. If a disconnect reason was returned, send the disconnect message and abort.
-        if let Some(reason) =
-            self.verify_challenge_response(peer_addr, peer_request.address, peer_response, our_nonce).await
+        if let Some(reason) = self
+            .verify_challenge_response(peer_addr, peer_request.address, peer_response, restrictions_id, our_nonce)
+            .await
         {
             send_event(&mut framed, peer_addr, reason.into()).await?;
             return Err(error(format!("Dropped '{peer_addr}' for reason: {reason:?}")));
@@ -1231,7 +1236,8 @@ impl<N: Network> Gateway<N> {
             return Err(error(format!("Failed to sign the challenge request nonce from '{peer_addr}'")));
         };
         // Send the challenge response.
-        let our_response = ChallengeResponse { signature: Data::Object(our_signature), nonce: response_nonce };
+        let our_response =
+            ChallengeResponse { restrictions_id, signature: Data::Object(our_signature), nonce: response_nonce };
         send_event(&mut framed, peer_addr, Event::ChallengeResponse(our_response)).await?;
 
         // Add the peer to the gateway.
@@ -1245,6 +1251,7 @@ impl<N: Network> Gateway<N> {
         &'a self,
         peer_addr: SocketAddr,
         peer_ip: &mut Option<SocketAddr>,
+        restrictions_id: Field<N>,
         stream: &'a mut TcpStream,
     ) -> io::Result<(SocketAddr, Framed<&mut TcpStream, EventCodec<N>>)> {
         // Construct the stream.
@@ -1286,7 +1293,8 @@ impl<N: Network> Gateway<N> {
             return Err(error(format!("Failed to sign the challenge request nonce from '{peer_addr}'")));
         };
         // Send the challenge response.
-        let our_response = ChallengeResponse { signature: Data::Object(our_signature), nonce: response_nonce };
+        let our_response =
+            ChallengeResponse { restrictions_id, signature: Data::Object(our_signature), nonce: response_nonce };
         send_event(&mut framed, peer_addr, Event::ChallengeResponse(our_response)).await?;
 
         // Sample a random nonce.
@@ -1300,8 +1308,9 @@ impl<N: Network> Gateway<N> {
         // Listen for the challenge response message.
         let peer_response = expect_event!(Event::ChallengeResponse, framed, peer_addr);
         // Verify the challenge response. If a disconnect reason was returned, send the disconnect message and abort.
-        if let Some(reason) =
-            self.verify_challenge_response(peer_addr, peer_request.address, peer_response, our_nonce).await
+        if let Some(reason) = self
+            .verify_challenge_response(peer_addr, peer_request.address, peer_response, restrictions_id, our_nonce)
+            .await
         {
             send_event(&mut framed, peer_addr, reason.into()).await?;
             return Err(error(format!("Dropped '{peer_addr}' for reason: {reason:?}")));
@@ -1340,10 +1349,17 @@ impl<N: Network> Gateway<N> {
         peer_addr: SocketAddr,
         peer_address: Address<N>,
         response: ChallengeResponse<N>,
+        expected_restrictions_id: Field<N>,
         expected_nonce: u64,
     ) -> Option<DisconnectReason> {
         // Retrieve the components of the challenge response.
-        let ChallengeResponse { signature, nonce } = response;
+        let ChallengeResponse { restrictions_id, signature, nonce } = response;
+
+        // Verify the restrictions ID.
+        if restrictions_id != expected_restrictions_id {
+            warn!("{CONTEXT} Gateway handshake with '{peer_addr}' failed (incorrect restrictions ID)");
+            return Some(DisconnectReason::InvalidChallengeResponse);
+        }
         // Perform the deferred non-blocking deserialization of the signature.
         let Ok(signature) = spawn_blocking!(signature.deserialize_blocking()) else {
             warn!("{CONTEXT} Gateway handshake with '{peer_addr}' failed (cannot deserialize the signature)");

node/bft/src/worker.rs

@@ -526,6 +526,7 @@ mod tests {
             fn latest_round(&self) -> u64;
             fn latest_block_height(&self) -> u32;
             fn latest_block(&self) -> Block<N>;
+            fn latest_restrictions_id(&self) -> Field<N>;
             fn latest_leader(&self) -> Option<(u64, Address<N>)>;
             fn update_latest_leader(&self, round: u64, leader: Address<N>);
             fn contains_block_height(&self, height: u32) -> bool;

node/bft/tests/gateway_e2e.rs

@@ -219,7 +219,7 @@ async fn handshake_responder_side_invalid_challenge_response() {
     let _ = test_peer.unicast(gateway.local_ip(), Event::ChallengeRequest(challenge_request));
 
     // Receive the gateway's challenge response.
-    let (peer_addr, Event::ChallengeResponse(ChallengeResponse { signature, nonce })) =
+    let (peer_addr, Event::ChallengeResponse(ChallengeResponse { restrictions_id, signature, nonce })) =
         test_peer.recv_timeout(Duration::from_secs(1)).await
     else {
         panic!("Expected challenge response")
@@ -251,6 +251,7 @@ async fn handshake_responder_side_invalid_challenge_response() {
     let _ = test_peer.unicast(
         gateway.local_ip(),
         Event::ChallengeResponse(ChallengeResponse {
+            restrictions_id,
             signature: Data::Object(
                 accounts.get(2).unwrap().sign_bytes(&challenge_request.nonce.to_le_bytes(), &mut rng).unwrap(),
             ),

node/router/messages/src/challenge_response.rs

@@ -16,14 +16,15 @@ use super::*;
 
 use snarkvm::{
     ledger::narwhal::Data,
-    prelude::{FromBytes, ToBytes},
+    prelude::{Field, FromBytes, ToBytes},
 };
 
 use std::borrow::Cow;
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct ChallengeResponse<N: Network> {
     pub genesis_header: Header<N>,
+    pub restrictions_id: Field<N>,
     pub signature: Data<Signature<N>>,
     pub nonce: u64,
 }
@@ -39,6 +40,7 @@ impl<N: Network> MessageTrait for ChallengeResponse<N> {
 impl<N: Network> ToBytes for ChallengeResponse<N> {
     fn write_le<W: io::Write>(&self, mut writer: W) -> io::Result<()> {
         self.genesis_header.write_le(&mut writer)?;
+        self.restrictions_id.write_le(&mut writer)?;
         self.signature.write_le(&mut writer)?;
         self.nonce.write_le(&mut writer)
     }
@@ -48,6 +50,7 @@ impl<N: Network> FromBytes for ChallengeResponse<N> {
     fn read_le<R: io::Read>(mut reader: R) -> io::Result<Self> {
         Ok(Self {
             genesis_header: Header::read_le(&mut reader)?,
+            restrictions_id: Field::read_le(&mut reader)?,
             signature: Data::read_le(&mut reader)?,
             nonce: u64::read_le(reader)?,
         })
@@ -60,7 +63,7 @@ pub mod prop_tests {
     use snarkvm::{
         console::prelude::{FromBytes, ToBytes},
         ledger::{ledger_test_helpers::sample_genesis_block, narwhal::Data},
-        prelude::{block::Header, PrivateKey, Signature},
+        prelude::{block::Header, Field, PrivateKey, Signature},
         utilities::rand::{TestRng, Uniform},
     };
 
@@ -70,6 +73,10 @@ pub mod prop_tests {
 
     type CurrentNetwork = snarkvm::prelude::MainnetV0;
 
+    pub fn any_restrictions_id() -> Field<CurrentNetwork> {
+        Uniform::rand(&mut TestRng::default())
+    }
+
     pub fn any_signature() -> BoxedStrategy<Signature<CurrentNetwork>> {
         (0..64)
             .prop_map(|message_size| {
@@ -86,10 +93,11 @@ pub mod prop_tests {
     }
 
     pub fn any_challenge_response() -> BoxedStrategy<ChallengeResponse<CurrentNetwork>> {
-        (any_signature(), any_genesis_header(), any::<u64>())
-            .prop_map(|(sig, genesis_header, nonce)| ChallengeResponse {
-                signature: Data::Object(sig),
+        (any_genesis_header(), any_signature(), any::<u64>())
+            .prop_map(|(genesis_header, sig, nonce)| ChallengeResponse {
                 genesis_header,
+                restrictions_id: any_restrictions_id(),
+                signature: Data::Object(sig),
                 nonce,
             })
             .boxed()