generated from codecrafters-io/tester-template
-
Notifications
You must be signed in to change notification settings - Fork 3
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
add test for send extension handshake
- Loading branch information
Showing
3 changed files
with
289 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,112 @@ | ||
// Helper methods to read and send BitTorrent messages | ||
package internal | ||
|
||
import ( | ||
"bytes" | ||
"encoding/binary" | ||
"io" | ||
"net" | ||
|
||
logger "github.com/codecrafters-io/tester-utils/logger" | ||
"github.com/jackpal/bencode-go" | ||
) | ||
|
||
type messageID uint8 | ||
type Message struct { | ||
ID messageID | ||
Payload []byte | ||
} | ||
|
||
const ( | ||
HandshakeExtendedID uint8 = 0 | ||
RequestMetadataExtensionMsgType uint8 = 0 | ||
DataMetadataExtensionMsgType uint8 = 1 | ||
|
||
MsgBitfield messageID = 5 | ||
MsgExtended messageID = 20 | ||
) | ||
|
||
func sendBitfieldMessage(conn net.Conn, payload []byte, logger *logger.Logger) (err error) { | ||
defer logOnExit(logger, &err) | ||
|
||
logger.Debugln("Sending bitfield message") | ||
req := Message { ID: MsgBitfield, Payload: payload } | ||
serialized := req.Serialize() | ||
_, err = conn.Write(serialized) | ||
return err | ||
} | ||
|
||
// Serialize serializes a message into a buffer of the form | ||
// <length prefix><message ID><payload> | ||
// Interprets `nil` as a keep-alive message | ||
func (m *Message) Serialize() []byte { | ||
if m == nil { | ||
return make([]byte, 4) | ||
} | ||
length := uint32(len(m.Payload) + 1) // +1 for id | ||
buf := make([]byte, 4+length) | ||
binary.BigEndian.PutUint32(buf[0:4], length) | ||
buf[4] = byte(m.ID) | ||
copy(buf[5:], m.Payload) | ||
return buf | ||
} | ||
|
||
func sendExtensionHandshake(conn net.Conn, metadataID uint8, metadataSize int, logger *logger.Logger) (err error) { | ||
defer logOnExit(logger, &err) | ||
|
||
logger.Debugln("Sending extension handshake") | ||
req := createExtensionHandshake(metadataID, metadataSize, logger) | ||
serialized := req.Serialize() | ||
_, err = conn.Write(serialized) | ||
return err | ||
} | ||
|
||
func createExtensionHandshake(metadataID uint8, metadataSize int, logger *logger.Logger) *Message { | ||
dict := make(map[string]interface{}) | ||
inner := make(map[string]int64) | ||
inner["ut_metadata"] = int64(metadataID) | ||
dict["m"] = inner | ||
dict["metadata_size"] = metadataSize | ||
var buf bytes.Buffer | ||
err := bencode.Marshal(&buf, dict) | ||
if err != nil { | ||
logger.Errorf("Error encoding: %v", err) | ||
} | ||
payload := formatExtendedPayload(buf, HandshakeExtendedID) | ||
return &Message{ID: MsgExtended, Payload: payload} | ||
} | ||
|
||
func formatExtendedPayload(buf bytes.Buffer, extensionId uint8) []byte { | ||
payload := make([]byte, 1+buf.Len()) | ||
payload[0] = uint8(extensionId) | ||
copy(payload[1:], buf.Bytes()) | ||
return payload | ||
} | ||
|
||
// Read parses a message from a stream. Returns `nil` on keep-alive message | ||
func readMessage(r io.Reader) (*Message, error) { | ||
lengthBuf := make([]byte, 4) | ||
_, err := io.ReadFull(r, lengthBuf) | ||
if err != nil { | ||
return nil, err | ||
} | ||
length := binary.BigEndian.Uint32(lengthBuf) | ||
|
||
// keep-alive message | ||
if length == 0 { | ||
return nil, nil | ||
} | ||
|
||
messageBuf := make([]byte, length) | ||
_, err = io.ReadFull(r, messageBuf) | ||
if err != nil { | ||
return nil, err | ||
} | ||
|
||
m := Message{ | ||
ID: messageID(messageBuf[0]), | ||
Payload: messageBuf[1:], | ||
} | ||
|
||
return &m, nil | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,73 @@ | ||
package internal | ||
|
||
import ( | ||
"errors" | ||
"fmt" | ||
"net" | ||
|
||
"github.com/codecrafters-io/tester-utils/test_case_harness" | ||
) | ||
|
||
var handshakeChannel = make(chan bool) | ||
|
||
func testMagnetSendExtendedHandshake(stageHarness *test_case_harness.TestCaseHarness) error { | ||
initRandom() | ||
|
||
logger := stageHarness.Logger | ||
executable := stageHarness.Executable | ||
|
||
magnetLink := randomMagnetLink() | ||
params, err := NewMagnetTestParams(magnetLink, logger) | ||
if err != nil { | ||
return err | ||
} | ||
|
||
go listenAndServeTrackerResponse(params.toTrackerParams()) | ||
go waitAndHandlePeerConnection(params.toPeerConnectionParams(), handleSendExtensionHandshake) | ||
|
||
logger.Infof("Running ./your_bittorrent.sh magnet_handshake %q", params.MagnetUrlEncoded) | ||
result, err := executable.Run("magnet_handshake", params.MagnetUrlEncoded) | ||
if err != nil { | ||
return err | ||
} | ||
|
||
if err = assertExitCode(result, 0); err != nil { | ||
return err | ||
} | ||
|
||
expected := fmt.Sprintf("Peer Metadata Extension ID: %d\n", params.MyMetadataExtensionID) | ||
|
||
if err = assertStdoutContains(result, expected); err != nil { | ||
return err | ||
} | ||
|
||
success := <-handshakeChannel | ||
if success { | ||
return nil | ||
} | ||
|
||
return errors.New("extension handshake was not received") | ||
} | ||
|
||
func handleSendExtensionHandshake(conn net.Conn, params PeerConnectionParams) { | ||
defer conn.Close() | ||
logger := params.logger | ||
|
||
if err := receiveAndSendHandshake(conn, params); err != nil { | ||
return | ||
} | ||
|
||
if err := sendBitfieldMessage(conn, params.bitfield, logger); err != nil { | ||
return | ||
} | ||
|
||
if err := sendExtensionHandshake(conn, params.myMetadataExtensionID, params.metadataSizeBytes, logger); err != nil { | ||
return | ||
} | ||
|
||
if _, err := receiveAndAssertExtensionHandshake(conn, logger); err != nil { | ||
return | ||
} | ||
|
||
handshakeChannel <- true | ||
} |