Move reth_eth_wire::ProtocolVersion and reth_eth_wire::Capabilities to reth-eth-wire-types (#10071)

Co-authored-by: Emilia Hane <emiliaha95@gmail.com>

crates/net/eth-wire/src/capability.rs

@@ -5,18 +5,15 @@ use crate::{
     p2pstream::MAX_RESERVED_MESSAGE_ID,
     protocol::{ProtoVersion, Protocol},
     version::ParseVersionError,
-    EthMessage, EthMessageID, EthVersion,
+    Capability, EthMessage, EthMessageID, EthVersion,
 };
-use alloy_rlp::{Decodable, Encodable, RlpDecodable, RlpEncodable};
 use derive_more::{Deref, DerefMut};
-use reth_codecs::add_arbitrary_tests;
-use reth_primitives::bytes::{BufMut, Bytes};
+use reth_primitives::bytes::Bytes;
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
 use std::{
     borrow::Cow,
     collections::{BTreeSet, HashMap},
-    fmt,
 };
 
 /// A Capability message consisting of the message-id and the payload
@@ -40,173 +37,6 @@ pub enum CapabilityMessage {
     Other(RawCapabilityMessage),
 }
 
-/// A message indicating a supported capability and capability version.
-#[add_arbitrary_tests(rlp)]
-#[derive(Clone, Debug, PartialEq, Eq, RlpEncodable, RlpDecodable, Default, Hash)]
-#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
-pub struct Capability {
-    /// The name of the subprotocol
-    pub name: Cow<'static, str>,
-    /// The version of the subprotocol
-    pub version: usize,
-}
-
-impl Capability {
-    /// Create a new `Capability` with the given name and version.
-    pub const fn new(name: String, version: usize) -> Self {
-        Self { name: Cow::Owned(name), version }
-    }
-
-    /// Create a new `Capability` with the given static name and version.
-    pub const fn new_static(name: &'static str, version: usize) -> Self {
-        Self { name: Cow::Borrowed(name), version }
-    }
-
-    /// Returns the corresponding eth capability for the given version.
-    pub const fn eth(version: EthVersion) -> Self {
-        Self::new_static("eth", version as usize)
-    }
-
-    /// Returns the [`EthVersion::Eth66`] capability.
-    pub const fn eth_66() -> Self {
-        Self::eth(EthVersion::Eth66)
-    }
-
-    /// Returns the [`EthVersion::Eth67`] capability.
-    pub const fn eth_67() -> Self {
-        Self::eth(EthVersion::Eth67)
-    }
-
-    /// Returns the [`EthVersion::Eth68`] capability.
-    pub const fn eth_68() -> Self {
-        Self::eth(EthVersion::Eth68)
-    }
-
-    /// Whether this is eth v66 protocol.
-    #[inline]
-    pub fn is_eth_v66(&self) -> bool {
-        self.name == "eth" && self.version == 66
-    }
-
-    /// Whether this is eth v67.
-    #[inline]
-    pub fn is_eth_v67(&self) -> bool {
-        self.name == "eth" && self.version == 67
-    }
-
-    /// Whether this is eth v68.
-    #[inline]
-    pub fn is_eth_v68(&self) -> bool {
-        self.name == "eth" && self.version == 68
-    }
-
-    /// Whether this is any eth version.
-    #[inline]
-    pub fn is_eth(&self) -> bool {
-        self.is_eth_v66() || self.is_eth_v67() || self.is_eth_v68()
-    }
-}
-
-impl fmt::Display for Capability {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}/{}", self.name, self.version)
-    }
-}
-
-impl From<EthVersion> for Capability {
-    #[inline]
-    fn from(value: EthVersion) -> Self {
-        Self::eth(value)
-    }
-}
-
-#[cfg(any(test, feature = "arbitrary"))]
-impl<'a> arbitrary::Arbitrary<'a> for Capability {
-    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
-        let version = u.int_in_range(0..=32)?; // TODO: What's the max?
-        let name = String::arbitrary(u)?; // TODO: what possible values?
-        Ok(Self::new(name, version))
-    }
-}
-
-/// Represents all capabilities of a node.
-#[derive(Debug, Clone, Eq, PartialEq)]
-pub struct Capabilities {
-    /// All Capabilities and their versions
-    inner: Vec<Capability>,
-    eth_66: bool,
-    eth_67: bool,
-    eth_68: bool,
-}
-
-impl Capabilities {
-    /// Returns all capabilities.
-    #[inline]
-    pub fn capabilities(&self) -> &[Capability] {
-        &self.inner
-    }
-
-    /// Consumes the type and returns the all capabilities.
-    #[inline]
-    pub fn into_inner(self) -> Vec<Capability> {
-        self.inner
-    }
-
-    /// Whether the peer supports `eth` sub-protocol.
-    #[inline]
-    pub const fn supports_eth(&self) -> bool {
-        self.eth_68 || self.eth_67 || self.eth_66
-    }
-
-    /// Whether this peer supports eth v66 protocol.
-    #[inline]
-    pub const fn supports_eth_v66(&self) -> bool {
-        self.eth_66
-    }
-
-    /// Whether this peer supports eth v67 protocol.
-    #[inline]
-    pub const fn supports_eth_v67(&self) -> bool {
-        self.eth_67
-    }
-
-    /// Whether this peer supports eth v68 protocol.
-    #[inline]
-    pub const fn supports_eth_v68(&self) -> bool {
-        self.eth_68
-    }
-}
-
-impl From<Vec<Capability>> for Capabilities {
-    fn from(value: Vec<Capability>) -> Self {
-        Self {
-            eth_66: value.iter().any(Capability::is_eth_v66),
-            eth_67: value.iter().any(Capability::is_eth_v67),
-            eth_68: value.iter().any(Capability::is_eth_v68),
-            inner: value,
-        }
-    }
-}
-
-impl Encodable for Capabilities {
-    fn encode(&self, out: &mut dyn BufMut) {
-        self.inner.encode(out)
-    }
-}
-
-impl Decodable for Capabilities {
-    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
-        let inner = Vec::<Capability>::decode(buf)?;
-
-        Ok(Self {
-            eth_66: inner.iter().any(Capability::is_eth_v66),
-            eth_67: inner.iter().any(Capability::is_eth_v67),
-            eth_68: inner.iter().any(Capability::is_eth_v68),
-            inner,
-        })
-    }
-}
-
 /// This represents a shared capability, its version, and its message id offset.
 ///
 /// The [offset](SharedCapability::message_id_offset) is the message ID offset for this shared
