feat(message): torrent message representation

pkg/message/message.go

@@ -0,0 +1,85 @@
+package message
+
+import (
+	"encoding/binary"
+	"fmt"
+)
+
+type messageID uint8
+
+const (
+	// MsgChoke chokes the receiver
+	MsgChoke messageID = 0
+	// MsgUnchoke unchokes the receiver
+	MsgUnchoke messageID = 1
+	// MsgInterested expresses interest in receiving data
+	MsgInterested messageID = 2
+	// MsgNotInterested expresses disinterest in receiving data
+	MsgNotInterested messageID = 3
+	// MsgHave alerts the receiver that the sender has downloaded a piece
+	MsgHave messageID = 4
+	// MsgBitfield encodes which pieces that the sender has downloaded
+	MsgBitfield messageID = 5
+	// MsgRequest requests a block of data from the receiver
+	MsgRequest messageID = 6
+	// MsgPiece delivers a block of data to fulfill a request
+	MsgPiece messageID = 7
+	// MsgCancel cancels a request
+	MsgCancel messageID = 8
+)
+
+type Message struct {
+	ID      messageID
+	Payload []byte
+}
+
+// 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 (m *Message) name() string {
+	if m == nil {
+		return "KeepAlive"
+	}
+	switch m.ID {
+	case MsgChoke:
+		return "Choke"
+	case MsgUnchoke:
+		return "Unchoke"
+	case MsgInterested:
+		return "Interested"
+	case MsgNotInterested:
+		return "NotInterested"
+	case MsgHave:
+		return "Have"
+	case MsgBitfield:
+		return "Bitfield"
+	case MsgRequest:
+		return "Request"
+	case MsgPiece:
+		return "Piece"
+	case MsgCancel:
+		return "Cancel"
+	default:
+		return fmt.Sprintf("Unknown#%d", m.ID)
+	}
+}
+
+func (m *Message) String() string {
+	if m == nil {
+		return m.name()
+	}
+	return fmt.Sprintf("%s [%d]", m.name(), len(m.Payload))
+}

pkg/message/message_test.go

@@ -0,0 +1,52 @@
+package message
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSerialize(t *testing.T) {
+	tests := map[string]struct {
+		input  *Message
+		output []byte
+	}{
+		"serialize message": {
+			input:  &Message{ID: MsgHave, Payload: []byte{1, 2, 3, 4}},
+			output: []byte{0, 0, 0, 5, 4, 1, 2, 3, 4},
+		},
+		"serialize keep-alive": {
+			input:  nil,
+			output: []byte{0, 0, 0, 0},
+		},
+	}
+
+	for _, test := range tests {
+		buf := test.input.Serialize()
+		assert.Equal(t, test.output, buf)
+	}
+}
+
+func TestString(t *testing.T) {
+	tests := []struct {
+		input  *Message
+		output string
+	}{
+		{nil, "KeepAlive"},
+		{&Message{MsgChoke, []byte{1, 2, 3}}, "Choke [3]"},
+		{&Message{MsgUnchoke, []byte{1, 2, 3}}, "Unchoke [3]"},
+		{&Message{MsgInterested, []byte{1, 2, 3}}, "Interested [3]"},
+		{&Message{MsgNotInterested, []byte{1, 2, 3}}, "NotInterested [3]"},
+		{&Message{MsgHave, []byte{1, 2, 3}}, "Have [3]"},
+		{&Message{MsgBitfield, []byte{1, 2, 3}}, "Bitfield [3]"},
+		{&Message{MsgRequest, []byte{1, 2, 3}}, "Request [3]"},
+		{&Message{MsgPiece, []byte{1, 2, 3}}, "Piece [3]"},
+		{&Message{MsgCancel, []byte{1, 2, 3}}, "Cancel [3]"},
+		{&Message{99, []byte{1, 2, 3}}, "Unknown#99 [3]"},
+	}
+
+	for _, test := range tests {
+		s := test.input.String()
+		assert.Equal(t, test.output, s)
+	}
+}

pkg/message/utils.go

@@ -0,0 +1,92 @@
+package message
+
+import (
+	"encoding/binary"
+	"fmt"
+	"io"
+)
+
+// FormatRequest creates a REQUEST Message
+func FormatRequest(index, begin, length int) *Message {
+	payload := make([]byte, 12)
+	binary.BigEndian.PutUint32(payload[0:4], uint32(index))
+	binary.BigEndian.PutUint32(payload[4:8], uint32(begin))
+	binary.BigEndian.PutUint32(payload[8:12], uint32(length))
+	return &Message{ID: MsgRequest, Payload: payload}
+}
+
+// FormatHave creates a HAVE message
+func FormatHave(index int) *Message {
+	payload := make([]byte, 4)
+	binary.BigEndian.PutUint32(payload, uint32(index))
+	return &Message{ID: MsgHave, Payload: payload}
+}
+
+// ParseHave parse a HAVE message
+func ParseHave(msg *Message) (int, error) {
+	if msg.ID != MsgHave {
+		return 0, fmt.Errorf("Expected HAVE (ID %d), go ID %d", MsgHave, msg.ID)
+	}
+
+	if len(msg.Payload) != 4 {
+		return 0, fmt.Errorf("Expected payload length 4, got length %d", len(msg.Payload))
+	}
+
+	index := int(binary.BigEndian.Uint32(msg.Payload))
+	return index, nil
+}
+
+// ParsePiece parses a PIECE message and copies its payload into a buffer
+func ParsePiece(index int, buf []byte, msg *Message) (int, error) {
+	if msg.ID != MsgPiece {
+		return 0, fmt.Errorf("Expected Piece (ID %d, got ID %d", MsgPiece, msg.ID)
+	}
+	if len(msg.Payload) < 8 {
+		return 0, fmt.Errorf("Payload too short. %d < 8", len(msg.Payload))
+	}
+
+	parsedIndex := int(binary.BigEndian.Uint32(msg.Payload[0:4]))
+	if parsedIndex != index {
+		return 0, fmt.Errorf("Expected index %d, got %d", index, parsedIndex)
+	}
+
+	begin := int(binary.BigEndian.Uint32(msg.Payload[4:8]))
+	if begin >= len(buf) {
+		return 0, fmt.Errorf("Begin offset too high. %d >= %d", begin, len(buf))
+	}
+
+	data := msg.Payload[8:]
+	if begin+len(data) > len(buf) {
+		return 0, fmt.Errorf("Data too long [%d] for offset %d with length %d", len(data), begin, len(buf))
+	}
+	copy(buf[begin:], data)
+	return len(data), nil
+}
+
+// Read parses a message from a stream. Returns `nil` on keep-alive message
+func Read(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
+}