improve fuzz tests (#3596)

bluenviron · Aug 1, 2024 · c9a938a · c9a938a
1 parent 59ae3ad
commit c9a938a
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 12 deletions.
diff --git a/internal/protocols/rtmp/amf0/unmarshal_test.go b/internal/protocols/rtmp/amf0/unmarshal_test.go
@@ -322,6 +322,9 @@ func FuzzUnmarshal(f *testing.F) {
 	}
 
 	f.Fuzz(func(_ *testing.T, b []byte) {
-		Unmarshal(b) //nolint:errcheck
+		what, err := Unmarshal(b)
+		if err == nil {
+			Marshal(what) //nolint:errcheck
+		}
 	})
 }
diff --git a/internal/protocols/rtmp/chunk/chunk_test.go b/internal/protocols/rtmp/chunk/chunk_test.go
@@ -160,27 +160,39 @@ func TestChunkMarshal(t *testing.T) {
 func FuzzChunk0Read(f *testing.F) {
 	f.Fuzz(func(_ *testing.T, b []byte) {
 		var chunk Chunk0
-		chunk.Read(bytes.NewReader(b), 65536, false) //nolint:errcheck
+		err := chunk.Read(bytes.NewReader(b), 65536, false)
+		if err == nil {
+			chunk.Marshal(false) //nolint:errcheck
+		}
 	})
 }
 
 func FuzzChunk1Read(f *testing.F) {
 	f.Fuzz(func(_ *testing.T, b []byte) {
 		var chunk Chunk1
-		chunk.Read(bytes.NewReader(b), 65536, false) //nolint:errcheck
+		err := chunk.Read(bytes.NewReader(b), 65536, false)
+		if err == nil {
+			chunk.Marshal(false) //nolint:errcheck
+		}
 	})
 }
 
 func FuzzChunk2Read(f *testing.F) {
 	f.Fuzz(func(_ *testing.T, b []byte) {
 		var chunk Chunk2
-		chunk.Read(bytes.NewReader(b), 65536, false) //nolint:errcheck
+		err := chunk.Read(bytes.NewReader(b), 65536, false)
+		if err == nil {
+			chunk.Marshal(false) //nolint:errcheck
+		}
 	})
 }
 
 func FuzzChunk3Read(f *testing.F) {
 	f.Fuzz(func(_ *testing.T, b []byte) {
 		var chunk Chunk3
-		chunk.Read(bytes.NewReader(b), 65536, true) //nolint:errcheck
+		err := chunk.Read(bytes.NewReader(b), 65536, true)
+		if err == nil {
+			chunk.Marshal(false) //nolint:errcheck
+		}
 	})
 }
diff --git a/internal/protocols/rtmp/handshake/handshake_test.go b/internal/protocols/rtmp/handshake/handshake_test.go
@@ -38,6 +38,7 @@ func TestHandshake(t *testing.T) {
 
 			clientInKey, clientOutKey, err := DoClient(rw, ca == "encrypted", true)
 			require.NoError(t, err)
+
 			<-done
 
 			if ca == "encrypted" {

diff --git a/internal/protocols/rtmp/message/reader_test.go b/internal/protocols/rtmp/message/reader_test.go
@@ -297,9 +297,18 @@ func FuzzReader(f *testing.F) {
 		0x01, 0x00, 0x00, 0x00, 0x88, 0x68, 0x76, 0x63,
 		0x31, 0x01, 0x02, 0x03,
 	})
+
 	f.Fuzz(func(_ *testing.T, b []byte) {
-		bc := bytecounter.NewReader(bytes.NewReader(b))
-		r := NewReader(bc, bc, nil)
-		r.Read() //nolint:errcheck
+		bcr := bytecounter.NewReader(bytes.NewReader(b))
+		r := NewReader(bcr, bcr, nil)
+
+		var buf bytes.Buffer
+		bcw := bytecounter.NewWriter(&buf)
+		w := NewWriter(bcw, bcw, true)
+
+		msg, err := r.Read()
+		if err == nil {
+			w.Write(msg) //nolint:errcheck
+		}
 	})
 }
diff --git a/internal/protocols/rtmp/rawmessage/reader_test.go b/internal/protocols/rtmp/rawmessage/reader_test.go
@@ -270,14 +270,20 @@ func TestReaderAcknowledge(t *testing.T) {
 
 func FuzzReader(f *testing.F) {
 	f.Fuzz(func(_ *testing.T, b []byte) {
-		br := bytecounter.NewReader(bytes.NewReader(b))
-		r := NewReader(br, br, func(_ uint32) error {
+		bcr := bytecounter.NewReader(bytes.NewReader(b))
+		r := NewReader(bcr, bcr, func(_ uint32) error {
 			return nil
 		})
 
+		var buf bytes.Buffer
+		bcw := bytecounter.NewWriter(&buf)
+		w := NewWriter(bcw, bcw, true)
+
 		for {
-			_, err := r.Read()
-			if err != nil {
+			msg, err := r.Read()
+			if err == nil {
+				w.Write(msg) //nolint:errcheck
+			} else {
 				break
 			}
 		}