From: Matt Joiner <anacrolix@gmail.com>
Date: Wed, 29 Sep 2021 23:59:18 +0000 (+1000)
Subject: Remove an allocation reading message length
X-Git-Tag: v1.32.0~6
X-Git-Url: http://www.git.stargrave.org/?a=commitdiff_plain;h=60a622ff4b8742bdd65fe5c497f77faa8fbd7324;p=btrtrc.git

Remove an allocation reading message length
---

diff --git a/peer_protocol/decoder.go b/peer_protocol/decoder.go
index 870716c5..decc8e8f 100644
--- 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 2f94e0ee..2800f3ca 100644
--- 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 f0203162..28191ee8 100644
--- 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)