From 160a726c1924c6f0166e4aaf258d3da9e14950b3 Mon Sep 17 00:00:00 2001
From: Luca Versari <veluca@google.com>
Date: Thu, 21 Nov 2024 15:28:22 +0100
Subject: [PATCH] Check modular tree height limit.

---
 jxl/src/error.rs         | 2 ++
 jxl/src/frame/modular.rs | 9 +++++++++
 2 files changed, 11 insertions(+)

diff --git a/jxl/src/error.rs b/jxl/src/error.rs
index 9caa410..d832a73 100644
--- a/jxl/src/error.rs
+++ b/jxl/src/error.rs
@@ -113,6 +113,8 @@ pub enum Error {
     InvalidProperty(u32),
     #[error("Modular tree too large: {0}, limit is {1}")]
     TreeTooLarge(usize, usize),
+    #[error("Modular tree too tall: {0}, limit is {1}")]
+    TreeTooTall(usize, usize),
     #[error("Modular tree multiplier too large: {0}, limit is {1}")]
     TreeMultiplierTooLarge(u32, u32),
     #[error("Modular tree multiplier too large: {0}, multiplier log is {1}")]
diff --git a/jxl/src/frame/modular.rs b/jxl/src/frame/modular.rs
index d493be5..2ae6fa0 100644
--- a/jxl/src/frame/modular.rs
+++ b/jxl/src/frame/modular.rs
@@ -153,7 +153,14 @@ impl Tree {
         let mut property_ranges = vec![];
         property_ranges.try_reserve(num_properties * tree.len())?;
         property_ranges.resize(num_properties * tree.len(), (i32::MIN, i32::MAX));
+        let mut height = vec![];
+        height.try_reserve(tree.len())?;
+        height.resize(tree.len(), 0);
         for i in 0..tree.len() {
+            const HEIGHT_LIMIT: usize = 2048;
+            if height[i] > HEIGHT_LIMIT {
+                return Err(Error::TreeTooLarge(height[i], HEIGHT_LIMIT));
+            }
             if let TreeNode::Split {
                 property,
                 val,
@@ -161,6 +168,8 @@ impl Tree {
                 right,
             } = tree[i]
             {
+                height[left as usize] = height[i] + 1;
+                height[right as usize] = height[i] + 1;
                 for p in 0..num_properties {
                     if p == property as usize {
                         let (l, u) = property_ranges[i * num_properties + p];