Skip to content

Commit

Permalink
fix: ignore extra data in discv4 messages (#110)
Browse files Browse the repository at this point in the history
**Motivation**

According to the spec:
> Implementations should ignore any additional elements in the
`packet-data` list as well as any extra data after the list.

We're currently failing if any of that happens.

**Description**

This PR introduces `Decoder::finish_unchecked`, which ignores any
trailing bytes instead of failing, and uses that to allow for additional
elements inside the packet-data. Also, it replaces some usage of
`RLPDecode::decode` with `RLPDecode::decode_unfinished` to allow for
extra data after the message.

As an extra, it refactors out the big `match` inside of both `Message`
encode and decode functions, since it's expected those will change often
in the next PRs.

Related to #17
  • Loading branch information
MegaRedHand authored Jul 3, 2024
1 parent 449bc51 commit eaeb9ce
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 27 deletions.
8 changes: 8 additions & 0 deletions crates/core/src/rlp/structs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,13 +80,21 @@ impl<'a> Decoder<'a> {
}
}

/// Finishes encoding the struct and returns the remaining bytes after the item.
/// If the item's payload is not empty, returns an error.
pub fn finish(self) -> Result<&'a [u8], RLPDecodeError> {
if self.payload.is_empty() {
Ok(self.remaining)
} else {
Err(RLPDecodeError::MalformedData)
}
}

/// Same as [`finish`](Self::finish), but discards the item's remaining payload
/// instead of failing.
pub fn finish_unchecked(self) -> &'a [u8] {
self.remaining
}
}

fn field_decode_error<T>(field_name: &str, err: RLPDecodeError) -> RLPDecodeError {
Expand Down
66 changes: 41 additions & 25 deletions crates/net/src/discv4.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,9 @@ use k256::ecdsa::{signature::Signer, SigningKey};
use std::net::IpAddr;

#[derive(Debug, Eq, PartialEq)]
// TODO: remove when all variants are used
// NOTE: All messages could have more fields than specified by the spec.
// Those additional fields should be ignored, and the message must be accepted.
// TODO: remove when all variants are used
#[allow(dead_code)]
pub(crate) enum Message {
/// A ping message. Should be responded to with a Pong message.
Expand All @@ -29,13 +29,8 @@ impl Message {
let signature_size = 65_usize;
let mut data: Vec<u8> = Vec::with_capacity(signature_size.next_power_of_two());
data.resize(signature_size, 0);
data.push(self.packet_type());
match self {
Message::Ping(msg) => msg.encode(&mut data),
Message::Pong(msg) => msg.encode(&mut data),
Message::FindNode(msg) => msg.encode(&mut data),
_ => todo!(),
}

self.encode_with_type(&mut data);

let digest = keccak_hash::keccak_buffer(&mut &data[signature_size..]).unwrap();

Expand All @@ -50,35 +45,56 @@ impl Message {
buf.put_slice(&data[..]);
}

fn packet_type(&self) -> u8 {
fn encode_with_type(&self, buf: &mut dyn BufMut) {
buf.put_u8(self.packet_type());
match self {
Message::Ping(_) => 0x01,
Message::Pong(_) => 0x02,
Message::FindNode(_) => 0x03,
Message::Neighbors(_) => 0x04,
Message::ENRRequest(_) => 0x05,
Message::ENRResponse(_) => 0x06,
Message::Ping(msg) => msg.encode(buf),
Message::Pong(msg) => msg.encode(buf),
Message::FindNode(msg) => msg.encode(buf),
_ => todo!(),
}
}
#[allow(unused)]

pub fn decode_with_header(encoded_msg: &[u8]) -> Result<Message, RLPDecodeError> {
let signature_len = 65;
let hash_len = 32;
let packet_index = signature_len + hash_len;
let signature_len = 65;
let packet_index = hash_len + signature_len;

// TODO: verify hash and signature
let _hash = &encoded_msg[..hash_len];
let _signature = &encoded_msg[hash_len..packet_index];

let packet_type = encoded_msg[packet_index];
let msg = &encoded_msg[packet_index + 1..];

Self::decode_with_type(packet_type, &encoded_msg[packet_index + 1..])
}

fn decode_with_type(packet_type: u8, msg: &[u8]) -> Result<Message, RLPDecodeError> {
// NOTE: extra elements inside the message should be ignored, along with extra data
// after the message.
match packet_type {
0x01 => {
let ping = PingMessage::decode(msg)?;
let (ping, _rest) = PingMessage::decode_unfinished(msg)?;
Ok(Message::Ping(ping))
}
0x02 => {
let pong = PongMessage::decode(msg)?;
let (pong, _rest) = PongMessage::decode_unfinished(msg)?;
Ok(Message::Pong(pong))
}
_ => todo!(),
}
}

fn packet_type(&self) -> u8 {
match self {
Message::Ping(_) => 0x01,
Message::Pong(_) => 0x02,
Message::FindNode(_) => 0x03,
Message::Neighbors(_) => 0x04,
Message::ENRRequest(_) => 0x05,
Message::ENRResponse(_) => 0x06,
}
}
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
Expand Down Expand Up @@ -203,8 +219,8 @@ impl RLPDecode for PingMessage {
expiration,
enr_seq,
};

let remaining = decoder.finish()?;
// NOTE: as per the spec, any additional elements should be ignored.
let remaining = decoder.finish_unchecked();
Ok((ping, remaining))
}
}
Expand Down Expand Up @@ -268,8 +284,8 @@ impl RLPDecode for PongMessage {
expiration,
enr_seq,
};
let remaining = decoder.finish()?;

// NOTE: as per the spec, any additional elements should be ignored.
let remaining = decoder.finish_unchecked();
Ok((pong, remaining))
}
}
Expand Down
4 changes: 2 additions & 2 deletions justfile
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@ lint:
test-all:
cargo test --workspace

test crate:
cargo test -p {{crate}}
test crate='*':
cargo test -p '{{crate}}'

clean:
cargo clean
Expand Down

0 comments on commit eaeb9ce

Please sign in to comment.