Skip to content

Commit aaccfe3

Browse files
AgeManningemhane
andauthored
Fix decoding of values as RLP lists (#76)
Co-authored-by: Emilia Hane <[email protected]> Co-authored-by: Emilia Hane <[email protected]>
1 parent df661bb commit aaccfe3

File tree

1 file changed

+151
-21
lines changed

1 file changed

+151
-21
lines changed

src/lib.rs

+151-21
Original file line numberDiff line numberDiff line change
@@ -265,10 +265,11 @@ impl<K: EnrKey> Enr<K> {
265265
self.seq
266266
}
267267

268-
/// Reads a custom key from the record if it exists, decoded as data.
268+
/// Reads a custom key from the record if it exists, decoded as data. Caution! Only use for
269+
/// data that is not an aggregate type. Default to using [`get_decodable`](Enr::get_decodable).
270+
#[deprecated(since = "0.11.2", note = "use get_decodable")]
269271
#[allow(clippy::missing_panics_doc)]
270272
pub fn get(&self, key: impl AsRef<[u8]>) -> Option<Bytes> {
271-
// It's ok to decode any valid RLP value as data
272273
self.get_raw_rlp(key).map(|mut rlp_data| {
273274
let raw_data = &mut rlp_data;
274275
let header = Header::decode(raw_data).expect("All data is sanitized");
@@ -298,7 +299,7 @@ impl<K: EnrKey> Enr<K> {
298299
/// Returns the IPv4 address of the ENR record if it is defined.
299300
#[must_use]
300301
pub fn ip4(&self) -> Option<Ipv4Addr> {
301-
if let Some(ip_bytes) = self.get("ip") {
302+
if let Some(Ok(ip_bytes)) = self.get_decodable::<Bytes>("ip") {
302303
return match ip_bytes.as_ref().len() {
303304
4 => {
304305
let mut ip = [0_u8; 4];
@@ -314,7 +315,7 @@ impl<K: EnrKey> Enr<K> {
314315
/// Returns the IPv6 address of the ENR record if it is defined.
315316
#[must_use]
316317
pub fn ip6(&self) -> Option<Ipv6Addr> {
317-
if let Some(ip_bytes) = self.get("ip6") {
318+
if let Some(Ok(ip_bytes)) = self.get_decodable::<Bytes>("ip6") {
318319
return match ip_bytes.as_ref().len() {
319320
16 => {
320321
let mut ip = [0_u8; 16];
@@ -330,7 +331,7 @@ impl<K: EnrKey> Enr<K> {
330331
/// The `id` of ENR record if it is defined.
331332
#[must_use]
332333
pub fn id(&self) -> Option<String> {
333-
if let Some(id_bytes) = self.get("id") {
334+
if let Some(Ok(id_bytes)) = self.get_decodable::<Bytes>("id") {
334335
return Some(String::from_utf8_lossy(id_bytes.as_ref()).to_string());
335336
}
336337
None
@@ -1031,24 +1032,15 @@ impl<K: EnrKey> Decodable for Enr<K> {
10311032
let keys = Header::decode_bytes(payload, false)?;
10321033
alloy_rlp::encode(keys)
10331034
}
1034-
b"eth" => {
1035-
let eth_header = Header::decode(payload)?;
1036-
let value = &mut &payload[..eth_header.payload_length];
1037-
payload.advance(eth_header.payload_length);
1038-
let val_header = Header {
1039-
list: true,
1040-
payload_length: value.len(),
1041-
};
1042-
let mut out = Vec::<u8>::new();
1043-
val_header.encode(&mut out);
1044-
out.extend_from_slice(value);
1045-
out
1046-
}
10471035
_ => {
10481036
let other_header = Header::decode(payload)?;
10491037
let value = &payload[..other_header.payload_length];
1038+
// Preserve the valid encoding
10501039
payload.advance(other_header.payload_length);
1051-
alloy_rlp::encode(value)
1040+
let mut out = Vec::<u8>::new();
1041+
other_header.encode(&mut out);
1042+
out.extend_from_slice(value);
1043+
out
10521044
}
10531045
};
10541046
content.insert(key.to_vec(), Bytes::from(value));
@@ -1143,10 +1135,46 @@ fn check_spec_reserved_keys(key: &[u8], mut value: &[u8]) -> Result<(), Error> {
11431135
#[cfg(feature = "k256")]
11441136
mod tests {
11451137
use super::*;
1138+
use alloy_rlp::{RlpDecodable, RlpEncodable};
11461139
use std::convert::TryFrom;
11471140

11481141
type DefaultEnr = Enr<k256::ecdsa::SigningKey>;
11491142

1143+
// Struct for testing the case that an ENR value is an RLP encodable list.
1144+
#[derive(RlpEncodable, RlpDecodable, Debug, PartialEq, Eq)]
1145+
struct GenericListValue {
1146+
first_field: [u8; 4],
1147+
second_field: [u8; 4],
1148+
third_field: u64,
1149+
}
1150+
1151+
impl GenericListValue {
1152+
fn gen_random() -> Self {
1153+
Self {
1154+
first_field: rand::random(),
1155+
second_field: rand::random(),
1156+
third_field: rand::random(),
1157+
}
1158+
}
1159+
}
1160+
1161+
// Struct for testing the case that an ENR value is a recursive RLP encodable list (just to
1162+
// depth 2).
1163+
#[derive(RlpEncodable, RlpDecodable, Debug, PartialEq, Eq)]
1164+
struct DoubleListValue {
1165+
first_field: [u8; 4],
1166+
second_field: GenericListValue,
1167+
}
1168+
1169+
impl DoubleListValue {
1170+
fn gen_random() -> Self {
1171+
Self {
1172+
first_field: rand::random(),
1173+
second_field: GenericListValue::gen_random(),
1174+
}
1175+
}
1176+
}
1177+
11501178
#[cfg(feature = "k256")]
11511179
#[test]
11521180
fn test_vector_k256() {
@@ -1227,6 +1255,78 @@ mod tests {
12271255
assert!(enr.verify());
12281256
}
12291257

1258+
#[test]
1259+
fn test_encode_decode_list_value() {
1260+
let key = k256::ecdsa::SigningKey::random(&mut rand::rngs::OsRng);
1261+
let ip = Ipv4Addr::new(127, 0, 0, 1);
1262+
let tcp = 3000;
1263+
let list_value = GenericListValue::gen_random();
1264+
1265+
let enr = Enr::builder()
1266+
.ip4(ip)
1267+
.tcp4(tcp)
1268+
.add_value("list_value", &list_value)
1269+
.build(&key)
1270+
.unwrap();
1271+
1272+
let mut encoded_enr = BytesMut::new();
1273+
enr.encode(&mut encoded_enr);
1274+
1275+
let decoded_enr =
1276+
Enr::<k256::ecdsa::SigningKey>::decode(&mut encoded_enr.to_vec().as_slice()).unwrap();
1277+
1278+
assert_eq!(decoded_enr.id(), Some("v4".into()));
1279+
assert_eq!(decoded_enr.ip4(), Some(ip));
1280+
assert_eq!(decoded_enr.tcp4(), Some(tcp));
1281+
assert_eq!(
1282+
decoded_enr
1283+
.get_decodable::<GenericListValue>("list_value")
1284+
.unwrap()
1285+
.unwrap(),
1286+
list_value
1287+
);
1288+
// Must compare encoding as the public key itself can be different
1289+
assert_eq!(decoded_enr.public_key().encode(), key.public().encode());
1290+
decoded_enr.public_key().encode_uncompressed();
1291+
assert!(decoded_enr.verify());
1292+
}
1293+
1294+
#[test]
1295+
fn test_encode_decode_double_list_value() {
1296+
let key = k256::ecdsa::SigningKey::random(&mut rand::rngs::OsRng);
1297+
let ip = Ipv4Addr::new(127, 0, 0, 1);
1298+
let tcp = 3000;
1299+
let list_value = DoubleListValue::gen_random();
1300+
1301+
let enr = Enr::builder()
1302+
.ip4(ip)
1303+
.tcp4(tcp)
1304+
.add_value("list_value", &list_value)
1305+
.build(&key)
1306+
.unwrap();
1307+
1308+
let mut encoded_enr = BytesMut::new();
1309+
enr.encode(&mut encoded_enr);
1310+
1311+
let decoded_enr =
1312+
Enr::<k256::ecdsa::SigningKey>::decode(&mut encoded_enr.to_vec().as_slice()).unwrap();
1313+
1314+
assert_eq!(decoded_enr.id(), Some("v4".into()));
1315+
assert_eq!(decoded_enr.ip4(), Some(ip));
1316+
assert_eq!(decoded_enr.tcp4(), Some(tcp));
1317+
assert_eq!(
1318+
decoded_enr
1319+
.get_decodable::<DoubleListValue>("list_value")
1320+
.unwrap()
1321+
.unwrap(),
1322+
list_value
1323+
);
1324+
// Must compare encoding as the public key itself can be different
1325+
assert_eq!(decoded_enr.public_key().encode(), key.public().encode());
1326+
decoded_enr.public_key().encode_uncompressed();
1327+
assert!(decoded_enr.verify());
1328+
}
1329+
12301330
// the values in the content are rlp lists
12311331
#[test]
12321332
fn test_rlp_list_value() {
@@ -1519,6 +1619,36 @@ mod tests {
15191619
assert_eq!(decoded_proto, proto);
15201620
}
15211621

1622+
#[test]
1623+
fn test_add_content_value_decoding() {
1624+
#[derive(PartialEq, Eq, Debug, alloy_rlp::RlpEncodable, alloy_rlp::RlpDecodable)]
1625+
struct Proto {
1626+
name: String,
1627+
version: u64,
1628+
}
1629+
1630+
let mut rng = rand::thread_rng();
1631+
let key = k256::ecdsa::SigningKey::random(&mut rng);
1632+
let proto = Proto {
1633+
name: "test".to_string(),
1634+
version: 1,
1635+
};
1636+
1637+
let enr = Enr::builder()
1638+
.add_value("proto", &proto)
1639+
.build(&key)
1640+
.unwrap();
1641+
1642+
let encoded_enr = alloy_rlp::encode(enr.clone());
1643+
let decoded_enr = Enr::<k256::ecdsa::SigningKey>::decode(&mut &encoded_enr[..]).unwrap();
1644+
1645+
let decoded_proto = decoded_enr
1646+
.get_decodable::<Proto>("proto")
1647+
.unwrap()
1648+
.unwrap();
1649+
assert_eq!(decoded_proto, proto);
1650+
}
1651+
15221652
#[test]
15231653
fn test_add_key() {
15241654
let mut rng = rand::thread_rng();
@@ -1626,7 +1756,7 @@ mod tests {
16261756
let mut enr = Enr::builder().tcp4(tcp).build(&key).unwrap();
16271757

16281758
assert_eq!(enr.tcp4(), Some(tcp));
1629-
assert_eq!(enr.get("topics").as_ref().map(AsRef::as_ref), None);
1759+
assert_eq!(enr.get_decodable::<BytesMut>("topics"), None);
16301760

16311761
let topics: &[u8] = &out;
16321762

@@ -1639,7 +1769,7 @@ mod tests {
16391769
assert_eq!(inserted[0], None);
16401770

16411771
assert_eq!(enr.tcp4(), None);
1642-
assert_eq!(enr.get("topics").as_ref().map(AsRef::as_ref), Some(topics));
1772+
assert_eq!(enr.get_decodable("topics"), Some(Ok(out)));
16431773

16441774
// Compare the encoding as the key itself can be different
16451775
assert_eq!(enr.public_key().encode(), key.public().encode());

0 commit comments

Comments
 (0)