Merge pull request AleoNet#2115 from AleoHQ/feat/auth-handshake

Adds barebones authenticated handshake

node/messages/src/challenge_request.rs

@@ -22,6 +22,7 @@ pub struct ChallengeRequest<N: Network> {
     pub listener_port: u16,
     pub node_type: NodeType,
     pub address: Address<N>,
+    pub nonce: u64,
 }
 
 impl<N: Network> MessageTrait for ChallengeRequest<N> {
@@ -34,13 +35,16 @@ impl<N: Network> MessageTrait for ChallengeRequest<N> {
     /// Serializes the message into the buffer.
     #[inline]
     fn serialize<W: Write>(&self, writer: &mut W) -> Result<()> {
-        Ok(bincode::serialize_into(writer, &(self.version, self.listener_port, self.node_type, self.address))?)
+        Ok(bincode::serialize_into(
+            writer,
+            &(self.version, self.listener_port, self.node_type, self.address, self.nonce),
+        )?)
     }
 
     /// Deserializes the given buffer into a message.
     #[inline]
     fn deserialize(bytes: BytesMut) -> Result<Self> {
-        let (version, listener_port, node_type, address) = bincode::deserialize_from(&mut bytes.reader())?;
-        Ok(Self { version, listener_port, node_type, address })
+        let (version, listener_port, node_type, address, nonce) = bincode::deserialize_from(&mut bytes.reader())?;
+        Ok(Self { version, listener_port, node_type, address, nonce })
     }
 }

node/messages/src/challenge_response.rs

@@ -18,7 +18,8 @@ use super::*;
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct ChallengeResponse<N: Network> {
-    pub header: Data<Header<N>>,
+    pub genesis_header: Header<N>,
+    pub signature: Data<Signature<N>>,
 }
 
 impl<N: Network> MessageTrait for ChallengeResponse<N> {
@@ -31,12 +32,17 @@ impl<N: Network> MessageTrait for ChallengeResponse<N> {
     /// Serializes the message into the buffer.
     #[inline]
     fn serialize<W: Write>(&self, writer: &mut W) -> Result<()> {
-        self.header.serialize_blocking_into(writer)
+        writer.write_all(&self.genesis_header.to_bytes_le()?)?;
+        self.signature.serialize_blocking_into(writer)
     }
 
     /// Deserializes the given buffer into a message.
     #[inline]
     fn deserialize(bytes: BytesMut) -> Result<Self> {
-        Ok(Self { header: Data::Buffer(bytes.freeze()) })
+        let mut reader = bytes.reader();
+        Ok(Self {
+            genesis_header: Header::read_le(&mut reader)?,
+            signature: Data::Buffer(reader.into_inner().freeze()),
+        })
     }
 }

node/messages/src/helpers/noise_codec.rs

@@ -357,6 +357,7 @@ mod tests {
             listener_port: 0,
             node_type: NodeType::Client,
             address: Address::new(Group::rand(rng)),
+            nonce: 0,
         })));
 
         assert_roundtrip(challenge_request);

node/router/Cargo.toml

@@ -60,6 +60,9 @@ version = "0.11"
 [dependencies.serde]
 version = "1"
 
+[dependencies.snarkos-account]
+path = "../../account"
+
 [dependencies.snarkos-node-messages]
 path = "../messages"
 

node/router/src/handshake.rs

@@ -26,10 +26,11 @@ use snarkos_node_messages::{
     MessageTrait,
 };
 use snarkos_node_tcp::{ConnectionSide, Tcp, P2P};
-use snarkvm::prelude::{error, Header, Network};
+use snarkvm::prelude::{error, Address, Header, Network};
 
 use anyhow::{bail, Result};
 use futures::SinkExt;
+use rand::{rngs::OsRng, Rng};
 use std::{io, net::SocketAddr};
 use tokio::net::TcpStream;
 use tokio_stream::StreamExt;
