feat: support any node format in admin API (#7247)

Cargo.lock

@@ -6291,6 +6291,7 @@ dependencies = [
  "clap",
  "criterion",
  "derive_more",
+ "enr",
  "hash-db",
  "itertools 0.12.1",
  "modular-bitfield",
@@ -6311,6 +6312,7 @@ dependencies = [
  "secp256k1 0.27.0",
  "serde",
  "serde_json",
+ "serde_with",
  "sha2",
  "strum 0.26.2",
  "sucds",

Cargo.toml

@@ -233,6 +233,7 @@ tracing-appender = "0.2"
 thiserror = "1.0"
 serde_json = "1.0.94"
 serde = { version = "1.0", default-features = false }
+serde_with = "3.3.0"
 humantime-serde = "1.1"
 rand = "0.8.5"
 schnellru = "0.2"

crates/net/nat/Cargo.toml

@@ -21,7 +21,7 @@ tracing.workspace = true
 pin-project-lite = "0.2.9"
 tokio = { workspace = true, features = ["time"] }
 thiserror.workspace = true
-serde_with = { version = "3.3.0", optional = true }
+serde_with = { workspace = true, optional = true }
 
 [dev-dependencies]
 reth-tracing.workspace = true

crates/primitives/Cargo.toml

@@ -27,9 +27,10 @@ alloy-trie = { workspace = true, features = ["serde"] }
 nybbles = { workspace = true, features = ["serde", "rlp"] }
 alloy-genesis.workspace = true
 alloy-eips.workspace = true
+enr = { workspace = true, features = ["rust-secp256k1"] }
+
 # crypto
 secp256k1 = { workspace = true, features = ["global-context", "recovery"] }
-
 # for eip-4844
 c-kzg = { workspace = true, features = ["serde"], optional = true }
 
@@ -42,6 +43,7 @@ itertools.workspace = true
 modular-bitfield = "0.11.2"
 once_cell.workspace = true
 rayon.workspace = true
+serde_with.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 sha2 = { version = "0.10.7", optional = true }

crates/primitives/src/lib.rs

@@ -72,7 +72,7 @@ pub use net::{
     goerli_nodes, holesky_nodes, mainnet_nodes, parse_nodes, sepolia_nodes, NodeRecord,
     GOERLI_BOOTNODES, HOLESKY_BOOTNODES, MAINNET_BOOTNODES, SEPOLIA_BOOTNODES,
 };
-pub use peer::{PeerId, WithPeerId};
+pub use peer::{AnyNode, PeerId, WithPeerId};
 pub use prune::{
     PruneCheckpoint, PruneMode, PruneModes, PruneProgress, PrunePurpose, PruneSegment,
     PruneSegmentError, ReceiptsLogPruneConfig, MINIMUM_PRUNING_DISTANCE,

crates/primitives/src/peer.rs

@@ -1,6 +1,99 @@
+use enr::Enr;
+use reth_rpc_types::NodeRecord;
+use secp256k1::SecretKey;
+use std::{net::IpAddr, str::FromStr};
 // Re-export PeerId for ease of use.
 pub use reth_rpc_types::PeerId;
 
+/// A peer that can come in ENR or [NodeRecord] form.
+#[derive(
+    Debug, Clone, Eq, PartialEq, Hash, serde_with::SerializeDisplay, serde_with::DeserializeFromStr,
+)]
+pub enum AnyNode {
+    /// An "enode:" peer with full ip
+    NodeRecord(NodeRecord),
+    /// An "enr:"
+    Enr(Enr<SecretKey>),
+    /// An incomplete "enode" with only a peer id
+    PeerId(PeerId),
+}
+
+impl AnyNode {
+    /// Returns the peer id of the node.
+    pub fn peer_id(&self) -> PeerId {
+        match self {
+            AnyNode::NodeRecord(record) => record.id,
+            AnyNode::Enr(enr) => {
+                PeerId::from_slice(&enr.public_key().serialize_uncompressed()[1..])
+            }
+            AnyNode::PeerId(peer_id) => *peer_id,
+        }
+    }
+
+    /// Returns the full node record if available.
+    pub fn node_record(&self) -> Option<NodeRecord> {
+        match self {
+            AnyNode::NodeRecord(record) => Some(*record),
+            AnyNode::Enr(enr) => {
+                let node_record = NodeRecord {
+                    address: enr.ip4().map(IpAddr::from).or_else(|| enr.ip6().map(IpAddr::from))?,
+                    tcp_port: enr.tcp4().or_else(|| enr.tcp6())?,
+                    udp_port: enr.udp4().or_else(|| enr.udp6())?,
+                    id: PeerId::from_slice(&enr.public_key().serialize_uncompressed()[1..]),
+                }
+                .into_ipv4_mapped();
+                Some(node_record)
+            }
+            _ => None,
+        }
+    }
+}
+
+impl From<NodeRecord> for AnyNode {
+    fn from(value: NodeRecord) -> Self {
+        Self::NodeRecord(value)
+    }
+}
+
+impl From<Enr<SecretKey>> for AnyNode {
+    fn from(value: Enr<SecretKey>) -> Self {
+        Self::Enr(value)
+    }
+}
+
+impl FromStr for AnyNode {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if let Some(rem) = s.strip_prefix("enode://") {
+            if let Ok(record) = NodeRecord::from_str(s) {
+                return Ok(AnyNode::NodeRecord(record));
+            }
+            // incomplete enode
+            if let Ok(peer_id) = PeerId::from_str(rem) {
+                return Ok(AnyNode::PeerId(peer_id));
+            }
+            return Err(format!("invalid public key: {rem}"));
+        }
+        if s.starts_with("enr:") {
+            return Enr::from_str(s).map(AnyNode::Enr)
+        }
+        Err("missing 'enr:' prefix for base64-encoded record".to_string())
+    }
+}
+
+impl std::fmt::Display for AnyNode {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            AnyNode::NodeRecord(record) => write!(f, "{record}"),
+            AnyNode::Enr(enr) => write!(f, "{enr}"),
+            AnyNode::PeerId(peer_id) => {
+                write!(f, "enode://{}", crate::hex::encode(peer_id.as_slice()))
+            }
+        }
+    }
+}
+
 /// Generic wrapper with peer id
 #[derive(Debug)]
 pub struct WithPeerId<T>(PeerId, pub T);
