peer_protocol/protocol.go | 62 ++++++++++++++++++++++++++--------------------------- peer_protocol/protocol_test.go | 14 ++++++++++++++ diff --git a/peer_protocol/protocol.go b/peer_protocol/protocol.go index 75e33cab453abd5db59bbf9a0dd1fef07cc9f4c8..90b01359e6db8b539cb532d8526375e88056d82e 100644 --- a/peer_protocol/protocol.go +++ b/peer_protocol/protocol.go @@ -45,43 +45,41 @@ } func (msg Message) MarshalBinary() (data []byte, err error) { buf := &bytes.Buffer{} - if msg.Keepalive { - data = buf.Bytes() - return - } - err = buf.WriteByte(byte(msg.Type)) - if err != nil { - return - } - switch msg.Type { - case Choke, Unchoke, Interested, NotInterested: - case Have: - err = binary.Write(buf, binary.BigEndian, msg.Index) - case Request, Cancel: - for _, i := range []Integer{msg.Index, msg.Begin, msg.Length} { - err = binary.Write(buf, binary.BigEndian, i) + if !msg.Keepalive { + err = buf.WriteByte(byte(msg.Type)) + if err != nil { + return + } + switch msg.Type { + case Choke, Unchoke, Interested, NotInterested: + case Have: + err = binary.Write(buf, binary.BigEndian, msg.Index) + case Request, Cancel: + for _, i := range []Integer{msg.Index, msg.Begin, msg.Length} { + err = binary.Write(buf, binary.BigEndian, i) + if err != nil { + break + } + } + case Bitfield: + _, err = buf.Write(marshalBitfield(msg.Bitfield)) + case Piece: + for _, i := range []Integer{msg.Index, msg.Begin} { + err = binary.Write(buf, binary.BigEndian, i) + if err != nil { + return + } + } + n, err := buf.Write(msg.Piece) if err != nil { break } - } - case Bitfield: - _, err = buf.Write(marshalBitfield(msg.Bitfield)) - case Piece: - for _, i := range []Integer{msg.Index, msg.Begin} { - err = binary.Write(buf, binary.BigEndian, i) - if err != nil { - return + if n != len(msg.Piece) { + panic(n) } + default: + err = fmt.Errorf("unknown message type: %s", msg.Type) } - n, err := buf.Write(msg.Piece) - if err != nil { - break - } - if n != len(msg.Piece) { - panic(n) - } - default: - err = fmt.Errorf("unknown message type: %s", msg.Type) } data = make([]byte, 4+buf.Len()) binary.BigEndian.PutUint32(data, uint32(buf.Len())) diff --git a/peer_protocol/protocol_test.go b/peer_protocol/protocol_test.go index 169399adebdde884cfeb62a018f859b269382f5d..ca65d32b57e93fd51e3795f4758eb4656ca1e8cf 100644 --- a/peer_protocol/protocol_test.go +++ b/peer_protocol/protocol_test.go @@ -108,3 +108,17 @@ t.Fatalf("expected ErrUnexpectedEOF decoding %q, got %s", stream, err) } } } + +func TestMarshalKeepalive(t *testing.T) { + b, err := (Message{ + Keepalive: true, + }).MarshalBinary() + if err != nil { + t.Fatalf("error marshalling keepalive: %s", err) + } + bs := string(b) + const expected = "\x00\x00\x00\x00" + if bs != expected { + t.Fatalf("marshalled keepalive is %q, expected %q", bs, expected) + } +}