Skip to content

Commit

Permalink
fix: decode rlp header on Enr::build (#72)
Browse files Browse the repository at this point in the history
Co-authored-by: Age Manning <Age@AgeManning.com>
  • Loading branch information
mattsse and AgeManning committed Apr 9, 2024
1 parent 1143e4b commit c76c9e5
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 4 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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" }

Expand Down
4 changes: 2 additions & 2 deletions src/builder.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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);
Expand Down
31 changes: 29 additions & 2 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand Down Expand Up @@ -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();
Expand Down

0 comments on commit c76c9e5

Please sign in to comment.