@@ -62,58 +63,73 @@ impl<N: Network> Router<N> {
 
         /* Step 1: Send the challenge request. */
 
+        // Initialize an RNG.
+        let rng = &mut OsRng;
+        // Sample a random nonce.
+        let nonce_a = rng.gen();
+
         // Send a challenge request to the peer.
-        let message = Message::<N>::ChallengeRequest(ChallengeRequest {
+        let message_a = Message::<N>::ChallengeRequest(ChallengeRequest {
             version: Message::<N>::VERSION,
             listener_port: self.local_ip.port(),
             node_type: self.node_type,
-            address: self.address,
+            address: self.address(),
+            nonce: nonce_a,
         });
-        trace!("Sending '{}-A' to '{peer_addr}'", message.name());
-        framed.send(message).await?;
+        trace!("Sending '{}-A' to '{peer_addr}'", message_a.name());
+        framed.send(message_a).await?;
 
         /* Step 2: Receive the challenge request. */
 
         // Listen for the challenge request message.
-        let challenge_request = match framed.try_next().await? {
+        let request_b = match framed.try_next().await? {
             // Received the challenge request message, proceed.
             Some(Message::ChallengeRequest(data)) => data,
             // Received a disconnect message, abort.
             Some(Message::Disconnect(reason)) => return Err(error(format!("'{peer_addr}' disconnected: {reason:?}"))),
             // Received an unexpected message, abort.
             _ => return Err(error(format!("'{peer_addr}' did not send a challenge request"))),
         };
-        trace!("Received '{}-B' from '{peer_addr}'", challenge_request.name());
+        trace!("Received '{}-B' from '{peer_addr}'", request_b.name());
 
         // Verify the challenge request. If a disconnect reason was returned, send the disconnect message and abort.
-        if let Some(reason) = self.verify_challenge_request(peer_addr, &challenge_request) {
+        if let Some(reason) = self.verify_challenge_request(peer_addr, &request_b) {
             trace!("Sending 'Disconnect' to '{peer_addr}'");
             framed.send(Message::Disconnect(Disconnect { reason: reason.clone() })).await?;
             return Err(error(format!("Dropped '{peer_addr}' for reason: {reason:?}")));
         }
 
         /* Step 3: Send the challenge response. */
 
+        // Sign the counterparty nonce.
+        let signature_b = self
+            .account
+            .sign_bytes(&request_b.nonce.to_le_bytes(), rng)
+            .map_err(|_| error(format!("Failed to sign the challenge request nonce from '{peer_addr}'")))?;
+
         // Send the challenge response.
-        let message = Message::ChallengeResponse(ChallengeResponse { header: Data::Object(genesis_header) });
-        trace!("Sending '{}-B' to '{peer_addr}'", message.name());
-        framed.send(message).await?;
+        let message_b =
+            Message::ChallengeResponse(ChallengeResponse { genesis_header, signature: Data::Object(signature_b) });
+        trace!("Sending '{}-B' to '{peer_addr}'", message_b.name());
+        framed.send(message_b).await?;
 
         /* Step 4: Receive the challenge response. */
 
         // Listen for the challenge response message.
-        let challenge_response = match framed.try_next().await? {
+        let response_a = match framed.try_next().await? {
             // Received the challenge response message, proceed.
             Some(Message::ChallengeResponse(data)) => data,
             // Received a disconnect message, abort.
             Some(Message::Disconnect(reason)) => return Err(error(format!("'{peer_addr}' disconnected: {reason:?}"))),
             // Received an unexpected message, abort.
             _ => return Err(error(format!("'{peer_addr}' did not send a challenge response"))),
         };
-        trace!("Received '{}-A' from '{peer_addr}'", challenge_response.name());
+        trace!("Received '{}-A' from '{peer_addr}'", response_a.name());
 
         // 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, challenge_response, genesis_header).await {
+        if let Some(reason) =
+            self.verify_challenge_response(peer_addr, request_b.address, response_a, genesis_header, nonce_a).await
+        {
             trace!("Sending 'Disconnect' to '{peer_addr}'");
             framed.send(Message::Disconnect(Disconnect { reason: reason.clone() })).await?;
             return Err(error(format!("Dropped '{peer_addr}' for reason: {reason:?}")));
@@ -124,13 +140,13 @@ impl<N: Network> Router<N> {
         // Prepare the peer.
         let peer_ip = match peer_side {
             // The peer initiated the connection.
-            ConnectionSide::Initiator => SocketAddr::new(peer_addr.ip(), challenge_request.listener_port),
+            ConnectionSide::Initiator => SocketAddr::new(peer_addr.ip(), request_b.listener_port),
             // This node initiated the connection.
             ConnectionSide::Responder => peer_addr,
         };
-        let peer_address = challenge_request.address;
-        let peer_version = challenge_request.version;
-        let peer_type = challenge_request.node_type;
+        let peer_address = request_b.address;
+        let peer_version = request_b.version;
+        let peer_type = request_b.node_type;
 
         // Construct the peer.
         let peer = Peer::new(peer_ip, peer_address, peer_version, peer_type);
@@ -180,39 +196,56 @@ impl<N: Network> Router<N> {
         message: &ChallengeRequest<N>,
     ) -> Option<DisconnectReason> {
         // Retrieve the components of the challenge request.
-        let &ChallengeRequest { version, listener_port: _, node_type: _, address: _ } = message;
+        let &ChallengeRequest { version, listener_port: _, node_type, address, nonce: _ } = message;
 
         // Ensure the message protocol version is not outdated.
         if version < Message::<N>::VERSION {
             warn!("Dropping '{peer_addr}' on version {version} (outdated)");
             return Some(DisconnectReason::OutdatedClientVersion);
         }
 
+        // TODO (howardwu): Remove this after Phase 2.
+        if !self.is_dev
+            && node_type.is_beacon()
+            && address.to_string() != "aleo1q6qstg8q8shwqf5m6q5fcenuwsdqsvp4hhsgfnx5chzjm3secyzqt9mxm8"
+        {
+            warn!("Dropping '{peer_addr}' for an invalid {node_type}");
+            return Some(DisconnectReason::ProtocolViolation);
+        }
+
         None
     }
 
     /// Verifies the given challenge response. Returns a disconnect reason if the response is invalid.
     async fn verify_challenge_response(
         &self,
         peer_addr: SocketAddr,
-        message: ChallengeResponse<N>,
-        genesis_header: Header<N>,
+        peer_address: Address<N>,
+        response: ChallengeResponse<N>,
+        expected_genesis_header: Header<N>,
+        expected_nonce: u64,
     ) -> Option<DisconnectReason> {
         // Retrieve the components of the challenge response.
-        let ChallengeResponse { header } = message;
+        let ChallengeResponse { genesis_header, signature } = response;
+
+        // Verify the challenge response, by checking that the block header matches.
+        if genesis_header != expected_genesis_header {
+            warn!("Handshake with '{peer_addr}' failed (incorrect block header)");
+            return Some(DisconnectReason::InvalidChallengeResponse);
+        }
 
-        // Perform the deferred non-blocking deserialization of the block header.
-        let block_header = match header.deserialize().await {
-            Ok(block_header) => block_header,
+        // Perform the deferred non-blocking deserialization of the signature.
+        let signature = match signature.deserialize().await {
+            Ok(signature) => signature,
             Err(_) => {
-                warn!("Handshake with '{peer_addr}' failed (cannot deserialize the block header)");
+                warn!("Handshake with '{peer_addr}' failed (cannot deserialize the signature)");
                 return Some(DisconnectReason::InvalidChallengeResponse);
             }
         };
 
-        // Verify the challenge response, by checking that the block header matches.
-        if block_header != genesis_header {
-            warn!("Handshake with '{peer_addr}' failed (incorrect block header)");
+        // Verify the signature.
+        if !signature.verify_bytes(&peer_address, &expected_nonce.to_le_bytes()) {
+            warn!("Handshake with '{peer_addr}' failed (invalid signature)");
             return Some(DisconnectReason::InvalidChallengeResponse);
         }
 

node/router/src/lib.rs

@@ -39,6 +39,7 @@ pub use outbound::*;
 mod routing;
 pub use routing::*;
 
+use snarkos_account::Account;
 use snarkos_node_messages::NodeType;
 use snarkos_node_tcp::{Config, Tcp};
 use snarkvm::prelude::{Address, Network};
@@ -58,8 +59,8 @@ pub struct Router<N: Network> {
     local_ip: SocketAddr,
     /// The node type.
     node_type: NodeType,
-    /// The address of the node.
-    address: Address<N>,
+    /// The account of the node.
+    account: Account<N>,
     /// The cache.
     cache: Cache<N>,
     /// The resolver.
@@ -95,7 +96,7 @@ impl<N: Network> Router<N> {
     pub async fn new(
         node_ip: SocketAddr,
         node_type: NodeType,
-        address: Address<N>,
+        account: Account<N>,
         trusted_peers: &[SocketAddr],
         max_peers: u16,
         is_dev: bool,
@@ -109,7 +110,7 @@ impl<N: Network> Router<N> {
             tcp,
             local_ip,
             node_type,
-            address,
+            account,
             cache: Default::default(),
             resolver: Default::default(),
             sync: Sync::new(local_ip),
@@ -172,7 +173,7 @@ impl<N: Network> Router<N> {
 
     /// Returns the Aleo address of the node.
     pub const fn address(&self) -> Address<N> {
-        self.address
+        self.account.address()
     }
 
     /// Returns the sync pool.

node/src/beacon/mod.rs

@@ -122,7 +122,7 @@ impl<N: Network, C: ConsensusStorage<N>> Beacon<N, C> {
         let router = Router::new(
             node_ip,
             NodeType::Beacon,
-            account.address(),
+            account.clone(),
             trusted_peers,
             Self::MAXIMUM_NUMBER_OF_PEERS as u16,
             dev.is_some(),

node/src/client/mod.rs

@@ -74,7 +74,7 @@ impl<N: Network, C: ConsensusStorage<N>> Client<N, C> {
         let router = Router::new(
             node_ip,
             NodeType::Client,
-            account.address(),
+            account.clone(),
             trusted_peers,
             Self::MAXIMUM_NUMBER_OF_PEERS as u16,
             dev.is_some(),

node/src/prover/mod.rs

@@ -92,7 +92,7 @@ impl<N: Network, C: ConsensusStorage<N>> Prover<N, C> {
         let router = Router::new(
             node_ip,
             NodeType::Prover,
-            account.address(),
+            account.clone(),
             trusted_peers,
             Self::MAXIMUM_NUMBER_OF_PEERS as u16,
             dev.is_some(),

node/src/validator/mod.rs

@@ -100,7 +100,7 @@ impl<N: Network, C: ConsensusStorage<N>> Validator<N, C> {
         let router = Router::new(
             node_ip,
             NodeType::Validator,
-            account.address(),
+            account.clone(),
             trusted_peers,
             Self::MAXIMUM_NUMBER_OF_PEERS as u16,
             dev.is_some(),