@@ -54,3 +147,47 @@ impl<T> WithPeerId<Option<T>> {
         self.1.map(|v| WithPeerId(self.0, v))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_node_record_parse() {
+        let url = "enode://6f8a80d14311c39f35f516fa664deaaaa13e85b2f7493f37f6144d86991ec012937307647bd3b9a82abe2974e1407241d54947bbb39763a4cac9f77166ad92a0@10.3.58.6:30303?discport=30301";
+        let node: AnyNode = url.parse().unwrap();
+        assert_eq!(node, AnyNode::NodeRecord(NodeRecord {
+            address: IpAddr::V4([10,3,58,6].into()),
+            tcp_port: 30303,
+            udp_port: 30301,
+            id: "6f8a80d14311c39f35f516fa664deaaaa13e85b2f7493f37f6144d86991ec012937307647bd3b9a82abe2974e1407241d54947bbb39763a4cac9f77166ad92a0".parse().unwrap(),
+        }));
+        assert_eq!(node.to_string(), url)
+    }
+
+    #[test]
+    fn test_peer_id_parse() {
+        let url = "enode://6f8a80d14311c39f35f516fa664deaaaa13e85b2f7493f37f6144d86991ec012937307647bd3b9a82abe2974e1407241d54947bbb39763a4cac9f77166ad92a0";
+        let node: AnyNode = url.parse().unwrap();
+        assert_eq!(node, AnyNode::PeerId("6f8a80d14311c39f35f516fa664deaaaa13e85b2f7493f37f6144d86991ec012937307647bd3b9a82abe2974e1407241d54947bbb39763a4cac9f77166ad92a0".parse().unwrap()));
+        assert_eq!(node.to_string(), url);
+
+        let url = "enode://";
+        let err = url.parse::<AnyNode>().unwrap_err();
+        assert_eq!(err, "invalid public key: ");
+    }
+
+    // <https://eips.ethereum.org/EIPS/eip-778>
+    #[test]
+    fn test_enr_parse() {
+        let url = "enr:-IS4QHCYrYZbAKWCBRlAy5zzaDZXJBGkcnh4MHcBFZntXNFrdvJjX04jRzjzCBOonrkTfj499SZuOh8R33Ls8RRcy5wBgmlkgnY0gmlwhH8AAAGJc2VjcDI1NmsxoQPKY0yuDUmstAHYpMa2_oxVtw0RW_QAdpzBQA8yWM0xOIN1ZHCCdl8";
+        let node: AnyNode = url.parse().unwrap();
+        assert_eq!(
+            node.peer_id(),
+            "0xca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd31387574077f301b421bc84df7266c44e9e6d569fc56be00812904767bf5ccd1fc7f"
+                .parse::<PeerId>()
+                .unwrap()
+        );
+        assert_eq!(node.to_string(), url);
+    }
+}

