h264: apply NALU count limit during Annex-B and AVCC unmarshaling

bluenviron · Oct 9, 2022 · e97d290 · e97d290
1 parent 74f941b
commit e97d290
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 1 deletion.
diff --git a/pkg/h264/annexb.go b/pkg/h264/annexb.go
@@ -46,6 +46,11 @@ outer:
 		}
 	}
 
+	if (n + 1) > MaxNALUsPerGroup {
+		return nil, fmt.Errorf("number of NALUs contained inside a single group (%d) is too big (maximum is %d)",
+			n+1, MaxNALUsPerGroup)
+	}
+
 	ret := make([][]byte, n+1)
 	pos := 0
 	start = initZeroCount + 1

diff --git a/pkg/h264/avcc.go b/pkg/h264/avcc.go
@@ -26,6 +26,11 @@ func AVCCUnmarshal(buf []byte) ([][]byte, error) {
 			return nil, fmt.Errorf("NALU size (%d) is too big (maximum is %d)", bl-pos, MaxNALUSize)
 		}
 
+		if (len(ret) + 1) > MaxNALUsPerGroup {
+			return nil, fmt.Errorf("number of NALUs contained inside a single group (%d) is too big (maximum is %d)",
+				len(ret)+1, MaxNALUsPerGroup)
+		}
+
 		ret = append(ret, buf[pos:pos+le])
 		pos += le
 

diff --git a/pkg/rtph264/decoder.go b/pkg/rtph264/decoder.go
@@ -200,7 +200,7 @@ func (d *Decoder) DecodeUntilMarker(pkt *rtp.Packet) ([][]byte, time.Duration, e
 		return nil, 0, err
 	}
 
-	if (len(d.naluBuffer) + len(nalus)) >= h264.MaxNALUsPerGroup {
+	if (len(d.naluBuffer) + len(nalus)) > h264.MaxNALUsPerGroup {
 		return nil, 0, fmt.Errorf("number of NALUs contained inside a single group (%d) is too big (maximum is %d)",
 			len(d.naluBuffer)+len(nalus), h264.MaxNALUsPerGroup)
 	}