feat(eth-wire): add td bitlen check in handshake (#1490)

crates/net/eth-wire/src/errors/eth.rs

@@ -74,4 +74,6 @@ pub enum EthHandshakeError {
     MismatchedProtocolVersion { expected: u8, got: u8 },
     #[error("mismatched chain in Status message. expected: {expected:?}, got: {got:?}")]
     MismatchedChain { expected: Chain, got: Chain },
+    #[error("total difficulty bitlen is too large. maximum: {maximum:?}, got: {got:?}")]
+    TotalDifficultyBitLenTooLarge { maximum: usize, got: usize },
 }

crates/net/eth-wire/src/ethstream.rs

@@ -86,7 +86,7 @@ where
             }
         };
 
-        // TODO: Add any missing checks
+        // The following checks should match the checks in go-ethereum:
         // https://github.com/ethereum/go-ethereum/blob/9244d5cd61f3ea5a7645fdf2a1a96d53421e412f/eth/protocols/eth/handshake.go#L87-L89
         match msg.message {
             EthMessage::Status(resp) => {
@@ -118,6 +118,16 @@ where
                     .into())
                 }
 
+                // TD at mainnet block #7753254 is 76 bits. If it becomes 100 million times
+                // larger, it will still fit within 100 bits
+                if status.total_difficulty.bit_len() > 100 {
+                    return Err(EthHandshakeError::TotalDifficultyBitLenTooLarge {
+                        maximum: 100,
+                        got: status.total_difficulty.bit_len(),
+                    }
+                    .into())
+                }
+
                 fork_filter.validate(resp.forkid).map_err(EthHandshakeError::InvalidFork)?;
 
                 // now we can create the `EthStream` because the peer has successfully completed
@@ -268,6 +278,7 @@ mod tests {
     use super::UnauthedEthStream;
     use crate::{
         capability::Capability,
+        errors::{EthHandshakeError, EthStreamError},
         hello::HelloMessage,
         p2pstream::{ProtocolVersion, UnauthedP2PStream},
         types::{broadcast::BlockHashNumber, EthMessage, EthVersion, Status},
@@ -328,6 +339,107 @@ mod tests {
         handle.await.unwrap();
     }
 
+    #[tokio::test]
+    async fn pass_handshake_on_low_td_bitlen() {
+        let genesis = H256::random();
+        let fork_filter = ForkFilter::new(Head::default(), genesis, Vec::new());
+
+        let status = Status {
+            version: EthVersion::Eth67 as u8,
+            chain: Chain::Mainnet.into(),
+            total_difficulty: U256::from(2).pow(U256::from(100)) - U256::from(1),
+            blockhash: H256::random(),
+            genesis,
+            // Pass the current fork id.
+            forkid: fork_filter.current(),
+        };
+
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let local_addr = listener.local_addr().unwrap();
+
+        let status_clone = status;
+        let fork_filter_clone = fork_filter.clone();
+        let handle = tokio::spawn(async move {
+            // roughly based off of the design of tokio::net::TcpListener
+            let (incoming, _) = listener.accept().await.unwrap();
+            let stream = PassthroughCodec::default().framed(incoming);
+            let (_, their_status) = UnauthedEthStream::new(stream)
+                .handshake(status_clone, fork_filter_clone)
+                .await
+                .unwrap();
+
+            // just make sure it equals our status, and that the handshake succeeded
+            assert_eq!(their_status, status_clone);
+        });
+
+        let outgoing = TcpStream::connect(local_addr).await.unwrap();
+        let sink = PassthroughCodec::default().framed(outgoing);
+
+        // try to connect
+        let (_, their_status) =
+            UnauthedEthStream::new(sink).handshake(status, fork_filter).await.unwrap();
+
+        // their status is a clone of our status, these should be equal
+        assert_eq!(their_status, status);
+
+        // await the other handshake
+        handle.await.unwrap();
+    }
+
+    #[tokio::test]
+    async fn fail_handshake_on_high_td_bitlen() {
+        let genesis = H256::random();
+        let fork_filter = ForkFilter::new(Head::default(), genesis, Vec::new());
+
+        let status = Status {
+            version: EthVersion::Eth67 as u8,
+            chain: Chain::Mainnet.into(),
+            total_difficulty: U256::from(2).pow(U256::from(100)),
+            blockhash: H256::random(),
+            genesis,
+            // Pass the current fork id.
+            forkid: fork_filter.current(),
+        };
+
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let local_addr = listener.local_addr().unwrap();
+
+        let status_clone = status;
+        let fork_filter_clone = fork_filter.clone();
+        let handle = tokio::spawn(async move {
+            // roughly based off of the design of tokio::net::TcpListener
+            let (incoming, _) = listener.accept().await.unwrap();
+            let stream = PassthroughCodec::default().framed(incoming);
+            let handshake_res =
+                UnauthedEthStream::new(stream).handshake(status_clone, fork_filter_clone).await;
+
+            // make sure the handshake fails due to td too high
+            assert!(matches!(
+                handshake_res,
+                Err(EthStreamError::EthHandshakeError(
+                    EthHandshakeError::TotalDifficultyBitLenTooLarge { maximum: 100, got: 101 }
+                ))
+            ));
+        });
+
+        let outgoing = TcpStream::connect(local_addr).await.unwrap();
+        let sink = PassthroughCodec::default().framed(outgoing);
+
+        // try to connect
+        let handshake_res = UnauthedEthStream::new(sink).handshake(status, fork_filter).await;
+
+        // this handshake should also fail due to td too high
+        assert!(matches!(
+            handshake_res,
+            Err(EthStreamError::EthHandshakeError(
+                EthHandshakeError::TotalDifficultyBitLenTooLarge { maximum: 100, got: 101 }
+            ))
+        ));
+
+        // await the other handshake
+        handle.await.unwrap();
+    }
+
     #[tokio::test]
     async fn can_write_and_read_cleartext() {
         let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();