feat(handshake): handshake

Represents a handshake between peers perfoming serialization
BrianLusina · Jul 18, 2022 · f6dee87 · f6dee87
1 parent edd1683
commit f6dee87
Show file tree

Hide file tree

Showing 4 changed files with 170 additions and 0 deletions.
diff --git a/pkg/handshake/handshake.go b/pkg/handshake/handshake.go
@@ -0,0 +1,28 @@
+package handshake
+
+// Handshake is a special message that a peer uses to identifier itself
+type Handshake struct {
+	Pstr     string
+	InfoHash [20]byte
+	PeerID   [20]byte
+}
+
+func New(infohash, peerID [20]byte) *Handshake {
+	return &Handshake{
+		Pstr:     "BitTorrent protocol",
+		InfoHash: infohash,
+		PeerID:   peerID,
+	}
+}
+
+// Serialize serializes the handshake to a buffer
+func (h *Handshake) Serialize() []byte {
+	buf := make([]byte, len(h.Pstr)+49)
+	buf[0] = byte(len(h.Pstr))
+	curr := 1
+	curr += copy(buf[curr:], []byte(h.Pstr))
+	curr += copy(buf[curr:], make([]byte, 8)) // 8 reserved bytes
+	curr += copy(buf[curr:], h.InfoHash[:])
+	curr += copy(buf[curr:], h.PeerID[:])
+	return buf
+}
diff --git a/pkg/handshake/handshake_test.go b/pkg/handshake/handshake_test.go
@@ -0,0 +1,48 @@
+package handshake
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNew(t *testing.T) {
+	infoHash := [20]byte{134, 212, 200, 0, 36, 164, 105, 190, 76, 80, 188, 90, 16, 44, 247, 23, 128, 49, 0, 116}
+	peerID := [20]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20}
+	h := New(infoHash, peerID)
+	expected := &Handshake{
+		Pstr:     "BitTorrent protocol",
+		InfoHash: [20]byte{134, 212, 200, 0, 36, 164, 105, 190, 76, 80, 188, 90, 16, 44, 247, 23, 128, 49, 0, 116},
+		PeerID:   [20]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20},
+	}
+	assert.Equal(t, expected, h)
+}
+
+func TestSerialize(t *testing.T) {
+	tests := map[string]struct {
+		input  *Handshake
+		output []byte
+	}{
+		"serialize message": {
+			input: &Handshake{
+				Pstr:     "BitTorrent protocol",
+				InfoHash: [20]byte{134, 212, 200, 0, 36, 164, 105, 190, 76, 80, 188, 90, 16, 44, 247, 23, 128, 49, 0, 116},
+				PeerID:   [20]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20},
+			},
+			output: []byte{19, 66, 105, 116, 84, 111, 114, 114, 101, 110, 116, 32, 112, 114, 111, 116, 111, 99, 111, 108, 0, 0, 0, 0, 0, 0, 0, 0, 134, 212, 200, 0, 36, 164, 105, 190, 76, 80, 188, 90, 16, 44, 247, 23, 128, 49, 0, 116, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20},
+		},
+		"different pstr": {
+			input: &Handshake{
+				Pstr:     "BitTorrent protocol, but cooler?",
+				InfoHash: [20]byte{134, 212, 200, 0, 36, 164, 105, 190, 76, 80, 188, 90, 16, 44, 247, 23, 128, 49, 0, 116},
+				PeerID:   [20]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20},
+			},
+			output: []byte{32, 66, 105, 116, 84, 111, 114, 114, 101, 110, 116, 32, 112, 114, 111, 116, 111, 99, 111, 108, 44, 32, 98, 117, 116, 32, 99, 111, 111, 108, 101, 114, 63, 0, 0, 0, 0, 0, 0, 0, 0, 134, 212, 200, 0, 36, 164, 105, 190, 76, 80, 188, 90, 16, 44, 247, 23, 128, 49, 0, 116, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20},
+		},
+	}
+
+	for _, test := range tests {
+		buf := test.input.Serialize()
+		assert.Equal(t, test.output, buf)
+	}
+}
diff --git a/pkg/handshake/utils.go b/pkg/handshake/utils.go
@@ -0,0 +1,42 @@
+package handshake
+
+import (
+	"fmt"
+	"io"
+)
+
+// Read parses a handshake from a stream
+func Read(r io.Reader) (*Handshake, error) {
+	lengthBuf := make([]byte, 1)
+
+	_, err := io.ReadFull(r, lengthBuf)
+	if err != nil {
+		return nil, err
+	}
+
+	pstrLen := int(lengthBuf[0])
+
+	if pstrLen == 0 {
+		err := fmt.Errorf("pstrlen cannot be 0")
+		return nil, err
+	}
+
+	handshakeBuf := make([]byte, 48+pstrLen)
+	_, err = io.ReadFull(r, handshakeBuf)
+	if err != nil {
+		return nil, err
+	}
+
+	var infoHash, peerID [20]byte
+
+	copy(infoHash[:], handshakeBuf[pstrLen+8:pstrLen+8+20])
+	copy(peerID[:], handshakeBuf[pstrLen+8+20:])
+
+	h := Handshake{
+		Pstr:     string(handshakeBuf[0:pstrLen]),
+		InfoHash: infoHash,
+		PeerID:   peerID,
+	}
+
+	return &h, nil
+}
diff --git a/pkg/handshake/utils_test.go b/pkg/handshake/utils_test.go
@@ -0,0 +1,52 @@
+package handshake
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRead(t *testing.T) {
+	tests := map[string]struct {
+		input  []byte
+		output *Handshake
+		fails  bool
+	}{
+		"parse handshake into struct": {
+			input: []byte{19, 66, 105, 116, 84, 111, 114, 114, 101, 110, 116, 32, 112, 114, 111, 116, 111, 99, 111, 108, 0, 0, 0, 0, 0, 0, 0, 0, 134, 212, 200, 0, 36, 164, 105, 190, 76, 80, 188, 90, 16, 44, 247, 23, 128, 49, 0, 116, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20},
+			output: &Handshake{
+				Pstr:     "BitTorrent protocol",
+				InfoHash: [20]byte{134, 212, 200, 0, 36, 164, 105, 190, 76, 80, 188, 90, 16, 44, 247, 23, 128, 49, 0, 116},
+				PeerID:   [20]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20},
+			},
+			fails: false,
+		},
+		"empty": {
+			input:  []byte{},
+			output: nil,
+			fails:  true,
+		},
+		"Not enough bytes": {
+			input:  []byte{19, 66, 105, 116, 84, 111, 114, 114, 101, 110, 116, 32, 112, 114, 111, 116, 111, 99, 111},
+			output: nil,
+			fails:  true,
+		},
+		"pstrlen is 0": {
+			input:  []byte{0, 0, 0},
+			output: nil,
+			fails:  true,
+		},
+	}
+
+	for _, test := range tests {
+		reader := bytes.NewReader(test.input)
+		m, err := Read(reader)
+		if test.fails {
+			assert.NotNil(t, err)
+		} else {
+			assert.Nil(t, err)
+		}
+		assert.Equal(t, test.output, m)
+	}
+}