From 7b71af24097f16b0337fc97f8a7c83bfb05cb2d9 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Sat, 4 Mar 2023 21:51:35 +0100
Subject: [PATCH] test: add empty decode buffer check (#1628)

---
 crates/net/eth-wire/src/p2pstream.rs          | 18 ++++++++++++++++--
 crates/storage/codecs/derive/src/arbitrary.rs | 10 ++++++----
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/crates/net/eth-wire/src/p2pstream.rs b/crates/net/eth-wire/src/p2pstream.rs
index a599095018..9a487cc378 100644
--- a/crates/net/eth-wire/src/p2pstream.rs
+++ b/crates/net/eth-wire/src/p2pstream.rs
@@ -698,8 +698,22 @@ impl Encodable for P2PMessage {
 /// The [`Decodable`](reth_rlp::Decodable) implementation for [`P2PMessage`] assumes that each of
 /// the message variants are snappy compressed, except for the [`P2PMessage::Hello`] variant since
 /// the hello message is never compressed in the `p2p` subprotocol.
+/// The [`Decodable`] implementation for [`P2PMessage::Ping`] and
+/// [`P2PMessage::Pong`] expects a snappy encoded payload, see [`Encodable`] implementation.
 impl Decodable for P2PMessage {
     fn decode(buf: &mut &[u8]) -> Result<Self, DecodeError> {
+        /// Removes the snappy prefix from the Ping/Pong buffer
+        fn advance_snappy_ping_pong_payload(buf: &mut &[u8]) -> Result<(), DecodeError> {
+            if buf.len() < 3 {
+                return Err(DecodeError::InputTooShort)
+            }
+            if buf[..3] != [0x01, 0x00, EMPTY_LIST_CODE] {
+                return Err(DecodeError::Custom("expected snappy payload"))
+            }
+            buf.advance(3);
+            Ok(())
+        }
+
         let message_id = u8::decode(&mut &buf[..])?;
         let id = P2PMessageID::try_from(message_id)
             .or(Err(DecodeError::Custom("unknown p2p message id")))?;
@@ -708,11 +722,11 @@ impl Decodable for P2PMessage {
             P2PMessageID::Hello => Ok(P2PMessage::Hello(HelloMessage::decode(buf)?)),
             P2PMessageID::Disconnect => Ok(P2PMessage::Disconnect(DisconnectReason::decode(buf)?)),
             P2PMessageID::Ping => {
-                buf.advance(1);
+                advance_snappy_ping_pong_payload(buf)?;
                 Ok(P2PMessage::Ping)
             }
             P2PMessageID::Pong => {
-                buf.advance(1);
+                advance_snappy_ping_pong_payload(buf)?;
                 Ok(P2PMessage::Pong)
             }
         }
diff --git a/crates/storage/codecs/derive/src/arbitrary.rs b/crates/storage/codecs/derive/src/arbitrary.rs
index 96c259b49f..61f82057b5 100644
--- a/crates/storage/codecs/derive/src/arbitrary.rs
+++ b/crates/storage/codecs/derive/src/arbitrary.rs
@@ -36,10 +36,12 @@ pub fn maybe_generate_tests(args: TokenStream, ast: &DeriveInput) -> TokenStream
                 {
                     let mut buf = vec![];
 
-                    let len = field.clone().encode(&mut buf);
-                    let decoded = super::#type_ident::decode(&mut buf.as_slice()).unwrap();
-
-                    assert!(field == decoded);
+                    let len = field.encode(&mut buf);
+                    let mut b = &mut buf.as_slice();
+                    let decoded = super::#type_ident::decode(b).unwrap();
+                    assert_eq!(field, decoded);
+                    // ensure buffer is fully consumed by decode
+                    assert!(b.is_empty());
                 }
             });
         } else if let Ok(num) = arg.to_string().parse() {