Repository: tangle-network/gadget (language: rust)

Issue:
[TASK] Add signature verification during the peer handshake
Add Signature here:
https://github.com/webb-tools/gadget/blob/281959b8b2c2bffeead19808d8ed328d2f0d0921/shell/src/network/handlers/p2p.rs#L104-L117

And verification here:
https://github.com/webb-tools/gadget/blob/281959b8b2c2bffeead19808d8ed328d2f0d0921/shell/src/network/handlers/p2p.rs#L81-L89


Discussion / hints:

Fix patch:
diff --git a/Cargo.lock b/Cargo.lock
index 1b7e4ed2a..6262835fb 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4109,6 +4109,7 @@ dependencies = [
  "sha2 0.10.8",
  "sp-application-crypto",
  "sp-core",
+ "sp-io",
  "sp-keystore",
  "structopt",
  "tangle-subxt",
diff --git a/shell/Cargo.toml b/shell/Cargo.toml
index bc8a58dbf..18c1c96b1 100644
--- a/shell/Cargo.toml
+++ b/shell/Cargo.toml
@@ -19,6 +19,7 @@ gadget-core = { workspace = true }
 gadget-common = { workspace = true }
 
 sp-core = { workspace = true, features = ["std"] }
+sp-io = { workspace = true, features = ["std"] }
 sp-keystore = { workspace = true }
 sp-application-crypto = { workspace = true }
 
diff --git a/shell/src/network/gossip.rs b/shell/src/network/gossip.rs
index 7d5b060f0..3c67a4d30 100644
--- a/shell/src/network/gossip.rs
+++ b/shell/src/network/gossip.rs
@@ -38,7 +38,7 @@ pub struct NetworkServiceWithoutSwarm<'a> {
     pub logger: &'a DebugLogger,
     pub inbound_mapping: &'a [InboundMapping],
     pub ecdsa_peer_id_to_libp2p_id: Arc<RwLock<HashMap<ecdsa::Public, PeerId>>>,
-    pub role_key: &'a ecdsa::Public,
+    pub role_key: &'a ecdsa::Pair,
     pub span: tracing::Span,
 }
 
@@ -63,7 +63,7 @@ pub struct NetworkService<'a> {
     pub logger: &'a DebugLogger,
     pub inbound_mapping: &'a [InboundMapping],
     pub ecdsa_peer_id_to_libp2p_id: &'a Arc<RwLock<HashMap<ecdsa::Public, PeerId>>>,
-    pub role_key: &'a ecdsa::Public,
+    pub role_key: &'a ecdsa::Pair,
     pub span: &'a tracing::Span,
 }
 