crates/rpc/rpc-api/src/admin.rs

@@ -1,5 +1,5 @@
 use jsonrpsee::{core::RpcResult, proc_macros::rpc};
-use reth_primitives::NodeRecord;
+use reth_primitives::{AnyNode, NodeRecord};
 use reth_rpc_types::{NodeInfo, PeerInfo};
 
 /// Admin namespace rpc interface that gives access to several non-standard RPC methods.
@@ -14,18 +14,18 @@ pub trait AdminApi {
     ///
     /// Returns true if the peer was successfully removed.
     #[method(name = "removePeer")]
-    fn remove_peer(&self, record: NodeRecord) -> RpcResult<bool>;
+    fn remove_peer(&self, record: AnyNode) -> RpcResult<bool>;
 
     /// Adds the given node record to the trusted peerset.
     #[method(name = "addTrustedPeer")]
-    fn add_trusted_peer(&self, record: NodeRecord) -> RpcResult<bool>;
+    fn add_trusted_peer(&self, record: AnyNode) -> RpcResult<bool>;
 
     /// Removes a remote node from the trusted peer set, but it does not disconnect it
     /// automatically.
     ///
     /// Returns true if the peer was successfully removed.
     #[method(name = "removeTrustedPeer")]
-    fn remove_trusted_peer(&self, record: NodeRecord) -> RpcResult<bool>;
+    fn remove_trusted_peer(&self, record: AnyNode) -> RpcResult<bool>;
 
     /// The peers administrative property can be queried for all the information known about the
     /// connected remote nodes at the networking granularity. These include general information

crates/rpc/rpc-builder/tests/it/http.rs

@@ -141,9 +141,9 @@ where
     let node: NodeRecord = url.parse().unwrap();
 
     AdminApiClient::add_peer(client, node).await.unwrap();
-    AdminApiClient::remove_peer(client, node).await.unwrap();
-    AdminApiClient::add_trusted_peer(client, node).await.unwrap();
-    AdminApiClient::remove_trusted_peer(client, node).await.unwrap();
+    AdminApiClient::remove_peer(client, node.into()).await.unwrap();
+    AdminApiClient::add_trusted_peer(client, node.into()).await.unwrap();
+    AdminApiClient::remove_trusted_peer(client, node.into()).await.unwrap();
     AdminApiClient::node_info(client).await.unwrap();
 }
 

crates/rpc/rpc/src/admin.rs

@@ -2,7 +2,7 @@ use crate::result::ToRpcResult;
 use async_trait::async_trait;
 use jsonrpsee::core::RpcResult;
 use reth_network_api::{NetworkInfo, PeerKind, Peers};
-use reth_primitives::{ChainSpec, NodeRecord};
+use reth_primitives::{AnyNode, ChainSpec, NodeRecord};
 use reth_rpc_api::AdminApiServer;
 use reth_rpc_types::{NodeInfo, PeerEthProtocolInfo, PeerInfo, PeerNetworkInfo, PeerProtocolsInfo};
 use std::sync::Arc;
@@ -36,20 +36,23 @@ where
     }
 
     /// Handler for `admin_removePeer`
-    fn remove_peer(&self, record: NodeRecord) -> RpcResult<bool> {
-        self.network.remove_peer(record.id, PeerKind::Basic);
+    fn remove_peer(&self, record: AnyNode) -> RpcResult<bool> {
+        self.network.remove_peer(record.peer_id(), PeerKind::Basic);
         Ok(true)
     }
 
     /// Handler for `admin_addTrustedPeer`
-    fn add_trusted_peer(&self, record: NodeRecord) -> RpcResult<bool> {
-        self.network.add_trusted_peer(record.id, record.tcp_addr());
+    fn add_trusted_peer(&self, record: AnyNode) -> RpcResult<bool> {
+        if let Some(record) = record.node_record() {
+            self.network.add_trusted_peer(record.id, record.tcp_addr())
+        }
+        self.network.add_trusted_peer_id(record.peer_id());
         Ok(true)
     }
 
     /// Handler for `admin_removeTrustedPeer`
-    fn remove_trusted_peer(&self, record: NodeRecord) -> RpcResult<bool> {
-        self.network.remove_peer(record.id, PeerKind::Trusted);
+    fn remove_trusted_peer(&self, record: AnyNode) -> RpcResult<bool> {
+        self.network.remove_peer(record.peer_id(), PeerKind::Trusted);
         Ok(true)
     }