From 5a1209595970eb531807dc00d5a7dc5ca27d608a Mon Sep 17 00:00:00 2001
From: Igor Katson <igor.katson@gmail.com>
Date: Sat, 9 Mar 2024 19:50:28 +0000
Subject: [PATCH] Update handshake deserialize to error on wrong messages

---
 crates/peer_binary_protocol/src/lib.rs | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/crates/peer_binary_protocol/src/lib.rs b/crates/peer_binary_protocol/src/lib.rs
index f1128b1..97e6dcc 100644
--- a/crates/peer_binary_protocol/src/lib.rs
+++ b/crates/peer_binary_protocol/src/lib.rs
@@ -531,12 +531,15 @@ impl Handshake<ByteBuf<'static>> {
                 expected_len,
                 "handshake",
             ))?;
-        Ok((
-            Self::bopts()
-                .deserialize(hbuf)
-                .map_err(|e| MessageDeserializeError::Other(e.into()))?,
-            expected_len,
-        ))
+        let h = Self::bopts()
+            .deserialize::<Handshake<ByteBuf<'_>>>(hbuf)
+            .map_err(|e| MessageDeserializeError::Other(e.into()))?;
+        if h.pstr.0 != PSTR_BT1.as_bytes() {
+            return Err(MessageDeserializeError::Other(anyhow::anyhow!(
+                "pstr doesn't match bittorrent V1"
+            )));
+        }
+        Ok((h, expected_len))
     }
 }