package message import ( "bufio" "encoding/binary" "fmt" "io" ) 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 ) // Message stores ID and payload of a message type Message struct { ID messageID Payload []byte } // 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} } // 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 } // ParseHave parses a HAVE message func ParseHave(msg *Message) (int, error) { if msg.ID != MsgHave { return 0, fmt.Errorf("Expected HAVE (ID %d), got 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 } // Serialize serializes a message into a buffer of the form // // 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 } // Read parses a message from a stream. Returns `nil` on keep-alive message func Read(r *bufio.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 } 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)) }