From d479e9ee7dc8ed72a74318903dbb59f485491363 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Sun, 10 Apr 2022 15:30:05 +0200
Subject: [PATCH] add an upper limit on AU/NALU sizes in order to prevent RAM
 exhaustion

---
 pkg/rtpaac/decoder.go  |  8 +++++++-
 pkg/rtpaac/encoder.go  |  4 ----
 pkg/rtpaac/rtpaac.go   |  7 +++++++
 pkg/rtph264/decoder.go | 12 +++++++++---
 pkg/rtph264/encoder.go |  5 -----
 pkg/rtph264/rtph264.go |  8 ++++++++
 6 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/pkg/rtpaac/decoder.go b/pkg/rtpaac/decoder.go
index dc4bc0e5..b539c261 100644
--- a/pkg/rtpaac/decoder.go
+++ b/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
diff --git a/pkg/rtpaac/encoder.go b/pkg/rtpaac/encoder.go
index d0a81626..c0350c4f 100644
--- a/pkg/rtpaac/encoder.go
+++ b/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[:])
diff --git a/pkg/rtpaac/rtpaac.go b/pkg/rtpaac/rtpaac.go
index 7311c652..01d95928 100644
--- a/pkg/rtpaac/rtpaac.go
+++ 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
+)
diff --git a/pkg/rtph264/decoder.go b/pkg/rtph264/decoder.go
index d3c9b985..ea884e63 100644
--- a/pkg/rtph264/decoder.go
+++ b/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 {
diff --git a/pkg/rtph264/encoder.go b/pkg/rtph264/encoder.go
index d286c322..5aa42939 100644
--- a/pkg/rtph264/encoder.go
+++ b/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[:])
diff --git a/pkg/rtph264/rtph264.go b/pkg/rtph264/rtph264.go
index ab0c3f6f..d0eab52c 100644
--- a/pkg/rtph264/rtph264.go
+++ 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
+)