From fe96889ca25284d661f5b725bcb8f10fb50c1da3 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Thu, 1 Aug 2024 15:36:40 +0200
Subject: [PATCH] fmp4: prevent RAM exhaustion by limiting max sample count

---
 pkg/formats/fmp4/fmp4.go                                 | 4 ++++
 pkg/formats/fmp4/parts.go                                | 9 +++++++++
 .../testdata/fuzz/FuzzPartsUnmarshal/971288027908400b    | 2 ++
 .../testdata/fuzz/FuzzPartsUnmarshal/aa157382146cae53    | 2 ++
 4 files changed, 17 insertions(+)
 create mode 100644 pkg/formats/fmp4/testdata/fuzz/FuzzPartsUnmarshal/971288027908400b
 create mode 100644 pkg/formats/fmp4/testdata/fuzz/FuzzPartsUnmarshal/aa157382146cae53

diff --git a/pkg/formats/fmp4/fmp4.go b/pkg/formats/fmp4/fmp4.go
index 1fe83d0..a18cb07 100644
--- a/pkg/formats/fmp4/fmp4.go
+++ b/pkg/formats/fmp4/fmp4.go
@@ -1,2 +1,6 @@
 // Package fmp4 contains a fragmented-MP4 reader and writer.
 package fmp4
+
+const (
+	maxSamplesPerTrun = 120 * 160 // 120fps * 60 seconds
+)
diff --git a/pkg/formats/fmp4/parts.go b/pkg/formats/fmp4/parts.go
index 4313b07..a28cf3f 100644
--- a/pkg/formats/fmp4/parts.go
+++ b/pkg/formats/fmp4/parts.go
@@ -110,6 +110,15 @@ func (ps *Parts) Unmarshal(byts []byte) error {
 					return nil, fmt.Errorf("unexpected trun")
 				}
 
+				// prevent RAM exhaustion due to unlimited Trun unmarshaling
+				rawBox := byts[h.BoxInfo.Offset:]
+				if len(rawBox) >= 16 {
+					sampleCount := uint32(rawBox[12])<<24 | uint32(rawBox[13])<<16 | uint32(rawBox[14])<<8 | uint32(rawBox[15])
+					if sampleCount > maxSamplesPerTrun {
+						return nil, fmt.Errorf("sample count (%d) exceeds maximum (%d)", sampleCount, maxSamplesPerTrun)
+					}
+				}
+
 				box, _, err := h.ReadPayload()
 				if err != nil {
 					return nil, err
diff --git a/pkg/formats/fmp4/testdata/fuzz/FuzzPartsUnmarshal/971288027908400b b/pkg/formats/fmp4/testdata/fuzz/FuzzPartsUnmarshal/971288027908400b
new file mode 100644
index 0000000..0c8b704
--- /dev/null
+++ b/pkg/formats/fmp4/testdata/fuzz/FuzzPartsUnmarshal/971288027908400b
@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("\x00\x00\x00\xc8moof\x00\x00\x00\x10mfhd\x00\x00\x00\x00\x00\x00\x00\x04\x00\x00\x00`traf\x00\x00\x00\x10tfhd\x00\x00\x00\x00\x00\x01_\x90\x00\x00\x004trun\x00\x01\x00\x00\xf9\xff\xff\x00")
diff --git a/pkg/formats/fmp4/testdata/fuzz/FuzzPartsUnmarshal/aa157382146cae53 b/pkg/formats/fmp4/testdata/fuzz/FuzzPartsUnmarshal/aa157382146cae53
new file mode 100644
index 0000000..67b25ad
--- /dev/null
+++ b/pkg/formats/fmp4/testdata/fuzz/FuzzPartsUnmarshal/aa157382146cae53
@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0000moof\x00\x00\x00\x10mfhd\x00000000000 0traf\x00\x00\x00\x10tfhd\x0000\x00000000\x000trun")