From 7f2864f3375781ad1c6c5fe0e560259187dd7cd8 Mon Sep 17 00:00:00 2001
From: Assaf Vayner <assaf@huggingface.co>
Date: Thu, 30 Jan 2025 11:11:54 -0800
Subject: [PATCH] check max chunk nums (#158)

FIX XET-201

- when validating a chunk header in a xorb ensure that the lengths that
chunk specifies are less than the max size (*2 for some buffer in case
the compression ends up increasing the stored size accidentally.)
---
 cas_object/src/cas_chunk_format.rs | 13 +++++++++++++
 merkledb/src/constants.rs          |  3 +++
 2 files changed, 16 insertions(+)

diff --git a/cas_object/src/cas_chunk_format.rs b/cas_object/src/cas_chunk_format.rs
index 35e77524..c6580513 100644
--- a/cas_object/src/cas_chunk_format.rs
+++ b/cas_object/src/cas_chunk_format.rs
@@ -2,6 +2,7 @@ use std::io::{Read, Write};
 use std::mem::size_of;
 
 use anyhow::anyhow;
+use merkledb::constants::MAXIMUM_CHUNK_SIZE;
 
 use crate::error::CasObjectError;
 use crate::CompressionScheme;
@@ -69,6 +70,18 @@ impl CASChunkHeader {
                 CURRENT_VERSION
             )));
         }
+        if self.get_compressed_length() as usize > MAXIMUM_CHUNK_SIZE * 2 {
+            return Err(CasObjectError::FormatError(anyhow!(
+                "chunk header compressed length too large at {}, maximum: {MAXIMUM_CHUNK_SIZE}",
+                self.get_compressed_length()
+            )));
+        }
+        if self.get_compressed_length() as usize > MAXIMUM_CHUNK_SIZE * 2 {
+            return Err(CasObjectError::FormatError(anyhow!(
+                "chunk header uncompressed length too large at {}, maximum: {MAXIMUM_CHUNK_SIZE}",
+                self.get_uncompressed_length()
+            )));
+        }
         Ok(())
     }
 }
diff --git a/merkledb/src/constants.rs b/merkledb/src/constants.rs
index c53ff486..7e8eafef 100644
--- a/merkledb/src/constants.rs
+++ b/merkledb/src/constants.rs
@@ -10,6 +10,9 @@ pub const MINIMUM_CHUNK_DIVISOR: usize = 8;
 /// TARGET_CDC_CHUNK_SIZE * MAXIMUM_CHUNK_MULTIPLIER is the largest chunk size
 pub const MAXIMUM_CHUNK_MULTIPLIER: usize = 2;
 
+/// no chunk may be larger than MAXIMUM_CHUNK_SIZE bytes
+pub const MAXIMUM_CHUNK_SIZE: usize = TARGET_CDC_CHUNK_SIZE * MAXIMUM_CHUNK_MULTIPLIER;
+
 /// Produce a CAS block when accumulated chunks exceeds TARGET_CAS_BLOCK_SIZE,
 /// this ensures that block sizes are always less than IDEAL_CAS_BLOCK_SIZE.
 pub const TARGET_CAS_BLOCK_SIZE: usize = IDEAL_CAS_BLOCK_SIZE - TARGET_CDC_CHUNK_SIZE * MAXIMUM_CHUNK_MULTIPLIER;