fix: decode rlp header on Enr::build (#72)

Co-authored-by: Age Manning <[email protected]>

Author: mattsse

Cargo.toml

@@ -28,6 +28,7 @@ secp256k1 = { version = "0.28", optional = true, default-features = false, featu
 ] }
 
 [dev-dependencies]
+alloy-rlp = { version = "0.3.4", features = ["derive"] }
 secp256k1 = { features = ["rand-std"], version = "0.28" }
 serde_json = { version = "1.0.114" }
 

src/builder.rs

@@ -1,5 +1,5 @@
 use crate::{Enr, EnrKey, EnrPublicKey, Error, Key, NodeId, MAX_ENR_SIZE};
-use alloy_rlp::{Decodable, Encodable, Header};
+use alloy_rlp::{Encodable, Header};
 use bytes::{Bytes, BytesMut};
 use std::{
     collections::BTreeMap,
@@ -158,7 +158,7 @@ impl<K: EnrKey> Builder<K> {
 
         // Sanitize all data, ensuring all RLP data is correctly formatted.
         for value in self.content.values() {
-            Bytes::decode(&mut value.as_ref())?;
+            Header::decode(&mut value.as_ref())?;
         }
 
         let mut id_bytes = BytesMut::with_capacity(3);

src/lib.rs

@@ -269,8 +269,11 @@ impl<K: EnrKey> Enr<K> {
     #[allow(clippy::missing_panics_doc)]
     pub fn get(&self, key: impl AsRef<[u8]>) -> Option<Bytes> {
         // It's ok to decode any valid RLP value as data
-        self.get_raw_rlp(key)
-            .map(|mut rlp_data| Bytes::decode(&mut rlp_data).expect("All data is sanitized"))
+        self.get_raw_rlp(key).map(|mut rlp_data| {
+            let raw_data = &mut rlp_data;
+            let header = Header::decode(raw_data).expect("All data is sanitized");
+            raw_data[..header.payload_length].to_vec().into()
+        })
     }
 
     /// Reads a custom key from the record if it exists, decoded as `T`.
@@ -1492,6 +1495,30 @@ mod tests {
         assert!(decoded_enr.verify());
     }
 
+    #[test]
+    fn test_add_content_value() {
+        #[derive(PartialEq, Eq, Debug, alloy_rlp::RlpEncodable, alloy_rlp::RlpDecodable)]
+        struct Proto {
+            name: String,
+            version: u64,
+        }
+
+        let mut rng = rand::thread_rng();
+        let key = k256::ecdsa::SigningKey::random(&mut rng);
+        let proto = Proto {
+            name: "test".to_string(),
+            version: 1,
+        };
+
+        let enr = Enr::builder()
+            .add_value("proto", &proto)
+            .build(&key)
+            .unwrap();
+
+        let decoded_proto = enr.get_decodable::<Proto>("proto").unwrap().unwrap();
+        assert_eq!(decoded_proto, proto);
+    }
+
     #[test]
     fn test_add_key() {
         let mut rng = rand::thread_rng();