add an upper limit on AU/NALU sizes in order to prevent RAM exhaustion

pkg/rtpaac/decoder.go

@@ -133,8 +133,14 @@ func (d *Decoder) Decode(pkt *rtp.Packet) ([][]byte, time.Duration, error) {
 		return nil, 0, fmt.Errorf("payload is too short")
 	}
 
-	d.fragmentedParts = append(d.fragmentedParts, payload)
 	d.fragmentedSize += len(payload)
+	if d.fragmentedSize > maxAUSize {
+		d.fragmentedParts = d.fragmentedParts[:0]
+		d.fragmentedMode = false
+		return nil, 0, fmt.Errorf("AU size (%d) is too big (maximum is %d)", d.fragmentedSize, maxAUSize)
+	}
+
+	d.fragmentedParts = append(d.fragmentedParts, payload)
 
 	if !pkt.Header.Marker {
 		return nil, 0, ErrMorePacketsNeeded

pkg/rtpaac/encoder.go

@@ -8,10 +8,6 @@ import (
 	"github.com/pion/rtp"
 )
 
-const (
-	rtpVersion = 0x02
-)
-
 func randUint32() uint32 {
 	var b [4]byte
 	rand.Read(b[:])

pkg/rtpaac/rtpaac.go

@@ -1,2 +1,9 @@
 // Package rtpaac contains a RTP/AAC decoder and encoder.
 package rtpaac
+
+const (
+	rtpVersion = 0x02
+
+	// i've never seen a 5kbit AU, but anyway....
+	maxAUSize = 5 * 1024
+)

pkg/rtph264/decoder.go

@@ -35,7 +35,7 @@ type Decoder struct {
 
 // Init initializes the decoder
 func (d *Decoder) Init() {
-	d.timeDecoder = rtptimedec.New(90000)
+	d.timeDecoder = rtptimedec.New(rtpClockRate)
 }
 
 // Decode decodes NALUs from a RTP/H264 packet.
@@ -99,9 +99,9 @@ func (d *Decoder) Decode(pkt *rtp.Packet) ([][]byte, time.Duration, error) {
 
 			nri := (pkt.Payload[0] >> 5) & 0x03
 			typ := pkt.Payload[1] & 0x1F
+			d.fragmentedSize = len(pkt.Payload) - 1
 			d.fragmentedParts = append(d.fragmentedParts, []byte{(nri << 5) | typ})
 			d.fragmentedParts = append(d.fragmentedParts, pkt.Payload[2:])
-			d.fragmentedSize = len(pkt.Payload) - 1
 			d.fragmentedMode = true
 
 			d.firstPacketReceived = true
@@ -138,8 +138,14 @@ func (d *Decoder) Decode(pkt *rtp.Packet) ([][]byte, time.Duration, error) {
 		return nil, 0, fmt.Errorf("invalid FU-A packet (decoded two starting packets in a row)")
 	}
 
-	d.fragmentedParts = append(d.fragmentedParts, pkt.Payload[2:])
 	d.fragmentedSize += len(pkt.Payload[2:])
+	if d.fragmentedSize > maxNALUSize {
+		d.fragmentedParts = d.fragmentedParts[:0]
+		d.fragmentedMode = false
+		return nil, 0, fmt.Errorf("NALU size (%d) is too big (maximum is %d)", d.fragmentedSize, maxNALUSize)
+	}
+
+	d.fragmentedParts = append(d.fragmentedParts, pkt.Payload[2:])
 
 	end := (pkt.Payload[1] >> 6) & 0x01
 	if end != 1 {

pkg/rtph264/encoder.go

@@ -8,11 +8,6 @@ import (
 	"github.com/pion/rtp"
 )
 
-const (
-	rtpVersion   = 0x02
-	rtpClockRate = 90000 // h264 always uses 90khz
-)
-
 func randUint32() uint32 {
 	var b [4]byte
 	rand.Read(b[:])

pkg/rtph264/rtph264.go

@@ -1,2 +1,10 @@
 // Package rtph264 contains a RTP/H264 decoder and encoder.
 package rtph264
+
+const (
+	rtpVersion   = 0x02
+	rtpClockRate = 90000 // h264 always uses 90khz
+
+	// with a 250 Mbps H264 video, the maximum NALU size is 2.2MB
+	maxNALUSize = 3 * 1024 * 1024
+)