package message import ( "encoding/binary" "errors" "fmt" "io" ) type messageID uint8 // Message ID const ( 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 stores ID and payload of a message type Message struct { ID messageID Payload []byte } // A Bitfield represents the pieces that a peer has type Bitfield []byte // FormatRequest formats the ID and payload for 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, } } // 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, errors.New("Payload too short") } 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 } // 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 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 } func (m *Message) String() 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" case MsgPort: return "Port" default: return fmt.Sprintf("Unknown#%d", m.ID) } } // HasPiece tells if a bitfield has a particular index set func (bf Bitfield) HasPiece(index int) bool { byteIndex := index / 8 offset := index % 8 return bf[byteIndex]>>(7-offset)&1 != 0 }