@@ -290,13 +290,22 @@ pub struct GossipMessage {
 
 #[derive(Serialize, Deserialize, Debug)]
 pub enum MyBehaviourRequest {
-    Handshake { ecdsa_public_key: ecdsa::Public },
-    Message { topic: String, raw_payload: Vec<u8> },
+    Handshake {
+        ecdsa_public_key: ecdsa::Public,
+        signature: ecdsa::Signature,
+    },
+    Message {
+        topic: String,
+        raw_payload: Vec<u8>,
+    },
 }
 
 #[derive(Serialize, Deserialize, Debug)]
 pub enum MyBehaviourResponse {
-    Handshaked { ecdsa_public_key: ecdsa::Public },
+    Handshaked {
+        ecdsa_public_key: ecdsa::Public,
+        signature: ecdsa::Signature,
+    },
     MessageHandled,
 }
 
@@ -379,7 +388,7 @@ mod tests {
     use crate::shell::wait_for_connection_to_bootnodes;
     use gadget_common::prelude::{DebugLogger, GadgetProtocolMessage, Network, WorkManager};
     use gadget_core::job_manager::WorkManagerInterface;
-    use sp_core::ecdsa;
+    use sp_core::{ecdsa, Pair};
 
     #[tokio::test]
     async fn test_gossip_network() {
@@ -518,7 +527,7 @@ mod tests {
                         assert_eq!(message.payload, b"Hello, world");
                         assert_eq!(
                             message.to_network_id,
-                            Some(get_dummy_role_key_from_index(my_idx))
+                            Some(get_dummy_role_key_from_index(my_idx).public())
                         );
                     }
                 }
@@ -537,7 +546,7 @@ mod tests {
         to_idx: usize,
     ) -> <WorkManager as WorkManagerInterface>::ProtocolMessage {
         let dummy_role_key = get_dummy_role_key_from_index(to_idx);
-        dummy_message_inner(input, Some(dummy_role_key))
+        dummy_message_inner(input, Some(dummy_role_key.public()))
     }
 
     fn dummy_message_inner(
@@ -557,7 +566,7 @@ mod tests {
         }
     }
 
-    fn get_dummy_role_key_from_index(index: usize) -> ecdsa::Public {
-        ecdsa::Public::from_raw([index as u8; 33])
+    fn get_dummy_role_key_from_index(index: usize) -> ecdsa::Pair {
+        ecdsa::Pair::from_seed(&[index as u8; 32])
     }
 }
diff --git a/shell/src/network/handlers/connections.rs b/shell/src/network/handlers/connections.rs
index 8379c5dc8..7d6b14f29 100644
--- a/shell/src/network/handlers/connections.rs
+++ b/shell/src/network/handlers/connections.rs
@@ -1,40 +1,49 @@
 use crate::network::gossip::{MyBehaviourRequest, NetworkService};
 use libp2p::PeerId;
+use sp_core::{keccak_256, Pair};
 
 impl NetworkService<'_> {
     #[tracing::instrument(skip(self))]
     pub(crate) async fn handle_connection_established(
         &mut self,
         peer_id: PeerId,
-        _num_established: u32,
+        num_established: u32,
     ) {
         self.logger.debug("Connection established");
-        // TODO: Sign the handshake message.
-        let handshake = MyBehaviourRequest::Handshake {
-            ecdsa_public_key: *self.role_key,
-        };
-        self.swarm
-            .behaviour_mut()
-            .p2p
-            .send_request(&peer_id, handshake);
-        self.swarm
-            .behaviour_mut()
-            .gossipsub
-            .add_explicit_peer(&peer_id);
+        if num_established == 1 {
+            let my_peer_id = self.swarm.local_peer_id();
+            let msg = my_peer_id.to_bytes();
+            let hash = keccak_256(&msg);
+            let signature = self.role_key.sign_prehashed(&hash);
+            let handshake = MyBehaviourRequest::Handshake {
+                ecdsa_public_key: self.role_key.public(),
+                signature,
+            };
+            self.swarm
+                .behaviour_mut()
+                .p2p
+                .send_request(&peer_id, handshake);
+            self.swarm
+                .behaviour_mut()
+                .gossipsub
+                .add_explicit_peer(&peer_id);
+        }
     }
 
     #[tracing::instrument(skip(self))]
     pub(crate) async fn handle_connection_closed(
         &mut self,
         peer_id: PeerId,
-        _num_established: u32,
+        num_established: u32,
         _cause: Option<libp2p::swarm::ConnectionError>,
     ) {
         self.logger.debug("Connection closed");
-        self.swarm
-            .behaviour_mut()
-            .gossipsub
-            .remove_explicit_peer(&peer_id);
+        if num_established == 0 {
+            self.swarm
+                .behaviour_mut()
+                .gossipsub
+                .remove_explicit_peer(&peer_id);
+        }
     }
 
     #[tracing::instrument(skip(self))]
diff --git a/shell/src/network/handlers/p2p.rs b/shell/src/network/handlers/p2p.rs
index df43f1b2f..29b1a6ce4 100644
--- a/shell/src/network/handlers/p2p.rs
+++ b/shell/src/network/handlers/p2p.rs
@@ -2,6 +2,8 @@ use crate::network::gossip::{MyBehaviourRequest, MyBehaviourResponse, NetworkSer
 use libp2p::gossipsub::IdentTopic;
 
 use libp2p::{request_response, PeerId};
+use sp_core::{keccak_256, Pair};
+use sp_io::crypto::ecdsa_verify_prehashed;
 
 impl NetworkService<'_> {
     #[tracing::instrument(skip(self, event))]
@@ -78,10 +80,24 @@ impl NetworkService<'_> {
     ) {
         use crate::network::gossip::MyBehaviourResponse::*;
         match message {
-            Handshaked { ecdsa_public_key } => {
-                // Their response to our handshake request.
-                // we should add them to our mapping.
-                // TODO: Add signature verification here.
+            Handshaked {
+                ecdsa_public_key,
+                signature,
+            } => {
+                let msg = peer.to_bytes();
+                let hash = keccak_256(&msg);
+                let valid = ecdsa_verify_prehashed(&signature, &hash, &ecdsa_public_key);
+                if !valid {
+                    self.logger
+                        .warn(format!("Invalid signature from peer: {peer}"));
+                    // TODO: report this peer.
+                    self.ecdsa_peer_id_to_libp2p_id
+                        .write()
+                        .await
+                        .remove(&ecdsa_public_key);
+                    let _ = self.swarm.disconnect_peer_id(peer);
+                    return;
+                }
                 self.ecdsa_peer_id_to_libp2p_id
                     .write()
                     .await
@@ -101,17 +117,36 @@ impl NetworkService<'_> {
     ) {
         use crate::network::gossip::MyBehaviourRequest::*;
         let result = match req {
-            Handshake { ecdsa_public_key } => {
+            Handshake {
+                ecdsa_public_key,
+                signature,
+            } => {
                 self.logger
                     .debug(format!("Received handshake from peer: {peer}"));
+                // Verify the signature
+                let msg = peer.to_bytes();
+                let hash = keccak_256(&msg);
+                let valid = ecdsa_verify_prehashed(&signature, &hash, &ecdsa_public_key);
+                if !valid {
+                    self.logger
+                        .warn(format!("Invalid signature from peer: {peer}"));
+                    let _ = self.swarm.disconnect_peer_id(peer);
+                    return;
+                }
                 self.ecdsa_peer_id_to_libp2p_id
                     .write()
                     .await
                     .insert(ecdsa_public_key, peer);
+                // Send response with our public key
+                let my_peer_id = self.swarm.local_peer_id();
+                let msg = my_peer_id.to_bytes();
+                let hash = keccak_256(&msg);
+                let signature = self.role_key.sign_prehashed(&hash);
                 self.swarm.behaviour_mut().p2p.send_response(
                     channel,
                     MyBehaviourResponse::Handshaked {
-                        ecdsa_public_key: *self.role_key,
+                        ecdsa_public_key: self.role_key.public(),
+                        signature,
                     },
                 )
             }
diff --git a/shell/src/network/setup.rs b/shell/src/network/setup.rs
index f999cc6fa..fb59434f2 100644
--- a/shell/src/network/setup.rs
+++ b/shell/src/network/setup.rs
@@ -25,7 +25,7 @@ pub async fn setup_libp2p_network(
     config: &ShellConfig,
     logger: DebugLogger,
     networks: Vec<&'static str>,
-    role_key: ecdsa::Public,
+    role_key: ecdsa::Pair,
 ) -> Result<(Vec<GossipHandle>, JoinHandle<()>), Box<dyn Error>> {
     // Setup both QUIC (UDP) and TCP transports the increase the chances of NAT traversal
     let mut swarm = libp2p::SwarmBuilder::with_existing_identity(identity)
diff --git a/shell/src/shell.rs b/shell/src/shell.rs
index d3943f238..3e62f9c10 100644
--- a/shell/src/shell.rs
+++ b/shell/src/shell.rs
@@ -55,7 +55,7 @@ pub async fn run_forever(config: ShellConfig) -> color_eyre::Result<()> {
             DFNS_CGGMP21_KEYREFRESH_PROTOCOL_NAME,
             DFNS_CGGMP21_KEYROTATE_PROTOCOL_NAME,
         ],
-        role_key.public(),
+        role_key,
     )
     .await
     .map_err(|e| color_eyre::eyre::eyre!("Failed to setup network: {e}"))?;


Test patch: