From 5b6343703965aa6a91ba9eab7f25d99ed9bdff28 Mon Sep 17 00:00:00 2001
From: Andrew Kirillov <20803092+akirillo@users.noreply.github.com>
Date: Thu, 12 Jan 2023 02:09:21 -0800
Subject: [PATCH] chore(net): `set_capability_offsets` tests and refactors
 (#763)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>
---
 crates/net/eth-wire/src/capability.rs |  2 +-
 crates/net/eth-wire/src/p2pstream.rs  | 69 +++++++++++++++++++++++----
 2 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/crates/net/eth-wire/src/capability.rs b/crates/net/eth-wire/src/capability.rs
index 1b6dd50c0c..3463892dc8 100644
--- a/crates/net/eth-wire/src/capability.rs
+++ b/crates/net/eth-wire/src/capability.rs
@@ -27,7 +27,7 @@ pub enum CapabilityMessage {
 
 /// A message indicating a supported capability and capability version.
 #[derive(
-    Clone, Debug, PartialEq, Eq, RlpEncodable, RlpDecodable, Serialize, Deserialize, Default,
+    Clone, Debug, PartialEq, Eq, RlpEncodable, RlpDecodable, Serialize, Deserialize, Default, Hash,
 )]
 pub struct Capability {
     /// The name of the subprotocol
diff --git a/crates/net/eth-wire/src/p2pstream.rs b/crates/net/eth-wire/src/p2pstream.rs
index 6819490602..a7fff8ca01 100644
--- a/crates/net/eth-wire/src/p2pstream.rs
+++ b/crates/net/eth-wire/src/p2pstream.rs
@@ -12,7 +12,7 @@ use pin_project::pin_project;
 use reth_rlp::{Decodable, DecodeError, Encodable, EMPTY_LIST_CODE};
 use serde::{Deserialize, Serialize};
 use std::{
-    collections::{BTreeSet, HashMap, VecDeque},
+    collections::{BTreeSet, HashMap, HashSet, VecDeque},
     io,
     pin::Pin,
     task::{ready, Context, Poll},
@@ -493,13 +493,14 @@ where
 /// capabilities and the input list of locally supported capabilities.
 ///
 /// Currently only `eth` versions 66 and 67 are supported.
+/// Additionally, the `p2p` capability version 5 is supported, but is
+/// expected _not_ to be in neither `local_capabilities` or `peer_capabilities`.
 pub fn set_capability_offsets(
     local_capabilities: Vec<Capability>,
     peer_capabilities: Vec<Capability>,
 ) -> Result<SharedCapability, P2PStreamError> {
     // find intersection of capabilities
-    let our_capabilities_map =
-        local_capabilities.into_iter().map(|c| (c.name, c.version)).collect::<HashMap<_, _>>();
+    let our_capabilities = local_capabilities.into_iter().collect::<HashSet<_>>();
 
     // map of capability name to version
     let mut shared_capabilities = HashMap::new();
@@ -519,14 +520,18 @@ pub fn set_capability_offsets(
     let mut shared_capability_names = BTreeSet::new();
 
     // find highest shared version of each shared capability
-    for capability in peer_capabilities {
+    for peer_capability in peer_capabilities {
         // if this is Some, we share this capability
-        if let Some(version) = our_capabilities_map.get(&capability.name) {
+        if our_capabilities.contains(&peer_capability) {
             // If multiple versions are shared of the same (equal name) capability, the numerically
             // highest wins, others are ignored
-            if capability.version <= *version {
-                shared_capabilities.insert(capability.name.clone(), capability.version);
-                shared_capability_names.insert(capability.name);
+
+            let version = shared_capabilities.get(&peer_capability.name);
+            if version.is_none() ||
+                (version.is_some() && peer_capability.version > *version.expect("is some; qed"))
+            {
+                shared_capabilities.insert(peer_capability.name.clone(), peer_capability.version);
+                shared_capability_names.insert(peer_capability.name);
             }
         }
     }
@@ -557,10 +562,10 @@ pub fn set_capability_offsets(
                 tracing::warn!("unknown capability: name={:?}, version={}", name, version,);
             }
             SharedCapability::Eth { .. } => {
-                shared_with_offsets.push(shared_capability.clone());
-
                 // increment the offset if the capability is known
                 offset += shared_capability.num_messages()?;
+
+                shared_with_offsets.push(shared_capability);
             }
         }
     }
@@ -851,6 +856,50 @@ mod tests {
         handle.await.unwrap();
     }
 
+    #[test]
+    fn test_peer_lower_capability_version() {
+        let local_capabilities: Vec<Capability> =
+            vec![EthVersion::Eth66.into(), EthVersion::Eth67.into()];
+        let peer_capabilities: Vec<Capability> = vec![EthVersion::Eth66.into()];
+
+        let shared_capability =
+            set_capability_offsets(local_capabilities, peer_capabilities).unwrap();
+
+        assert_eq!(
+            shared_capability,
+            SharedCapability::Eth {
+                version: EthVersion::Eth66,
+                offset: MAX_RESERVED_MESSAGE_ID + 1
+            }
+        )
+    }
+
+    #[test]
+    fn test_peer_capability_version_too_low() {
+        let local_capabilities: Vec<Capability> = vec![EthVersion::Eth67.into()];
+        let peer_capabilities: Vec<Capability> = vec![EthVersion::Eth66.into()];
+
+        let shared_capability = set_capability_offsets(local_capabilities, peer_capabilities);
+
+        assert!(matches!(
+            shared_capability,
+            Err(P2PStreamError::HandshakeError(P2PHandshakeError::NoSharedCapabilities))
+        ))
+    }
+
+    #[test]
+    fn test_peer_capability_version_too_high() {
+        let local_capabilities: Vec<Capability> = vec![EthVersion::Eth66.into()];
+        let peer_capabilities: Vec<Capability> = vec![EthVersion::Eth67.into()];
+
+        let shared_capability = set_capability_offsets(local_capabilities, peer_capabilities);
+
+        assert!(matches!(
+            shared_capability,
+            Err(P2PStreamError::HandshakeError(P2PHandshakeError::NoSharedCapabilities))
+        ))
+    }
+
     #[test]
     fn snappy_decode_encode_ping() {
         let snappy_ping = b"\x02\x01\0\xc0";