From b48d44859f985c8988ef758af7bca42c1d082506 Mon Sep 17 00:00:00 2001
From: Sarp Centel <sarpcentel@gmail.com>
Date: Tue, 27 Aug 2024 11:28:22 +0300
Subject: [PATCH] allow both {0,0,0,0,0,0,0,0} and {0,0,0,0,0,16,0,0} during
 handshake

---
 internal/stage_handshake.go      |  5 ++++-
 internal/stage_helpers.go        | 17 +++++++++++++----
 internal/stage_magnet_helpers.go |  2 +-
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/internal/stage_handshake.go b/internal/stage_handshake.go
index c6c8d6a..9ed7c3b 100644
--- a/internal/stage_handshake.go
+++ b/internal/stage_handshake.go
@@ -60,7 +60,10 @@ func testHandshake(stageHarness *test_case_harness.TestCaseHarness) error {
 		return err
 	}
 
-	expectedReservedBytes := []byte{0, 0, 0, 0, 0, 0, 0, 0}
+	expectedReservedBytes := [][]byte{
+		{0, 0, 0, 0, 0, 0, 0, 0},
+		{0, 0, 0, 0, 0, 16, 0, 0},
+	}
 
 	peersResponse := createPeersResponse("127.0.0.1", peerPort)
 
diff --git a/internal/stage_helpers.go b/internal/stage_helpers.go
index 3916d0d..a42d258 100644
--- a/internal/stage_helpers.go
+++ b/internal/stage_helpers.go
@@ -24,7 +24,7 @@ type PeerConnectionParams struct {
 	address string
 	myPeerID [20]byte
 	infoHash [20]byte
-	expectedReservedBytes []byte
+	expectedReservedBytes [][]byte
 	myMetadataExtensionID uint8
 	metadataSizeBytes int
 	bitfield []byte
@@ -414,8 +414,8 @@ func receiveAndSendHandshake(conn net.Conn, peer PeerConnectionParams) (err erro
 		return fmt.Errorf("error reading handshake: %s", err)
 	}
 	
-	if !bytes.Equal(handshake.Reserved[:], peer.expectedReservedBytes) {
-		return fmt.Errorf("did you send reserved bytes? expected bytes: %v but received: %v", peer.expectedReservedBytes, handshake.Reserved)
+	if !isEqualToOneOf(handshake.Reserved[:], peer.expectedReservedBytes...) {
+		return fmt.Errorf("did you send reserved bytes? expected bytes: %v but received: %v", peer.expectedReservedBytes[0], handshake.Reserved)
 	}
 
 	if !bytes.Equal(handshake.InfoHash[:], peer.infoHash[:]) {
@@ -426,7 +426,7 @@ func receiveAndSendHandshake(conn net.Conn, peer PeerConnectionParams) (err erro
 	logger.Debugf("Sending back handshake with peer_id: %x", peer.myPeerID)
 	
 	var reservedBytes [8]byte 
-	copy(reservedBytes[:], peer.expectedReservedBytes)
+	copy(reservedBytes[:], peer.expectedReservedBytes[0])
 
 	err = sendHandshake(conn, reservedBytes, handshake.InfoHash, peer.myPeerID)
 	if err != nil {
@@ -460,4 +460,13 @@ func randomHash() ([20]byte, error) {
 		return [20]byte{}, err
 	}
 	return hash, nil
+}
+
+func isEqualToOneOf(target []byte, arrays ...[]byte) bool {
+	for _, array := range arrays {
+		if bytes.Equal(target, array) {
+			return true
+		}
+	}
+	return false
 }
\ No newline at end of file
diff --git a/internal/stage_magnet_helpers.go b/internal/stage_magnet_helpers.go
index 386b9be..ff97c73 100644
--- a/internal/stage_magnet_helpers.go
+++ b/internal/stage_magnet_helpers.go
@@ -92,7 +92,7 @@ func (m *MagnetTestParams) toPeerConnectionParams() PeerConnectionParams {
 		address: m.PeerAddress,
 		myPeerID: m.ExpectedPeerID,
 		infoHash: m.ExpectedInfoHash,
-		expectedReservedBytes: m.ExpectedReservedBytes,
+		expectedReservedBytes: [][]byte{m.ExpectedReservedBytes},
 		myMetadataExtensionID: m.MyMetadataExtensionID,
 		metadataSizeBytes: m.MagnetLinkInfo.MetadataSizeBytes,
 		bitfield: m.MagnetLinkInfo.Bitfield,