Skip to content

Commit a5d5ff6

Browse files
committed
feat: update OverlayContentKey trait and types inheriting it
1 parent f635bb0 commit a5d5ff6

File tree

14 files changed

+77
-65
lines changed

14 files changed

+77
-65
lines changed

ethportal-api/src/types/content_key/beacon.rs

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
use crate::{
22
types::content_key::{error::ContentKeyError, overlay::OverlayContentKey},
33
utils::bytes::hex_encode_compact,
4+
RawContentKey,
45
};
6+
use alloy_primitives::Bytes;
57
use serde::{Deserialize, Deserializer, Serialize, Serializer};
68
use sha2::{Digest, Sha256};
79
use ssz::{Decode, DecodeError, Encode};
@@ -81,22 +83,22 @@ pub struct HistoricalSummariesWithProofKey {
8183
pub epoch: u64,
8284
}
8385

84-
impl From<&BeaconContentKey> for Vec<u8> {
86+
impl From<&BeaconContentKey> for Bytes {
8587
fn from(val: &BeaconContentKey) -> Self {
8688
val.to_bytes()
8789
}
8890
}
8991

90-
impl From<BeaconContentKey> for Vec<u8> {
92+
impl From<BeaconContentKey> for Bytes {
9193
fn from(val: BeaconContentKey) -> Self {
9294
val.to_bytes()
9395
}
9496
}
9597

96-
impl TryFrom<Vec<u8>> for BeaconContentKey {
98+
impl TryFrom<RawContentKey> for BeaconContentKey {
9799
type Error = ContentKeyError;
98100

99-
fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
101+
fn try_from(value: RawContentKey) -> Result<Self, Self::Error> {
100102
let Some((&selector, key)) = value.split_first() else {
101103
return Err(ContentKeyError::InvalidLength {
102104
received: value.len(),
@@ -170,7 +172,7 @@ impl OverlayContentKey for BeaconContentKey {
170172
sha256.finalize().into()
171173
}
172174

173-
fn to_bytes(&self) -> Vec<u8> {
175+
fn to_bytes(&self) -> RawContentKey {
174176
let mut bytes: Vec<u8> = Vec::new();
175177

176178
match self {
@@ -196,7 +198,7 @@ impl OverlayContentKey for BeaconContentKey {
196198
}
197199
}
198200

199-
bytes
201+
bytes.into()
200202
}
201203
}
202204

ethportal-api/src/types/content_key/history.rs

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use std::{fmt, hash::Hash};
99
use crate::{
1010
types::content_key::{error::ContentKeyError, overlay::OverlayContentKey},
1111
utils::bytes::hex_encode_compact,
12+
RawContentKey,
1213
};
1314

1415
// Prefixes for the different types of history content keys:
@@ -44,6 +45,7 @@ impl HistoryContentKey {
4445
.choose(&mut rand::thread_rng())
4546
.ok_or_else(|| anyhow::Error::msg("Failed to choose random prefix"))?;
4647
random_bytes.insert(0, *random_prefix);
48+
let random_bytes: RawContentKey = random_bytes.into();
4749
Self::try_from(random_bytes).map_err(anyhow::Error::msg)
4850
}
4951
}
@@ -118,20 +120,20 @@ pub struct EpochAccumulatorKey {
118120

119121
impl From<&HistoryContentKey> for Vec<u8> {
120122
fn from(val: &HistoryContentKey) -> Self {
121-
val.to_bytes()
123+
val.to_bytes().to_vec()
122124
}
123125
}
124126

125127
impl From<HistoryContentKey> for Vec<u8> {
126128
fn from(val: HistoryContentKey) -> Self {
127-
val.to_bytes()
129+
val.to_bytes().to_vec()
128130
}
129131
}
130132

131-
impl TryFrom<Vec<u8>> for HistoryContentKey {
133+
impl TryFrom<RawContentKey> for HistoryContentKey {
132134
type Error = ContentKeyError;
133135

134-
fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
136+
fn try_from(value: RawContentKey) -> Result<Self, Self::Error> {
135137
let Some((&selector, key)) = value.split_first() else {
136138
return Err(ContentKeyError::InvalidLength {
137139
received: value.len(),
@@ -195,7 +197,7 @@ impl OverlayContentKey for HistoryContentKey {
195197
sha256.finalize().into()
196198
}
197199

198-
fn to_bytes(&self) -> Vec<u8> {
200+
fn to_bytes(&self) -> RawContentKey {
199201
let mut bytes: Vec<u8> = Vec::new();
200202

201203
match self {
@@ -217,7 +219,7 @@ impl OverlayContentKey for HistoryContentKey {
217219
}
218220
}
219221

220-
bytes
222+
bytes.into()
221223
}
222224
}
223225

@@ -326,7 +328,7 @@ mod test {
326328
});
327329

328330
// round trip
329-
let decoded = HistoryContentKey::try_from(key.to_bytes().to_vec()).unwrap();
331+
let decoded = HistoryContentKey::try_from(key.to_bytes()).unwrap();
330332
assert_eq!(decoded, key);
331333

332334
assert_eq!(key.to_bytes(), expected_content_key);

ethportal-api/src/types/content_key/overlay.rs

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,13 @@ use quickcheck::{Arbitrary, Gen};
55
use crate::{
66
types::content_key::error::ContentKeyError,
77
utils::bytes::{hex_decode, hex_encode, hex_encode_compact},
8+
RawContentKey,
89
};
910

1011
/// Types whose values represent keys to lookup content items in an overlay network.
1112
/// Keys are serializable.
1213
pub trait OverlayContentKey:
13-
Into<Vec<u8>>
14-
+ TryFrom<Vec<u8>, Error = ContentKeyError>
14+
TryFrom<RawContentKey, Error = ContentKeyError>
1515
+ Clone
1616
+ fmt::Debug
1717
+ fmt::Display
@@ -25,7 +25,7 @@ pub trait OverlayContentKey:
2525
fn content_id(&self) -> [u8; 32];
2626

2727
/// Returns the bytes of the content key.
28-
fn to_bytes(&self) -> Vec<u8>;
28+
fn to_bytes(&self) -> RawContentKey;
2929

3030
/// Returns the content key as a hex encoded "0x"-prefixed string.
3131
fn to_hex(&self) -> String {
@@ -34,8 +34,8 @@ pub trait OverlayContentKey:
3434

3535
/// Returns the content key from a hex encoded "0x"-prefixed string.
3636
fn from_hex(data: &str) -> anyhow::Result<Self> {
37-
let bytes = hex_decode(&data.to_lowercase())?;
38-
Ok(Self::try_from(bytes)?)
37+
let bytes = hex_decode(data)?;
38+
Ok(Self::try_from(bytes.into())?)
3939
}
4040
}
4141

@@ -67,10 +67,10 @@ impl Arbitrary for IdentityContentKey {
6767
}
6868
}
6969

70-
impl TryFrom<Vec<u8>> for IdentityContentKey {
70+
impl TryFrom<RawContentKey> for IdentityContentKey {
7171
type Error = ContentKeyError;
7272

73-
fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
73+
fn try_from(value: RawContentKey) -> Result<Self, Self::Error> {
7474
// Require that length of input is equal to 32.
7575
if value.len() != 32 {
7676
return Err(ContentKeyError::InvalidLength {
@@ -116,7 +116,7 @@ impl OverlayContentKey for IdentityContentKey {
116116
fn content_id(&self) -> [u8; 32] {
117117
self.value
118118
}
119-
fn to_bytes(&self) -> Vec<u8> {
120-
self.value.to_vec()
119+
fn to_bytes(&self) -> RawContentKey {
120+
self.value.to_vec().into()
121121
}
122122
}

ethportal-api/src/types/content_key/state.rs

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use std::{fmt, hash::Hash};
88
use crate::{
99
types::{content_key::overlay::OverlayContentKey, state_trie::nibbles::Nibbles},
1010
utils::bytes::hex_encode_compact,
11-
ContentKeyError,
11+
ContentKeyError, RawContentKey,
1212
};
1313

1414
// Prefixes for the different types of state content keys:
@@ -70,7 +70,7 @@ impl OverlayContentKey for StateContentKey {
7070
sha256.finalize().into()
7171
}
7272

73-
fn to_bytes(&self) -> Vec<u8> {
73+
fn to_bytes(&self) -> RawContentKey {
7474
let mut bytes: Vec<u8> = vec![];
7575

7676
match self {
@@ -88,25 +88,25 @@ impl OverlayContentKey for StateContentKey {
8888
}
8989
}
9090

91-
bytes
91+
bytes.into()
9292
}
9393
}
9494

9595
impl From<&StateContentKey> for Vec<u8> {
9696
fn from(val: &StateContentKey) -> Self {
97-
val.to_bytes()
97+
val.to_bytes().to_vec()
9898
}
9999
}
100100

101101
impl From<StateContentKey> for Vec<u8> {
102102
fn from(val: StateContentKey) -> Self {
103-
val.to_bytes()
103+
val.to_bytes().to_vec()
104104
}
105105
}
106106

107-
impl TryFrom<Vec<u8>> for StateContentKey {
107+
impl TryFrom<RawContentKey> for StateContentKey {
108108
type Error = ContentKeyError;
109-
fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
109+
fn try_from(value: RawContentKey) -> Result<Self, Self::Error> {
110110
let Some((&selector, key)) = value.split_first() else {
111111
return Err(ContentKeyError::from_decode_error(
112112
DecodeError::InvalidLengthPrefix {
@@ -187,7 +187,7 @@ impl fmt::Display for StateContentKey {
187187
mod test {
188188
use std::{path::PathBuf, str::FromStr};
189189

190-
use alloy_primitives::{keccak256, Address};
190+
use alloy_primitives::{bytes, keccak256, Address, Bytes};
191191
use anyhow::Result;
192192
use rstest::rstest;
193193
use serde_yaml::Value;
@@ -242,16 +242,20 @@ mod test {
242242
#[test]
243243
fn decode_empty_key_should_fail() {
244244
assert_eq!(
245-
StateContentKey::try_from(vec![]).unwrap_err().to_string(),
245+
StateContentKey::try_from(Bytes::new())
246+
.unwrap_err()
247+
.to_string(),
246248
"Unable to decode key SSZ bytes 0x due to InvalidLengthPrefix { len: 0, expected: 1 }",
247249
);
248250
}
249251

250252
#[test]
251253
fn decode_key_with_invalid_selector_should_fail() {
252-
let invalid_selector_content_key = "0x0024000000c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a4700005000000";
254+
let invalid_selector_content_key = bytes!(
255+
"0024000000c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a4700005000000"
256+
);
253257
assert_eq!(
254-
StateContentKey::try_from(hex_decode(invalid_selector_content_key).unwrap())
258+
StateContentKey::try_from(invalid_selector_content_key.clone())
255259
.unwrap_err()
256260
.to_string(),
257261
format!("Unable to decode key SSZ bytes {invalid_selector_content_key} due to UnionSelectorInvalid(0)"),
@@ -267,9 +271,9 @@ mod test {
267271
let yaml = yaml.as_mapping().unwrap();
268272

269273
let content_key_bytes = hex_decode(yaml["content_key"].as_str().unwrap())?;
270-
let content_key = StateContentKey::try_from(content_key_bytes.clone())?;
274+
let content_key = StateContentKey::try_from(Bytes::from(content_key_bytes.clone()))?;
271275

272-
assert_eq!(content_key.to_bytes(), content_key_bytes);
276+
assert_eq!(content_key.to_bytes(), Bytes::from(content_key_bytes));
273277
Ok(())
274278
}
275279

@@ -300,7 +304,7 @@ mod test {
300304
let yaml = yaml.as_mapping().unwrap();
301305

302306
let content_key_bytes = hex_decode(yaml["content_key"].as_str().unwrap())?;
303-
let content_key = StateContentKey::try_from(content_key_bytes)?;
307+
let content_key = StateContentKey::try_from(Bytes::from(content_key_bytes))?;
304308
let expected_content_id = yaml_as_b256(&yaml["content_id"]);
305309

306310
assert_eq!(B256::from(content_key.content_id()), expected_content_id);
@@ -337,7 +341,7 @@ mod test {
337341

338342
fn assert_content_key(value: &Value, expected_content_key: StateContentKey) -> Result<()> {
339343
assert_eq!(
340-
StateContentKey::try_from(yaml_as_hex(value))?,
344+
StateContentKey::try_from(Bytes::from(yaml_as_hex(value)))?,
341345
expected_content_key,
342346
"decoding from bytes {value:?} didn't match expected key {expected_content_key:?}"
343347
);

ethportal-api/src/types/portal_wire.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -622,6 +622,7 @@ impl From<Accept> for Value {
622622
#[allow(clippy::unwrap_used)]
623623
mod test {
624624
use super::*;
625+
use alloy_primitives::bytes;
625626
use std::str::FromStr;
626627
use test_log::test;
627628

@@ -821,7 +822,7 @@ mod test {
821822

822823
#[test]
823824
fn message_encoding_offer() {
824-
let content_keys = vec![hex_decode("0x010203").unwrap().into()];
825+
let content_keys = vec![bytes!("010203")];
825826
let offer = Offer { content_keys };
826827
let offer = Message::Offer(offer);
827828

portalnet/src/find/query_info.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ impl<TContentKey: OverlayContentKey> QueryInfo<TContentKey> {
7373
Request::FindNodes(FindNodes { distances })
7474
}
7575
QueryType::FindContent { ref target, .. } => Request::FindContent(FindContent {
76-
content_key: target.clone().into(),
76+
content_key: target.clone().to_bytes().to_vec(),
7777
}),
7878
};
7979

portalnet/src/gossip.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ pub fn propagate_gossip_cross_thread<TContentKey: OverlayContentKey>(
135135
// change content keys to raw content keys
136136
let interested_content = interested_content
137137
.into_iter()
138-
.map(|(key, value)| (key.into().into(), value))
138+
.map(|(key, value)| (key.to_bytes(), value))
139139
.collect();
140140
let offer_request = Request::PopulatedOffer(PopulatedOffer {
141141
content_items: interested_content,
@@ -197,7 +197,7 @@ pub async fn trace_propagate_gossip_cross_thread<TContentKey: OverlayContentKey>
197197
for enr in interested_enrs.into_iter() {
198198
let (result_tx, mut result_rx) = tokio::sync::mpsc::unbounded_channel();
199199
let offer_request = Request::PopulatedOfferWithResult(PopulatedOfferWithResult {
200-
content_item: (content_key.clone().into().into(), data.clone()),
200+
content_item: (content_key.clone().to_bytes(), data.clone()),
201201
result_tx,
202202
});
203203

portalnet/src/overlay/protocol.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ impl<
9494
TStore: 'static + ContentStore<Key = TContentKey> + Send + Sync,
9595
> OverlayProtocol<TContentKey, TMetric, TValidator, TStore>
9696
where
97-
<TContentKey as TryFrom<Vec<u8>>>::Error: Debug + Display + Send,
97+
<TContentKey as TryFrom<RawContentKey>>::Error: Debug + Display + Send,
9898
{
9999
pub async fn new(
100100
config: OverlayConfig,
@@ -441,7 +441,7 @@ where
441441
let direction = RequestDirection::Outgoing {
442442
destination: enr.clone(),
443443
};
444-
let content_key = TContentKey::try_from(content_key).map_err(|err| {
444+
let content_key = TContentKey::try_from(content_key.into()).map_err(|err| {
445445
OverlayRequestError::FailedValidation(format!(
446446
"Error decoding content key for received utp content: {err}"
447447
))

0 commit comments

Comments
 (0)