diff --git a/peer_protocol/decoder.go b/peer_protocol/decoder.go index decc8e8f..cba8dd08 100644 --- a/peer_protocol/decoder.go +++ b/peer_protocol/decoder.go @@ -20,7 +20,7 @@ type Decoder struct { // io.EOF is returned if the source terminates cleanly on a message boundary. func (d *Decoder) Decode(msg *Message) (err error) { var length Integer - err = length.ReadFrom(d.R) + err = length.Read(d.R) if err != nil { return fmt.Errorf("reading message length: %w", err) } diff --git a/peer_protocol/int.go b/peer_protocol/int.go index 28191ee8..80fb6624 100644 --- a/peer_protocol/int.go +++ b/peer_protocol/int.go @@ -1,27 +1,27 @@ package peer_protocol import ( - "bytes" "encoding/binary" "io" + + "github.com/pkg/errors" ) type Integer uint32 -func (i *Integer) Read(r io.Reader) error { - return binary.Read(r, binary.BigEndian, i) -} - func (i *Integer) UnmarshalBinary(b []byte) error { - return i.Read(bytes.NewReader(b)) + if len(b) != 4 { + return errors.New("expected 4 bytes") + } + *i = Integer(binary.BigEndian.Uint32(b)) + return nil } -func (i *Integer) ReadFrom(r io.Reader) error { +func (i *Integer) Read(r io.Reader) error { var b [4]byte n, err := r.Read(b[:]) if n == 4 { - *i = Integer(binary.BigEndian.Uint32(b[:])) - return nil + return i.UnmarshalBinary(b[:]) } if err == nil { return io.ErrUnexpectedEOF