Prevent ENR from being corruptable

sigp · Oct 28, 2024 · 6eeb694 · 6eeb694
1 parent 8963f58
commit 6eeb694
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 14 deletions.
diff --git a/src/builder.rs b/src/builder.rs
@@ -31,7 +31,7 @@ impl<K: EnrKey> Default for Builder<K> {
     /// Constructs a minimal [`Builder`] for the v4 identity scheme.
     fn default() -> Self {
         Self {
-            id: String::from_utf8_unchecked(ENR_VERSION.into()),
+            id: String::from_utf8(ENR_VERSION.into()).expect("Is a valid string"),
             seq: 1,
             content: BTreeMap::new(),
             phantom: PhantomData,

diff --git a/src/lib.rs b/src/lib.rs
@@ -739,6 +739,7 @@ impl<K: EnrKey> Enr<K> {
 
     /// Helper function for `set_tcp_socket()` and `set_udp_socket`.
     fn set_socket(&mut self, socket: SocketAddr, key: &K, is_tcp: bool) -> Result<(), Error> {
+        let enr_backup = self.clone();
         let (port_string, port_v6_string): (Key, Key) = if is_tcp {
             (TCP_ENR_KEY.into(), TCP6_ENR_KEY.into())
         } else {
@@ -813,10 +814,13 @@ impl<K: EnrKey> Enr<K> {
         }
 
         // increment the sequence number
-        self.seq = self
-            .seq
-            .checked_add(1)
-            .ok_or(Error::SequenceNumberTooHigh)?;
+        match self.seq.checked_add(1) {
+            Some(seq_no) => self.seq = seq_no,
+            None => {
+                *self = enr_backup;
+                return Err(Error::SequenceNumberTooHigh);
+            }
+        }
 
         // sign the record
         self.sign(key)?;
@@ -829,6 +833,7 @@ impl<K: EnrKey> Enr<K> {
 
     /// Removes a key from the ENR.
     pub fn remove_key(&mut self, content_key: impl AsRef<[u8]>, enr_key: &K) -> Result<(), Error> {
+        let enr_backup = self.clone();
         self.content.remove(content_key.as_ref());
 
         // add the new public key.
@@ -838,10 +843,13 @@ impl<K: EnrKey> Enr<K> {
         self.content.insert(public_key.enr_key(), pubkey.freeze());
 
         // increment the sequence number
-        self.seq = self
-            .seq
-            .checked_add(1)
-            .ok_or(Error::SequenceNumberTooHigh)?;
+        match self.seq.checked_add(1) {
+            Some(seq_no) => self.seq = seq_no,
+            None => {
+                *self = enr_backup;
+                return Err(Error::SequenceNumberTooHigh);
+            }
+        }
 
         // sign the record
         self.sign(enr_key)?;
@@ -888,17 +896,23 @@ impl<K: EnrKey> Enr<K> {
             let value = out.freeze();
             // Prevent inserting invalid RLP integers
             if is_keyof_u16(key.as_ref()) {
-                u16::decode(&mut value.as_ref())?;
+                if let Err(err) = u16::decode(&mut value.as_ref()) {
+                    *self = enr_backup;
+                    return Err(err.into());
+                }
             }
 
             inserted.push(self.content.insert(key.as_ref().to_vec(), value));
         }
 
         // increment the sequence number
-        self.seq = self
-            .seq
-            .checked_add(1)
-            .ok_or(Error::SequenceNumberTooHigh)?;
+        match self.seq.checked_add(1) {
+            Some(seq_no) => self.seq = seq_no,
+            None => {
+                *self = enr_backup;
+                return Err(Error::SequenceNumberTooHigh);
+            }
+        }
 
         // sign the record
         self.sign(enr_key)?;