Remove an allocation reading message length

diff --git a/peer_protocol/decoder.go b/peer_protocol/decoder.go
index 870716c50e1098d6a2455325fe526e3e07b5ee4d..decc8e8ffa5597352f3553f3e3aa17788ddd9783 100644 (file)
--- a/peer_protocol/decoder.go
+++ b/peer_protocol/decoder.go
@@ -20,12 +20,9 @@ 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 = binary.Read(d.R, binary.BigEndian, &length)
+       err = length.ReadFrom(d.R)
        if err != nil {
-               if err != io.EOF {
-                       err = fmt.Errorf("error reading message length: %w", err)
-               }
-               return
+               return fmt.Errorf("reading message length: %w", err)
        }
        if length > d.MaxLength {
                return errors.New("message too long")

diff --git a/peer_protocol/decoder_test.go b/peer_protocol/decoder_test.go
index 2f94e0ee1b81a7afeda41c0eaf2aa05c3e74ce4e..2800f3cab7d572adccd30b51dde51150a11de98e 100644 (file)
--- a/peer_protocol/decoder_test.go
+++ b/peer_protocol/decoder_test.go
@@ -70,7 +70,7 @@ func TestDecodeShortPieceEOF(t *testing.T) {
        var m Message
        require.NoError(t, d.Decode(&m))
        assert.Len(t, m.Piece, 1)
-       assert.Equal(t, io.EOF, d.Decode(&m))
+       assert.ErrorIs(t, d.Decode(&m), io.EOF)
 }
 
 func TestDecodeOverlongPiece(t *testing.T) {

diff --git a/peer_protocol/int.go b/peer_protocol/int.go
index f0203162cd461262849acdefd5b1124066d4684a..28191ee82600d7bca1b4f382674bf2b691c7998e 100644 (file)
--- a/peer_protocol/int.go
+++ b/peer_protocol/int.go
@@ -16,6 +16,19 @@ func (i *Integer) UnmarshalBinary(b []byte) error {
        return i.Read(bytes.NewReader(b))
 }
 
+func (i *Integer) ReadFrom(r io.Reader) error {
+       var b [4]byte
+       n, err := r.Read(b[:])
+       if n == 4 {
+               *i = Integer(binary.BigEndian.Uint32(b[:]))
+               return nil
+       }
+       if err == nil {
+               return io.ErrUnexpectedEOF
+       }
+       return err
+}
+
 // It's perfectly fine to cast these to an int. TODO: Or is it?
 func (i Integer) Int() int {
        return int(i)