Implement Message.String()

peer/message.go

@@ -2,26 +2,29 @@ package peer
 
 import (
 	"encoding/binary"
+	"fmt"
 	"io"
 )
 
+type messageID uint8
+
 // Message ID
 const (
-	MsgChoke         uint8 = 0
-	MsgUnchoke       uint8 = 1
-	MsgInterested    uint8 = 2
-	MsgNotInterested uint8 = 3
-	MsgHave          uint8 = 4
-	MsgBitfield      uint8 = 5
-	MsgRequest       uint8 = 6
-	MsgPiece         uint8 = 7
-	MsgCancel        uint8 = 8
-	MsgPort          uint8 = 9
+	MsgChoke         messageID = 0
+	MsgUnchoke       messageID = 1
+	MsgInterested    messageID = 2
+	MsgNotInterested messageID = 3
+	MsgHave          messageID = 4
+	MsgBitfield      messageID = 5
+	MsgRequest       messageID = 6
+	MsgPiece         messageID = 7
+	MsgCancel        messageID = 8
+	MsgPort          messageID = 9
 )
 
 // Message m
 type Message struct {
-	ID      uint8
+	ID      messageID
 	Payload []byte
 }
 
@@ -48,6 +51,7 @@ func Read(r io.Reader) (*Message, error) {
 	}
 	length := binary.BigEndian.Uint32(lengthBuf)
 
+	// keep-alive message
 	if length == 0 {
 		return nil, nil
 	}
@@ -59,9 +63,43 @@ func Read(r io.Reader) (*Message, error) {
 	}
 
 	m := Message{
-		ID:      messageBuf[0],
+		ID:      messageID(messageBuf[0]),
 		Payload: messageBuf[1:],
 	}
 
 	return &m, nil
 }
+
+func (m *Message) String() string {
+	if m == nil {
+		return "KeepAlive"
+	}
+
+	var idName string
+	switch m.ID {
+	case MsgChoke:
+		idName = "Choke"
+	case MsgUnchoke:
+		idName = "Unchoke"
+	case MsgInterested:
+		idName = "Interested"
+	case MsgNotInterested:
+		idName = "NotInterested"
+	case MsgHave:
+		idName = "Have"
+	case MsgBitfield:
+		idName = "Bitfield"
+	case MsgRequest:
+		idName = "Request"
+	case MsgPiece:
+		idName = "Piece"
+	case MsgCancel:
+		idName = "Cancel"
+	case MsgPort:
+		idName = "Port"
+	default:
+		idName = fmt.Sprintf("Unknown#%d", m.ID)
+	}
+
+	return fmt.Sprintf("%s\t[% x]", idName, m.Payload)
+}

peer/message_test.go

@@ -67,3 +67,28 @@ func TestRead(t *testing.T) {
 		assert.Equal(t, test.output, m)
 	}
 }
+
+func TestString(t *testing.T) {
+	tests := []struct {
+		input  *Message
+		output string
+	}{
+		{nil, "KeepAlive"},
+		{&Message{MsgChoke, []byte{1, 2, 3}}, "Choke\t[01 02 03]"},
+		{&Message{MsgUnchoke, []byte{1, 2, 3}}, "Unchoke\t[01 02 03]"},
+		{&Message{MsgInterested, []byte{1, 2, 3}}, "Interested\t[01 02 03]"},
+		{&Message{MsgNotInterested, []byte{1, 2, 3}}, "NotInterested\t[01 02 03]"},
+		{&Message{MsgHave, []byte{1, 2, 3}}, "Have\t[01 02 03]"},
+		{&Message{MsgBitfield, []byte{1, 2, 3}}, "Bitfield\t[01 02 03]"},
+		{&Message{MsgRequest, []byte{1, 2, 3}}, "Request\t[01 02 03]"},
+		{&Message{MsgPiece, []byte{1, 2, 3}}, "Piece\t[01 02 03]"},
+		{&Message{MsgCancel, []byte{1, 2, 3}}, "Cancel\t[01 02 03]"},
+		{&Message{MsgPort, []byte{1, 2, 3}}, "Port\t[01 02 03]"},
+		{&Message{99, []byte{1, 2, 3}}, "Unknown#99\t[01 02 03]"},
+	}
+
+	for _, test := range tests {
+		s := test.input.String()
+		assert.Equal(t, test.output, s)
+	}
+}