Fix H264Reader buffer overflow

Fixes #1734
pion · Jun 28, 2021 · 299e04a · 299e04a
1 parent 1ba27d8
commit 299e04a
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/pkg/media/h264reader/h264reader.go b/pkg/media/h264reader/h264reader.go
@@ -167,8 +167,10 @@ func (reader *H264Reader) processByte(readByte byte) (nalFound bool) {
 				countOfConsecutiveZeroBytesInPrefix = 3
 			}
 			nalUnitLength := len(reader.nalBuffer) - countOfConsecutiveZeroBytesInPrefix
-			reader.nalBuffer = reader.nalBuffer[0:nalUnitLength]
-			nalFound = true
+			if nalUnitLength > 0 {
+				reader.nalBuffer = reader.nalBuffer[0:nalUnitLength]
+				nalFound = true
+			}
 		} else {
 			reader.countOfConsecutiveZeroBytes = 0
 		}

diff --git a/pkg/media/h264reader/h264reader_test.go b/pkg/media/h264reader/h264reader_test.go
@@ -98,3 +98,24 @@ func TestSkipSEI(t *testing.T) {
 	assert.Nil(err)
 	assert.Equal(byte(0xAB), nal.Data[0])
 }
+
+func TestIssue1734_NextNal(t *testing.T) {
+	tt := [...][]byte{
+		[]byte("\x00\x00\x010\x00\x00\x01\x00\x00\x01"),
+		[]byte("\x00\x00\x00\x01\x00\x00\x01"),
+	}
+
+	for _, cur := range tt {
+		r, err := NewReader(bytes.NewReader(cur))
+		assert.NoError(t, err)
+
+		// Just make sure it doesn't crash
+		for {
+			nal, err := r.NextNAL()
+
+			if err != nil || nal == nil {
+				break
+			}
+		}
+	}
+}