From 8c26f5f00a5a19ab6892c63f73408a76acf8deb0 Mon Sep 17 00:00:00 2001
From: Alessandro Ros <aler9.dev@gmail.com>
Date: Sun, 5 May 2024 11:19:35 +0200
Subject: [PATCH] improve check on H264 padding (#559)

---
 pkg/format/rtph264/decoder.go                       | 13 +++++++++++--
 .../testdata/fuzz/FuzzDecoder/3b0600afabf53c93      |  3 +++
 pkg/format/rtph265/decoder.go                       |  4 ----
 3 files changed, 14 insertions(+), 6 deletions(-)
 create mode 100644 pkg/format/rtph264/testdata/fuzz/FuzzDecoder/3b0600afabf53c93

diff --git a/pkg/format/rtph264/decoder.go b/pkg/format/rtph264/decoder.go
index 9590262b..c047c1d4 100644
--- a/pkg/format/rtph264/decoder.go
+++ b/pkg/format/rtph264/decoder.go
@@ -29,6 +29,15 @@ func joinFragments(fragments [][]byte, size int) []byte {
 	return ret
 }
 
+func isAllZero(buf []byte) bool {
+	for _, b := range buf {
+		if b != 0 {
+			return false
+		}
+	}
+	return true
+}
+
 // Decoder is a RTP/H264 decoder.
 // Specification: https://datatracker.ietf.org/doc/html/rfc6184
 type Decoder struct {
@@ -125,8 +134,8 @@ func (d *Decoder) decodeNALUs(pkt *rtp.Packet) ([][]byte, error) {
 			size := uint16(payload[0])<<8 | uint16(payload[1])
 			payload = payload[2:]
 
-			// avoid final padding
-			if size == 0 {
+			// discard padding
+			if size == 0 && isAllZero(payload) {
 				break
 			}
 
diff --git a/pkg/format/rtph264/testdata/fuzz/FuzzDecoder/3b0600afabf53c93 b/pkg/format/rtph264/testdata/fuzz/FuzzDecoder/3b0600afabf53c93
new file mode 100644
index 00000000..89145c16
--- /dev/null
+++ b/pkg/format/rtph264/testdata/fuzz/FuzzDecoder/3b0600afabf53c93
@@ -0,0 +1,3 @@
+go test fuzz v1
+[]byte("0")
+[]byte("8\x00\x000")
diff --git a/pkg/format/rtph265/decoder.go b/pkg/format/rtph265/decoder.go
index 27769ad0..aa2caa4c 100644
--- a/pkg/format/rtph265/decoder.go
+++ b/pkg/format/rtph265/decoder.go
@@ -75,10 +75,6 @@ func (d *Decoder) decodeNALUs(pkt *rtp.Packet) ([][]byte, error) {
 			size := uint16(payload[0])<<8 | uint16(payload[1])
 			payload = payload[2:]
 
-			if size == 0 {
-				break
-			}
-
 			if int(size) > len(payload) {
 				return nil, fmt.Errorf("invalid aggregation unit (invalid size)")
 			}