From b033b929d0a685dba1d0fa8ef7ce17b3b2ccd622 Mon Sep 17 00:00:00 2001 From: Age Manning Date: Tue, 29 Oct 2024 14:40:47 +1100 Subject: [PATCH] Improves ergonomics to allow removing fields from an ENR (#84) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: João Oliveira --- src/builder.rs | 22 ++-- src/lib.rs | 265 ++++++++++++++++++++++++++++++++----------------- 2 files changed, 189 insertions(+), 98 deletions(-) diff --git a/src/builder.rs b/src/builder.rs index 3b7628b..935df24 100644 --- a/src/builder.rs +++ b/src/builder.rs @@ -1,4 +1,8 @@ use crate::{Enr, EnrKey, EnrPublicKey, Error, Key, NodeId, MAX_ENR_SIZE}; +use crate::{ + ENR_VERSION, ID_ENR_KEY, IP6_ENR_KEY, IP_ENR_KEY, TCP6_ENR_KEY, TCP_ENR_KEY, UDP6_ENR_KEY, + UDP_ENR_KEY, +}; use alloy_rlp::{Encodable, Header}; use bytes::{Bytes, BytesMut}; use std::{ @@ -27,7 +31,7 @@ impl Default for Builder { /// Constructs a minimal [`Builder`] for the v4 identity scheme. fn default() -> Self { Self { - id: String::from("v4"), + id: String::from_utf8(ENR_VERSION.into()).expect("Is a valid string"), seq: 1, content: BTreeMap::new(), phantom: PhantomData, @@ -65,13 +69,13 @@ impl Builder { /// Adds an `ip` field to the `ENRBuilder`. pub fn ip4(&mut self, ip: Ipv4Addr) -> &mut Self { - self.add_value("ip", &ip.octets().as_ref()); + self.add_value(IP_ENR_KEY, &ip.octets().as_ref()); self } /// Adds an `ip6` field to the `ENRBuilder`. pub fn ip6(&mut self, ip: Ipv6Addr) -> &mut Self { - self.add_value("ip6", &ip.octets().as_ref()); + self.add_value(IP6_ENR_KEY, &ip.octets().as_ref()); self } @@ -81,32 +85,32 @@ impl Builder { /// Adds an `Id` field to the `ENRBuilder`. pub fn id(&mut self, id: &str) -> &mut Self { - self.add_value("id", &id.as_bytes()); + self.add_value(ID_ENR_KEY, &id.as_bytes()); self } */ /// Adds a `tcp` field to the `ENRBuilder`. pub fn tcp4(&mut self, tcp: u16) -> &mut Self { - self.add_value("tcp", &tcp); + self.add_value(TCP_ENR_KEY, &tcp); self } /// Adds a `tcp6` field to the `ENRBuilder`. pub fn tcp6(&mut self, tcp: u16) -> &mut Self { - self.add_value("tcp6", &tcp); + self.add_value(TCP6_ENR_KEY, &tcp); self } /// Adds a `udp` field to the `ENRBuilder`. pub fn udp4(&mut self, udp: u16) -> &mut Self { - self.add_value("udp", &udp); + self.add_value(UDP_ENR_KEY, &udp); self } /// Adds a `udp6` field to the `ENRBuilder`. pub fn udp6(&mut self, udp: u16) -> &mut Self { - self.add_value("udp6", &udp); + self.add_value(UDP6_ENR_KEY, &udp); self } @@ -179,7 +183,7 @@ impl Builder { let mut id_bytes = BytesMut::with_capacity(self.id.length()); self.id.as_bytes().encode(&mut id_bytes); - self.add_value_rlp("id", id_bytes.freeze()); + self.add_value_rlp(ID_ENR_KEY, id_bytes.freeze()); self.add_public_key(&key.public()); let rlp_content = self.rlp_content(); diff --git a/src/lib.rs b/src/lib.rs index 66b2002..7e4761e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -221,6 +221,15 @@ type PreviousRlpEncodedValues = Vec>; const MAX_ENR_SIZE: usize = 300; +const ID_ENR_KEY: &[u8] = b"id"; +const ENR_VERSION: &[u8] = b"v4"; +pub const IP_ENR_KEY: &[u8] = b"ip"; +pub const IP6_ENR_KEY: &[u8] = b"ip6"; +pub const TCP_ENR_KEY: &[u8] = b"tcp"; +pub const TCP6_ENR_KEY: &[u8] = b"tcp6"; +pub const UDP_ENR_KEY: &[u8] = b"udp"; +pub const UDP6_ENR_KEY: &[u8] = b"udp6"; + /// The ENR, allowing for arbitrary signing algorithms. /// /// This struct will always have a valid signature, known public key type, sequence number and `NodeId`. All other parameters are variable/optional. @@ -304,7 +313,7 @@ impl Enr { /// Returns the IPv4 address of the ENR record if it is defined. #[must_use] pub fn ip4(&self) -> Option { - if let Some(Ok(ip_bytes)) = self.get_decodable::("ip") { + if let Some(Ok(ip_bytes)) = self.get_decodable::(IP_ENR_KEY) { return match ip_bytes.as_ref().len() { 4 => { let mut ip = [0_u8; 4]; @@ -320,7 +329,7 @@ impl Enr { /// Returns the IPv6 address of the ENR record if it is defined. #[must_use] pub fn ip6(&self) -> Option { - if let Some(Ok(ip_bytes)) = self.get_decodable::("ip6") { + if let Some(Ok(ip_bytes)) = self.get_decodable::(IP6_ENR_KEY) { return match ip_bytes.as_ref().len() { 16 => { let mut ip = [0_u8; 16]; @@ -375,25 +384,25 @@ impl Enr { /// The TCP port of ENR record if it is defined. #[must_use] pub fn tcp4(&self) -> Option { - self.get_decodable("tcp").and_then(Result::ok) + self.get_decodable(TCP_ENR_KEY).and_then(Result::ok) } /// The IPv6-specific TCP port of ENR record if it is defined. #[must_use] pub fn tcp6(&self) -> Option { - self.get_decodable("tcp6").and_then(Result::ok) + self.get_decodable(TCP6_ENR_KEY).and_then(Result::ok) } /// The UDP port of ENR record if it is defined. #[must_use] pub fn udp4(&self) -> Option { - self.get_decodable("udp").and_then(Result::ok) + self.get_decodable(UDP_ENR_KEY).and_then(Result::ok) } /// The IPv6-specific UDP port of ENR record if it is defined. #[must_use] pub fn udp6(&self) -> Option { - self.get_decodable("udp6").and_then(Result::ok) + self.get_decodable(UDP6_ENR_KEY).and_then(Result::ok) } /// Provides a socket (based on the UDP port), if the IPv4 and UDP fields are specified. @@ -593,7 +602,7 @@ impl Enr { pub fn set_ip(&mut self, ip: IpAddr, key: &K) -> Result, Error> { match ip { IpAddr::V4(addr) => { - let prev_value = self.insert("ip", &addr.octets().as_ref(), key)?; + let prev_value = self.insert(IP_ENR_KEY, &addr.octets().as_ref(), key)?; if let Some(bytes) = prev_value { if bytes.len() == 4 { let mut v = [0_u8; 4]; @@ -603,7 +612,7 @@ impl Enr { } } IpAddr::V6(addr) => { - let prev_value = self.insert("ip6", &addr.octets().as_ref(), key)?; + let prev_value = self.insert(IP6_ENR_KEY, &addr.octets().as_ref(), key)?; if let Some(bytes) = prev_value { if bytes.len() == 16 { let mut v = [0_u8; 16]; @@ -619,36 +628,56 @@ impl Enr { /// Sets the `udp` field of the ENR. Returns any pre-existing UDP port in the record. pub fn set_udp4(&mut self, udp: u16, key: &K) -> Result, Error> { - if let Some(udp_bytes) = self.insert("udp", &udp, key)? { + if let Some(udp_bytes) = self.insert(UDP_ENR_KEY, &udp, key)? { return Ok(u16::decode(&mut udp_bytes.as_ref()).ok()); } Ok(None) } + /// Unsets the udp field on the ENR. + pub fn remove_udp4(&mut self, key: &K) -> Result<(), Error> { + self.remove_key(UDP_ENR_KEY, key) + } + /// Sets the `udp6` field of the ENR. Returns any pre-existing UDP port in the record. pub fn set_udp6(&mut self, udp: u16, key: &K) -> Result, Error> { - if let Some(udp_bytes) = self.insert("udp6", &udp, key)? { + if let Some(udp_bytes) = self.insert(UDP6_ENR_KEY, &udp, key)? { return Ok(u16::decode(&mut udp_bytes.as_ref()).ok()); } Ok(None) } + /// Unsets the udp6 field on the ENR. + pub fn remove_udp6(&mut self, key: &K) -> Result<(), Error> { + self.remove_key(UDP6_ENR_KEY, key) + } + /// Sets the `tcp` field of the ENR. Returns any pre-existing tcp port in the record. pub fn set_tcp4(&mut self, tcp: u16, key: &K) -> Result, Error> { - if let Some(tcp_bytes) = self.insert("tcp", &tcp, key)? { + if let Some(tcp_bytes) = self.insert(TCP_ENR_KEY, &tcp, key)? { return Ok(u16::decode(&mut tcp_bytes.as_ref()).ok()); } Ok(None) } + /// Unsets the `tcp` field on the ENR. + pub fn remove_tcp(&mut self, key: &K) -> Result<(), Error> { + self.remove_key(TCP_ENR_KEY, key) + } + /// Sets the `tcp6` field of the ENR. Returns any pre-existing tcp6 port in the record. pub fn set_tcp6(&mut self, tcp: u16, key: &K) -> Result, Error> { - if let Some(tcp_bytes) = self.insert("tcp6", &tcp, key)? { + if let Some(tcp_bytes) = self.insert(TCP6_ENR_KEY, &tcp, key)? { return Ok(u16::decode(&mut tcp_bytes.as_ref()).ok()); } Ok(None) } + /// Unsets the `tcp6` field on the ENR. + pub fn remove_tcp6(&mut self, key: &K) -> Result<(), Error> { + self.remove_key(TCP6_ENR_KEY, key) + } + /// Sets the [EIP-7636](https://eips.ethereum.org/EIPS/eip-7636) `client` field in the record. pub fn set_client_info( &mut self, @@ -671,98 +700,142 @@ impl Enr { self.set_socket(socket, key, false) } + /// Unsets the `ip` and `udp` fields on the ENR. + pub fn remove_udp_socket(&mut self, key: &K) -> Result<(), Error> { + let keys_to_remove = [IP_ENR_KEY, UDP_ENR_KEY].iter(); + let keys_to_insert = std::iter::empty::<(Vec, &[u8])>(); + self.remove_insert(keys_to_remove, keys_to_insert, key) + .map(|_| ()) + } + + /// Unsets the `ip6` and `udp6` fields on the ENR. + pub fn remove_udp6_socket(&mut self, key: &K) -> Result<(), Error> { + let keys_to_remove = [IP6_ENR_KEY, UDP6_ENR_KEY].iter(); + let keys_to_insert = std::iter::empty::<(Vec, &[u8])>(); + self.remove_insert(keys_to_remove, keys_to_insert, key) + .map(|_| ()) + } + /// Sets the IP and TCP port in a single update with a single increment in sequence number. pub fn set_tcp_socket(&mut self, socket: SocketAddr, key: &K) -> Result<(), Error> { self.set_socket(socket, key, true) } + /// Unsets the `ip` and `tcp` fields on the ENR. + pub fn remove_tcp_socket(&mut self, key: &K) -> Result<(), Error> { + let keys_to_remove = [IP_ENR_KEY, TCP_ENR_KEY].iter(); + let keys_to_insert = std::iter::empty::<(Vec, &[u8])>(); + self.remove_insert(keys_to_remove, keys_to_insert, key) + .map(|_| ()) + } + + /// Unsets the `ip6` and `tcp6` fields on the ENR. + pub fn remove_tcp6_socket(&mut self, key: &K) -> Result<(), Error> { + let keys_to_remove = [IP6_ENR_KEY, TCP6_ENR_KEY].iter(); + let keys_to_insert = std::iter::empty::<(Vec, &[u8])>(); + self.remove_insert(keys_to_remove, keys_to_insert, key) + .map(|_| ()) + } + /// Helper function for `set_tcp_socket()` and `set_udp_socket`. fn set_socket(&mut self, socket: SocketAddr, key: &K, is_tcp: bool) -> Result<(), Error> { + // We work on this new version, allowing us to not mutate self on error. + let mut new_enr = self.clone(); let (port_string, port_v6_string): (Key, Key) = if is_tcp { - ("tcp".into(), "tcp6".into()) + (TCP_ENR_KEY.into(), TCP6_ENR_KEY.into()) } else { - ("udp".into(), "udp6".into()) + (UDP_ENR_KEY.into(), UDP6_ENR_KEY.into()) }; - let (prev_ip, prev_port) = match socket.ip() { + match socket.ip() { IpAddr::V4(addr) => { let mut ip = BytesMut::new(); addr.encode(&mut ip); let mut port = BytesMut::new(); socket.port().encode(&mut port); - ( - self.content.insert("ip".into(), ip.freeze()), - self.content.insert(port_string.clone(), port.freeze()), - ) + new_enr.content.insert(IP_ENR_KEY.into(), ip.freeze()); + new_enr.content.insert(port_string.clone(), port.freeze()); } IpAddr::V6(addr) => { let mut ip6 = BytesMut::new(); addr.encode(&mut ip6); let mut port = BytesMut::new(); socket.port().encode(&mut port); - ( - self.content.insert("ip6".into(), ip6.freeze()), - self.content.insert(port_v6_string.clone(), port.freeze()), - ) + new_enr.content.insert(IP6_ENR_KEY.into(), ip6.freeze()); + new_enr + .content + .insert(port_v6_string.clone(), port.freeze()); } }; let public_key = key.public(); let mut pubkey = BytesMut::new(); public_key.encode().as_ref().encode(&mut pubkey); - let previous_key = self.content.insert(public_key.enr_key(), pubkey.freeze()); + new_enr + .content + .insert(public_key.enr_key(), pubkey.freeze()); - // check the size and revert on failure - if self.size() > MAX_ENR_SIZE { - // if the size of the record is too large, revert and error - // revert the public key - if let Some(key) = previous_key { - self.content.insert(public_key.enr_key(), key); - } else { - self.content.remove(&public_key.enr_key()); - } - // revert the content - match socket.ip() { - IpAddr::V4(_) => { - if let Some(ip) = prev_ip { - self.content.insert("ip".into(), ip); - } else { - self.content.remove(b"ip".as_ref()); - } - if let Some(udp) = prev_port { - self.content.insert(port_string, udp); - } else { - self.content.remove(&port_string); - } - } - IpAddr::V6(_) => { - if let Some(ip) = prev_ip { - self.content.insert("ip6".into(), ip); - } else { - self.content.remove(b"ip6".as_ref()); - } - if let Some(udp) = prev_port { - self.content.insert(port_v6_string, udp); - } else { - self.content.remove(&port_v6_string); - } - } - } + // check the size + if new_enr.size() > MAX_ENR_SIZE { return Err(Error::ExceedsMaxSize); } // increment the sequence number - self.seq = self + new_enr.seq = new_enr .seq .checked_add(1) .ok_or(Error::SequenceNumberTooHigh)?; // sign the record - self.sign(key)?; + new_enr.sign(key)?; // update the node id - self.node_id = NodeId::from(key.public()); + new_enr.node_id = NodeId::from(key.public()); + + if new_enr.size() > MAX_ENR_SIZE { + // in case the signature size changes, inform the user the size has exceeded the + // maximum + return Err(Error::ExceedsMaxSize); + } + // Everything passed update the record + *self = new_enr; + + Ok(()) + } + + /// Removes a key from the ENR. + pub fn remove_key(&mut self, content_key: impl AsRef<[u8]>, enr_key: &K) -> Result<(), Error> { + // We work on this new version, allowing us to not mutate self on error. + let mut new_enr = self.clone(); + new_enr.content.remove(content_key.as_ref()); + + // add the new public key. + let public_key = enr_key.public(); + let mut pubkey = BytesMut::new(); + public_key.encode().as_ref().encode(&mut pubkey); + new_enr + .content + .insert(public_key.enr_key(), pubkey.freeze()); + + // increment the sequence number + new_enr.seq = new_enr + .seq + .checked_add(1) + .ok_or(Error::SequenceNumberTooHigh)?; + + // sign the record + new_enr.sign(enr_key)?; + // update the node id + new_enr.node_id = NodeId::from(enr_key.public()); + + if new_enr.size() > MAX_ENR_SIZE { + // in case the signature size changes, inform the user the size has exceeded the + // maximum + return Err(Error::ExceedsMaxSize); + } + + *self = new_enr; Ok(()) } @@ -778,24 +851,26 @@ impl Enr { insert_key_values: impl Iterator, &'a [u8])>, enr_key: &K, ) -> Result<(PreviousRlpEncodedValues, PreviousRlpEncodedValues), Error> { - let enr_backup = self.clone(); + // We work on this new version, allowing us to not mutate self on error. + let mut new_enr = self.clone(); let mut removed = Vec::new(); for key in remove_keys { - removed.push(self.content.remove(key.as_ref())); + removed.push(new_enr.content.remove(key.as_ref())); } // add the new public key let public_key = enr_key.public(); let mut pubkey = BytesMut::new(); public_key.encode().as_ref().encode(&mut pubkey); - self.content.insert(public_key.enr_key(), pubkey.freeze()); + new_enr + .content + .insert(public_key.enr_key(), pubkey.freeze()); let mut inserted = Vec::new(); for (key, value) in insert_key_values { // currently only support "v4" identity schemes - if key.as_ref() == b"id" && value != b"v4" { - *self = enr_backup; + if key.as_ref() == ID_ENR_KEY && value != ENR_VERSION { return Err(Error::UnsupportedIdentityScheme); } let mut out = BytesMut::new(); @@ -806,27 +881,27 @@ impl Enr { u16::decode(&mut value.as_ref())?; } - inserted.push(self.content.insert(key.as_ref().to_vec(), value)); + inserted.push(new_enr.content.insert(key.as_ref().to_vec(), value)); } // increment the sequence number - self.seq = self + new_enr.seq = new_enr .seq .checked_add(1) .ok_or(Error::SequenceNumberTooHigh)?; // sign the record - self.sign(enr_key)?; + new_enr.sign(enr_key)?; // update the node id - self.node_id = NodeId::from(enr_key.public()); + new_enr.node_id = NodeId::from(enr_key.public()); - if self.size() > MAX_ENR_SIZE { + if new_enr.size() > MAX_ENR_SIZE { // in case the signature size changes, inform the user the size has exceeded the // maximum - *self = enr_backup; return Err(Error::ExceedsMaxSize); } + *self = new_enr; Ok((removed, inserted)) } @@ -883,7 +958,7 @@ impl Enr { /// Compute the enr's signature with the given key. fn compute_signature(&self, signing_key: &K) -> Result, Error> { match self.id() { - Some(ref id) if id == "v4" => signing_key + Some(ref id) if id.as_bytes() == ENR_VERSION => signing_key .sign_v4(&self.rlp_content()) .map_err(|_| Error::SigningError), // other identity schemes are unsupported @@ -950,9 +1025,17 @@ impl std::fmt::Debug for Enr { .iter() .filter(|(key, _)| { // skip all pairs already covered as fields - !["id", "ip", "ip6", "udp", "udp6", "tcp", "tcp6"] - .iter() - .any(|k| k.as_bytes() == key.as_slice()) + ![ + ID_ENR_KEY, + IP_ENR_KEY, + IP6_ENR_KEY, + UDP_ENR_KEY, + UDP6_ENR_KEY, + TCP_ENR_KEY, + TCP6_ENR_KEY, + ] + .iter() + .any(|&k| k == key.as_slice()) }) .map(|(key, val)| (String::from_utf8_lossy(key), hex::encode(val))), ) @@ -1068,15 +1151,15 @@ impl Decodable for Enr { } alloy_rlp::encode(id) } - b"tcp" | b"tcp6" | b"udp" | b"udp6" => { + TCP_ENR_KEY | TCP6_ENR_KEY | UDP_ENR_KEY | UDP6_ENR_KEY => { let port = u16::decode(payload)?; alloy_rlp::encode(port) } - b"ip" => { + IP_ENR_KEY => { let ip = Ipv4Addr::decode(payload)?; alloy_rlp::encode(ip) } - b"ip6" => { + IP6_ENR_KEY => { let ip6 = Ipv6Addr::decode(payload)?; alloy_rlp::encode(ip6) } @@ -1162,12 +1245,12 @@ pub(crate) fn digest(b: &[u8]) -> [u8; 32] { } const fn is_keyof_u16(key: &[u8]) -> bool { - matches!(key, b"tcp" | b"tcp6" | b"udp" | b"udp6") + matches!(key, TCP_ENR_KEY | TCP6_ENR_KEY | UDP_ENR_KEY | UDP6_ENR_KEY) } fn check_spec_reserved_keys(key: &[u8], mut value: &[u8]) -> Result<(), Error> { match key { - b"tcp" | b"tcp6" | b"udp" | b"udp6" => { + TCP_ENR_KEY | TCP6_ENR_KEY | UDP_ENR_KEY | UDP6_ENR_KEY => { u16::decode(&mut value)?; } b"id" => { @@ -1176,10 +1259,10 @@ fn check_spec_reserved_keys(key: &[u8], mut value: &[u8]) -> Result<(), Error> { return Err(Error::UnsupportedIdentityScheme); } } - b"ip" => { + IP_ENR_KEY => { Ipv4Addr::decode(&mut value)?; } - b"ip6" => { + IP6_ENR_KEY => { Ipv6Addr::decode(&mut value)?; } b"secp256k1" => { @@ -1822,7 +1905,11 @@ mod tests { let topics: &[u8] = &out; let (removed, inserted) = enr - .remove_insert([b"tcp"].iter(), vec![(b"topics", topics)].into_iter(), &key) + .remove_insert( + [TCP_ENR_KEY].iter(), + vec![(b"topics", topics)].into_iter(), + &key, + ) .unwrap(); let mut buffer = Vec::::new(); tcp.encode(&mut buffer); @@ -1886,7 +1973,7 @@ mod tests { for tcp in LOW_INT_PORTS { let mut enr = Enr::empty(&key).unwrap(); - let res = enr.insert(b"tcp", &tcp.to_be_bytes().as_ref(), &key); + let res = enr.insert(TCP_ENR_KEY, &tcp.to_be_bytes().as_ref(), &key); if u8::try_from(tcp).is_ok() { assert_eq!(res.unwrap_err().to_string(), "invalid rlp data"); } else { @@ -1905,7 +1992,7 @@ mod tests { let res = enr.remove_insert( [b"none"].iter(), - vec![(b"tcp".as_slice(), tcp.to_be_bytes().as_slice())].into_iter(), + vec![(TCP_ENR_KEY, tcp.to_be_bytes().as_slice())].into_iter(), &key, ); if u8::try_from(tcp).is_ok() { @@ -1964,7 +2051,7 @@ mod tests { let mut buffer = Vec::::new(); tcp.encode(&mut buffer); assert!(enr.verify()); - assert_eq!(enr.get_raw_rlp("tcp").unwrap(), buffer); + assert_eq!(enr.get_raw_rlp(TCP_ENR_KEY).unwrap(), buffer); assert_eq!(enr.tcp4(), Some(tcp)); }