@@ -549,6 +379,7 @@ pub struct UnsupportedCapabilityError {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::{Capabilities, Capability};
 
     #[test]
     fn from_eth_68() {

crates/net/eth-wire/src/ethstream.rs

@@ -349,8 +349,9 @@ mod tests {
         broadcast::BlockHashNumber,
         errors::{EthHandshakeError, EthStreamError},
         hello::DEFAULT_TCP_PORT,
-        p2pstream::{ProtocolVersion, UnauthedP2PStream},
-        EthMessage, EthStream, EthVersion, HelloMessageWithProtocols, PassthroughCodec, Status,
+        p2pstream::UnauthedP2PStream,
+        EthMessage, EthStream, EthVersion, HelloMessageWithProtocols, PassthroughCodec,
+        ProtocolVersion, Status,
     };
     use futures::{SinkExt, StreamExt};
     use reth_chainspec::NamedChain;

crates/net/eth-wire/src/hello.rs

@@ -1,4 +1,4 @@
-use crate::{capability::Capability, EthVersion, ProtocolVersion};
+use crate::{Capability, EthVersion, ProtocolVersion};
 use alloy_rlp::{RlpDecodable, RlpEncodable};
 use reth_codecs::derive_arbitrary;
 use reth_network_peers::PeerId;
@@ -214,9 +214,7 @@ impl HelloMessageBuilder {
 
 #[cfg(test)]
 mod tests {
-    use crate::{
-        capability::Capability, p2pstream::P2PMessage, EthVersion, HelloMessage, ProtocolVersion,
-    };
+    use crate::{p2pstream::P2PMessage, Capability, EthVersion, HelloMessage, ProtocolVersion};
     use alloy_rlp::{Decodable, Encodable, EMPTY_STRING_CODE};
     use reth_network_peers::pk2id;
     use secp256k1::{SecretKey, SECP256K1};

crates/net/eth-wire/src/lib.rs

@@ -32,14 +32,14 @@ pub use tokio_util::codec::{
 };
 
 pub use crate::{
-    capability::Capability,
     disconnect::CanDisconnect,
     ethstream::{EthStream, UnauthedEthStream, MAX_MESSAGE_SIZE},
     hello::{HelloMessage, HelloMessageBuilder, HelloMessageWithProtocols},
     p2pstream::{
-        DisconnectP2P, P2PMessage, P2PMessageID, P2PStream, ProtocolVersion, UnauthedP2PStream,
+        DisconnectP2P, P2PMessage, P2PMessageID, P2PStream, UnauthedP2PStream,
         MAX_RESERVED_MESSAGE_ID,
     },
+    Capability, ProtocolVersion,
 };
 
 // Re-export wire types

crates/net/eth-wire/src/multiplex.rs

@@ -17,10 +17,10 @@ use std::{
 };
 
 use crate::{
-    capability::{Capability, SharedCapabilities, SharedCapability, UnsupportedCapabilityError},
+    capability::{SharedCapabilities, SharedCapability, UnsupportedCapabilityError},
     errors::{EthStreamError, P2PStreamError},
     p2pstream::DisconnectP2P,
-    CanDisconnect, DisconnectReason, EthStream, P2PStream, Status, UnauthedEthStream,
+    CanDisconnect, Capability, DisconnectReason, EthStream, P2PStream, Status, UnauthedEthStream,
 };
 use bytes::{Bytes, BytesMut};
 use futures::{Sink, SinkExt, Stream, StreamExt, TryStream, TryStreamExt};

crates/net/eth-wire/src/p2pstream.rs

@@ -16,7 +16,7 @@ use reth_primitives::{
 };
 use std::{
     collections::VecDeque,
-    fmt, io,
+    io,
     pin::Pin,
     task::{ready, Context, Poll},
     time::Duration,
@@ -788,49 +788,10 @@ impl TryFrom<u8> for P2PMessageID {
     }
 }
 
-/// RLPx `p2p` protocol version
-#[derive_arbitrary(rlp)]
-#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
-#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
-pub enum ProtocolVersion {
-    /// `p2p` version 4
-    V4 = 4,
-    /// `p2p` version 5
-    #[default]
-    V5 = 5,
-}
-
-impl fmt::Display for ProtocolVersion {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "v{}", *self as u8)
-    }
-}
-
-impl Encodable for ProtocolVersion {
-    fn encode(&self, out: &mut dyn BufMut) {
-        (*self as u8).encode(out)
-    }
-    fn length(&self) -> usize {
-        // the version should be a single byte
-        (*self as u8).length()
-    }
-}
-
-impl Decodable for ProtocolVersion {
-    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
-        let version = u8::decode(buf)?;
-        match version {
-            4 => Ok(Self::V4),
-            5 => Ok(Self::V5),
-            _ => Err(RlpError::Custom("unknown p2p protocol version")),
-        }
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{capability::SharedCapability, test_utils::eth_hello, EthVersion};
+    use crate::{capability::SharedCapability, test_utils::eth_hello, EthVersion, ProtocolVersion};
     use tokio::net::{TcpListener, TcpStream};
     use tokio_util::codec::Decoder;
 

crates/net/eth-wire/src/protocol.rs

@@ -1,6 +1,6 @@
 //! A Protocol defines a P2P subprotocol in a `RLPx` connection
 
-use crate::{capability::Capability, EthMessageID, EthVersion};
+use crate::{Capability, EthMessageID, EthVersion};
 
 /// Type that represents a [Capability] and the number of messages it uses.
 ///

crates/net/eth-wire/src/test_utils.rs

@@ -60,7 +60,7 @@ pub async fn connect_passthrough(
 /// A Rplx subprotocol for testing
 pub mod proto {
     use super::*;
-    use crate::{capability::Capability, protocol::Protocol};
+    use crate::{protocol::Protocol, Capability};
     use bytes::{Buf, BufMut, BytesMut};
 
     /// Returns a new testing `HelloMessage` with eth and the test protocol