add tests for announce extension support stage

add tests for announce extension support stage

internal/bittorrent.go

@@ -6,7 +6,6 @@ import (
  "encoding/hex"
  "fmt"
  "io"
- "log"
  "net"
  "os"
  "strings"
@@ -17,6 +16,7 @@ import (
 
 type Handshake struct {
  ProtocolStr string
+ Reserved [8]byte
  InfoHash [20]byte
  PeerID [20]byte
 }
@@ -153,28 +153,32 @@ func readHandshake(r io.Reader, logger *logger.Logger) (*Handshake, error) {
  logger.Infof("Did you send reserved bytes? expected bytes: %v but received: %v\n", expectedReservedBytes, actualReservedBytes)
  }
 
+ var reservedBytes [8]byte
+ copy(reservedBytes[:], handshakeBuffer[protocolNameLength : protocolNameLength+8])
  var infoHash, peerID [20]byte
  copy(infoHash[:], handshakeBuffer[protocolNameLength+8:protocolNameLength+8+20])
  copy(peerID[:], handshakeBuffer[protocolNameLength+8+20:])
 
  handshake := Handshake{
  ProtocolStr: string(handshakeBuffer[0:protocolNameLength]),
+ Reserved: reservedBytes,
  InfoHash: infoHash,
  PeerID: peerID,
  }
 
  return &handshake, nil
 }
 
-func sendHandshake(conn net.Conn, infoHash [20]byte, peerID [20]byte) {
+func sendHandshake(conn net.Conn, reserved [8]byte, infoHash [20]byte, peerID [20]byte) error {
  handshake := []byte{19} // Protocol name length
  handshake = append(handshake, []byte(ProtocolName)...) // Protocol name
- handshake = append(handshake, make([]byte, 8)...) // Reserved bytes
+ handshake = append(handshake, reserved[:]...)  // Reserved bytes
  handshake = append(handshake, infoHash[:]...) // Info hash
  handshake = append(handshake, peerID[:]...) // Peer ID
 
  _, err := conn.Write(handshake)
  if err != nil {
- log.Fatal("Error sending handshake:", err)
+ return fmt.Errorf("error sending handshake: %v", err)
  }
+ return nil
 }

internal/stage_handshake.go

@@ -1,14 +1,11 @@
 package internal
 
 import (
- "bytes"
  "fmt"
- "math/rand"
  "net"
  "os"
  "path"
 
- logger "github.com/codecrafters-io/tester-utils/logger"
  "github.com/codecrafters-io/tester-utils/test_case_harness"
 )
 
@@ -63,9 +60,32 @@ func testHandshake(stageHarness *test_case_harness.TestCaseHarness) error {
  return err
  }
 
+ 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)
- go listenAndServePeersResponse(trackerAddress, peersResponse, infoHash, fileLengthBytes, logger)
- go waitAndHandlePeerConnection(peerAddress, expectedPeerID, infoHash, logger)
+
+ go listenAndServeTrackerResponse(
+ TrackerParams {
+ trackerAddress: trackerAddress,
+ peersResponse: peersResponse,
+ expectedInfoHash: infoHash,
+ fileLengthBytes: fileLengthBytes,
+ logger: logger,
+ })
+
+ go waitAndHandlePeerConnection(
+ PeerConnectionParams {
+ address: peerAddress,
+ myPeerID: expectedPeerID,
+ infoHash: infoHash,
+ expectedReservedBytes: expectedReservedBytes,
+ logger: logger,
+ },
+ handleHandshake,
+ )
 
  logger.Infof("Running ./%s handshake %s %s", path.Base(executable.Path), torrentFilePath, peerAddress)
  result, err := executable.Run("handshake", torrentFilePath, peerAddress)
@@ -86,46 +106,11 @@ func testHandshake(stageHarness *test_case_harness.TestCaseHarness) error {
  return nil
 }
 
-func randomHash() ([20]byte, error) {
- var hash [20]byte
- if _, err := rand.Read(hash[:]); err != nil {
- return [20]byte{}, err
- }
- return hash, nil
-}
-
-func waitAndHandlePeerConnection(address string, myPeerID [20]byte, infoHash [20]byte, logger *logger.Logger) {
- listener, err := net.Listen("tcp", address)
- if err != nil {
- logger.Errorf("Error: %s", err)
- return
- }
- defer listener.Close()
-
- for {
- conn, err := listener.Accept()
- if err != nil {
- logger.Errorf("Error accepting connection: %s", err)
- }
- logger.Debugf("Waiting for handshake message")
- handleConnection(conn, myPeerID, infoHash, logger)
- }
-}
-
-func handleConnection(conn net.Conn, myPeerID [20]byte, infoHash [20]byte, logger *logger.Logger) {
+func handleHandshake(conn net.Conn, params PeerConnectionParams) {
  defer conn.Close()
 
- handshake, err := readHandshake(conn, logger)
+ err := receiveAndSendHandshake(conn, params)
  if err != nil {
- logger.Errorf("error reading handshake: %s", err)
- return
- }
- if !bytes.Equal(handshake.InfoHash[:], infoHash[:]) {
- logger.Errorf("expected infohash %x but got %x", infoHash, handshake.InfoHash)
  return
  }
-
- logger.Debugf("Received handshake: [infohash: %x, peer_id: %x]\n", handshake.InfoHash, handshake.PeerID)
- logger.Debugf("Sending back handshake with peer_id: %x", myPeerID)
- sendHandshake(conn, handshake.InfoHash, myPeerID)
 }

internal/stage_helpers.go

@@ -18,6 +18,29 @@ import (
  "github.com/jackpal/bencode-go"
 )
 
+type ConnectionHandler func(net.Conn, PeerConnectionParams)
+
+type PeerConnectionParams struct {
+ address string
+ myPeerID [20]byte
+ infoHash [20]byte
+ expectedReservedBytes [][]byte
+ myMetadataExtensionID uint8
+ metadataSizeBytes int
+ bitfield []byte
+ magnetLink MagnetTestTorrentInfo
+ logger *logger.Logger
+}
+
+type TrackerParams struct {
+ trackerAddress string
+ peersResponse []byte
+ expectedInfoHash [20]byte
+ fileLengthBytes int
+ logger *logger.Logger
+ myMetadataExtensionID uint8
+}
+
 var samplePieceHashes = []string{
  "ddf33172599fda84f0a209a3034f79f0b8aa5e22",
  "795a618a1ee5275e952843b01a56ae4e142752ef",
@@ -219,10 +242,11 @@ func calculateSHA1(filePath string) (string, error) {
  return hex.EncodeToString(hashBytes), nil
 }
 
-func listenAndServePeersResponse(address string, responseContent []byte, expectedInfoHash [20]byte, fileLengthBytes int, logger *logger.Logger) {
+func listenAndServeTrackerResponse(p TrackerParams) {
+ logger := p.logger
  mux := http.NewServeMux()
  mux.HandleFunc("/announce", func(w http.ResponseWriter, r *http.Request) {
- serveTrackerResponse(w, r, responseContent, expectedInfoHash, fileLengthBytes, logger)
+ serveTrackerResponse(w, r, p.peersResponse, p.expectedInfoHash, p.fileLengthBytes, p.logger)
  })
 
  // Redirect /announce/ to /announce while preserving query parameters 
@@ -237,8 +261,8 @@ func listenAndServePeersResponse(address string, responseContent []byte, expecte
  http.Redirect(w, r, parsedURL.String(), http.StatusMovedPermanently)
  })
 
- logger.Debugf("Server started on port %s...\n", address)
- err := http.ListenAndServe(address, mux)
+ logger.Debugf("Tracker started on address %s...\n", p.trackerAddress)
+ err := http.ListenAndServe(p.trackerAddress, mux)
  if err != nil {
  logger.Errorf("Error: %s", err)
  }
@@ -380,3 +404,69 @@ func createPeersResponse(peerIP string, peerPort int) []byte {
  }
  return buf.Bytes()
 }
+
+func receiveAndSendHandshake(conn net.Conn, peer PeerConnectionParams) (err error) {
+ defer logOnExit(peer.logger, &err)
+
+ logger := peer.logger
+ handshake, err := readHandshake(conn, logger)
+ if err != nil {
+ return fmt.Errorf("error reading handshake: %s", err)
+ }
+
+ 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[:]) {
+ return fmt.Errorf("expected infohash %x but got %x", peer.infoHash, handshake.InfoHash)
+ }
+
+ logger.Debugf("Received handshake: [infohash: %x, peer_id: %x]\n", handshake.InfoHash, handshake.PeerID)
+ logger.Debugf("Sending back handshake with peer_id: %x", peer.myPeerID)
+
+ var reservedBytes [8]byte 
+ copy(reservedBytes[:], peer.expectedReservedBytes[0])
+
+ err = sendHandshake(conn, reservedBytes, handshake.InfoHash, peer.myPeerID)
+ if err != nil {
+ return err
+ }
+ return nil
+}
+
+func waitAndHandlePeerConnection(p PeerConnectionParams, handler ConnectionHandler) {
+ logger := p.logger
+ logger.Debugf("Peer listening on address: %s", p.address)
+ listener, err := net.Listen("tcp", p.address)
+ if err != nil {
+ logger.Errorf("Error: %s", err)
+ return
+ }
+ defer listener.Close()
+
+ for {
+ conn, err := listener.Accept()
+ if err != nil {
+ logger.Errorf("Error accepting connection: %s", err)
+ }
+ handler(conn, p)
+ }
+}
+
+func randomHash() ([20]byte, error) {
+ var hash [20]byte
+ if _, err := rand.Read(hash[:]); err != nil {
+ 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
+}

internal/stage_magnet_helpers.go

@@ -1,5 +1,27 @@
 package internal
 
+import (
+ "encoding/hex"
+ "fmt"
+ "math/rand"
+
+ logger "github.com/codecrafters-io/tester-utils/logger"
+)
+
+type MagnetTestParams struct {
+ TrackerAddress string
+ PeerPort int
+ PeerAddress string
+ PeersResponse []byte
+ ExpectedInfoHash [20]byte
+ ExpectedReservedBytes []byte
+ ExpectedPeerID [20]byte
+ MyMetadataExtensionID uint8
+ MagnetUrlEncoded string
+ MagnetLinkInfo MagnetTestTorrentInfo
+ Logger *logger.Logger
+}
+
 type MagnetTestTorrentInfo struct {
  Filename string
  InfoHashStr string
@@ -53,3 +75,77 @@ var magnetTestTorrents = []MagnetTestTorrentInfo {
  ExpectedSha1: "b1807e3d7920a559df2a2f0f555a404dec66a63e",
  },
 }
+
+func (m *MagnetTestParams) toTrackerParams() TrackerParams {
+ return TrackerParams {
+ trackerAddress: m.TrackerAddress,
+ peersResponse: m.PeersResponse,
+ expectedInfoHash: m.ExpectedInfoHash,
+ fileLengthBytes: m.MagnetLinkInfo.FileLengthBytes,
+ logger: m.Logger,
+ myMetadataExtensionID: m.MyMetadataExtensionID,
+ }
+}
+
+func (m *MagnetTestParams) toPeerConnectionParams() PeerConnectionParams {
+ return PeerConnectionParams {
+ address: m.PeerAddress,
+ myPeerID: m.ExpectedPeerID,
+ infoHash: m.ExpectedInfoHash,
+ expectedReservedBytes: [][]byte{m.ExpectedReservedBytes},
+ myMetadataExtensionID: m.MyMetadataExtensionID,
+ metadataSizeBytes: m.MagnetLinkInfo.MetadataSizeBytes,
+ bitfield: m.MagnetLinkInfo.Bitfield,
+ magnetLink: m.MagnetLinkInfo,
+ logger: m.Logger,
+ }
+}
+
+func NewMagnetTestParams(magnetLink MagnetTestTorrentInfo, logger *logger.Logger) (*MagnetTestParams, error) {
+ params := MagnetTestParams{}
+
+ peerPort, err := findFreePort()
+ if err != nil {
+ return nil, fmt.Errorf("couldn't find free port: %s", err)
+ }
+ params.PeerPort = peerPort
+ params.PeerAddress = fmt.Sprintf("127.0.0.1:%d", peerPort)
+ params.PeersResponse = createPeersResponse("127.0.0.1", peerPort)
+
+ trackerPort, err := findFreePort()
+ if err != nil {
+ return nil, fmt.Errorf("couldn't find free port: %s", err)
+ }
+ trackerAddress := fmt.Sprintf("127.0.0.1:%d", trackerPort)
+ params.TrackerAddress = trackerAddress
+
+ infoHashStr := magnetLink.InfoHashStr
+ params.MagnetUrlEncoded = "magnet:?xt=urn:btih:" + infoHashStr + "&dn=" + magnetLink.Filename + "&tr=http%3A%2F%2F" + trackerAddress + "%2Fannounce"
+
+ infoHash, err := decodeInfoHash(infoHashStr)
+ if err != nil {
+ return nil, fmt.Errorf("error decoding infohash: %v", err)
+ }
+ params.ExpectedInfoHash = infoHash
+
+ expectedPeerID, err := randomHash()
+ if err != nil {
+ return nil, fmt.Errorf("error generating random peer id: %v", err)
+ }
+ params.ExpectedPeerID = expectedPeerID
+ params.ExpectedReservedBytes = []byte{0, 0, 0, 0, 0, 16, 0, 0}
+ params.MyMetadataExtensionID = uint8(rand.Intn(255) + 1)
+ params.MagnetLinkInfo = magnetLink
+ params.Logger = logger
+ return &params, nil
+}
+
+func decodeInfoHash(infoHashStr string) ([20]byte, error) {
+ var infoHash [20]byte
+ decodedBytes, err := hex.DecodeString(infoHashStr)
+ if err != nil {
+ return infoHash, err
+ }
+ copy(infoHash[:], decodedBytes)
+ return infoHash, nil
+}