codec/h264/lex.go: checking NAL packet lengths

codec/h264/lex.go

@@ -59,6 +59,13 @@ const (
 	typeFUB = 29
 )
 
+// Min NAL lengths.
+const (
+	minSingleNALLen = 1
+	minSTAPALen     = 4
+	minFUALen       = 2
+)
+
 var noDelay = make(chan time.Time)
 
 func init() {
@@ -205,6 +212,10 @@ func (l *RTPLexer) Lex(dst io.Writer, src io.Reader, delay time.Duration) error
 		}
 
 		if typeSingleNALULowBound <= nalType && nalType <= typeSingleNALUHighBound {
+			// If len too small, ignore.
+			if len(payload) < minSingleNALLen {
+				continue
+			}
 			l.writeWithPrefix(payload)
 		} else {
 			switch nalType {
@@ -241,6 +252,11 @@ func (l *RTPLexer) Lex(dst io.Writer, src io.Reader, delay time.Duration) error
 // handleSTAPA parses NAL units from an aggregation packet and writes
 // them to the Lexers buffer buf.
 func (l *RTPLexer) handleSTAPA(d []byte) {
+	// If the length is too small, ignore.
+	if len(d) < minSTAPALen {
+		return
+	}
+
 	for i := 1; i < len(d); {
 		size := int(binary.BigEndian.Uint16(d[i:]))
 
@@ -258,6 +274,11 @@ func (l *RTPLexer) handleSTAPA(d []byte) {
 // handleFUA parses NAL units from fragmentation packets and writes
 // them to the Lexer's buf.
 func (l *RTPLexer) handleFUA(d []byte) {
+	// If length is too small, ignore.
+	if len(d) < minFUALen {
+		return
+	}
+
 	// Get start and end indiciators from FU header.
 	const FUHeadIdx = 1
 	start := d[FUHeadIdx]&0x80 != 0