From 76256f1a9020029fa38c62a4755d2c51ec09e481 Mon Sep 17 00:00:00 2001
From: Diva M <divma@protonmail.com>
Date: Mon, 25 Sep 2023 17:07:15 -0500
Subject: [PATCH] unify encoding rlp content

---
 src/lib.rs | 34 +++++++++++++++++++---------------
 1 file changed, 19 insertions(+), 15 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 6dd1e35..242e91f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -767,10 +767,15 @@ impl<K: EnrKey> Enr<K> {
 
     // Private Functions //
 
-    /// Evaluates the RLP-encoding of the content of the ENR record.
-    fn rlp_content(&self) -> BytesMut {
-        let mut stream = RlpStream::new_with_buffer(BytesMut::with_capacity(MAX_ENR_SIZE));
-        stream.begin_list(self.content.len() * 2 + 1);
+    /// Encodes the ENR's content (sequence number + ordered (key, value) pairs) into the stream.
+    fn append_rlp_content(&self, stream: &mut RlpStream, include_signature: bool) {
+        // signature(optional) + sequence number + (key, value) pairs
+        let content_pairs = self.content.len() * 2;
+        let item_count = include_signature.then_some(1).unwrap_or(0) + 1 + content_pairs;
+        stream.begin_list(item_count);
+        if include_signature {
+            stream.append(&self.signature);
+        }
         stream.append(&self.seq);
         for (k, v) in &self.content {
             // Keys are bytes
@@ -778,6 +783,13 @@ impl<K: EnrKey> Enr<K> {
             // Values are raw RLP encoded data
             stream.append_raw(v, 1);
         }
+    }
+
+    /// Encodes the ENR's content (sequence number + ordered (key, value) pairs).
+    fn rlp_content(&self) -> BytesMut {
+        let mut stream = RlpStream::new_with_buffer(BytesMut::with_capacity(MAX_ENR_SIZE));
+        let include_signature = false;
+        self.append_rlp_content(&mut stream, include_signature);
         stream.out()
     }
 
@@ -918,17 +930,9 @@ impl<'de, K: EnrKey> Deserialize<'de> for Enr<K> {
 }
 
 impl<K: EnrKey> rlp::Encodable for Enr<K> {
-    fn rlp_append(&self, s: &mut RlpStream) {
-        s.begin_list(self.content.len() * 2 + 2);
-        s.append(&self.signature);
-        s.append(&self.seq);
-        // must use rlp_content to preserve ordering.
-        for (k, v) in &self.content {
-            // Keys are byte data
-            s.append(k);
-            // Values are raw RLP encoded data
-            s.append_raw(v, 1);
-        }
+    fn rlp_append(&self, stream: &mut RlpStream) {
+        let include_signature = true;
+        self.append_rlp_content(stream, include_signature)
     }
 }