From b69e6dff01f830506456e5510c05396b92723c09 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Thu, 30 Sep 2021 09:42:54 +1000
Subject: [PATCH] Improvements to decoder fuzzing

---
 peer_protocol/decoder.go           |  3 +--
 peer_protocol/decoder_fuzz_test.go | 32 +++++++++++++++++++++++++-----
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/peer_protocol/decoder.go b/peer_protocol/decoder.go
index 3b9d6e9c..870716c5 100644
--- a/peer_protocol/decoder.go
+++ b/peer_protocol/decoder.go
@@ -18,13 +18,12 @@ type Decoder struct {
 }
 
 // io.EOF is returned if the source terminates cleanly on a message boundary.
-// TODO: Is that before or after the message?
 func (d *Decoder) Decode(msg *Message) (err error) {
 	var length Integer
 	err = binary.Read(d.R, binary.BigEndian, &length)
 	if err != nil {
 		if err != io.EOF {
-			err = fmt.Errorf("error reading message length: %s", err)
+			err = fmt.Errorf("error reading message length: %w", err)
 		}
 		return
 	}
diff --git a/peer_protocol/decoder_fuzz_test.go b/peer_protocol/decoder_fuzz_test.go
index 0fed779c..2fd357e4 100644
--- a/peer_protocol/decoder_fuzz_test.go
+++ b/peer_protocol/decoder_fuzz_test.go
@@ -3,7 +3,11 @@ package peer_protocol
 import (
 	"bufio"
 	"bytes"
+	"errors"
+	"io"
 	"testing"
+
+	qt "github.com/frankban/quicktest"
 )
 
 func FuzzDecoder(f *testing.F) {
@@ -11,13 +15,31 @@ func FuzzDecoder(f *testing.F) {
 	f.Add([]byte("\x00\x00\x00\x01\x00"))
 	f.Add([]byte("\x00\x00\x00\x03\x14\x00"))
 	f.Fuzz(func(t *testing.T, b []byte) {
+		c := qt.New(t)
 		d := Decoder{
-			R: bufio.NewReader(bytes.NewReader(b)),
+			R:         bufio.NewReader(bytes.NewReader(b)),
+			MaxLength: 0x100,
+		}
+		var ms []Message
+		for {
+			var m Message
+			err := d.Decode(&m)
+			t.Log(err)
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			if err == nil {
+				c.Assert(m, qt.Not(qt.Equals), Message{})
+				ms = append(ms, m)
+				continue
+			} else {
+				t.Skip(err)
+			}
 		}
-		var m Message
-		err := d.Decode(&m)
-		if err != nil {
-			t.Skip(err)
+		var buf bytes.Buffer
+		for _, m := range ms {
+			buf.Write(m.MustMarshalBinary())
 		}
+		c.Assert(buf.Bytes(), qt.DeepEquals, b)
 	})
 }
-- 
2.48.1