From 356a800e572c20760c89af05052a2eabcc1e1b0f Mon Sep 17 00:00:00 2001
From: lenamonj <jlenamon@gmail.com>
Date: Mon, 7 Sep 2026 21:06:30 -0400
Subject: [PATCH] Check a container's payload fits the file before allocating
 for it

The uncompressed KTX and DDS loaders, the compressed KTX loader and the
.astc loader sized their allocation from header fields and only found out
afterwards that the file did not hold that much data. The allocation is
zero-filled, so a 68-byte KTX declaring 65535x65535 committed 4 GB of
resident memory before failing, and the compressed KTX loader returned
success with a truncated buffer. Every such site now compares the declared
payload against the bytes left in the stream first and reports a corrupt
header instead.
---
 Source/astcenccli_image_load_store.cpp | 57 ++++++++++++++++++++++++++
 Test/astc_test_functional.py           | 10 ++---
 2 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/Source/astcenccli_image_load_store.cpp b/Source/astcenccli_image_load_store.cpp
index 8708d30a..10a53f58 100644
--- a/Source/astcenccli_image_load_store.cpp
+++ b/Source/astcenccli_image_load_store.cpp
@@ -910,6 +910,39 @@ static void ktx_header_switch_endianness(ktx_header * kt)
 	#undef REV
 }
 
+/**
+ * @brief Determine whether a stream still holds a given number of bytes.
+ *
+ * Loaders size allocations from header fields, so a header declaring more
+ * payload than the file holds must be rejected before the allocation.
+ *
+ * @param[in,out] file    The stream, positioned at the start of the payload.
+ * @param         bytes   The payload size the header declared.
+ *
+ * @return @c true if the stream holds at least @c bytes more data.
+ */
+static bool stream_holds_bytes(
+	std::ifstream& file,
+	size_t bytes
+) {
+	std::streampos start = file.tellg();
+	if (start < 0)
+	{
+		return false;
+	}
+
+	file.seekg(0, std::ios::end);
+	std::streampos end = file.tellg();
+	file.seekg(start);
+
+	if (file.fail() || (end < start))
+	{
+		return false;
+	}
+
+	return static_cast<uint64_t>(end - start) >= static_cast<uint64_t>(bytes);
+}
+
 /**
  * @brief Load an uncompressed KTX image using the local custom loader.
  *
@@ -1201,6 +1234,12 @@ static astcenc_image_ptr load_ktx_uncompressed_image(
 		return nullptr;
 	}
 
+	if (!stream_holds_bytes(file, bytes_per_image))
+	{
+		print_error("ERROR: Image header corrupt '%s'\n", filename);
+		return nullptr;
+	}
+
 	std::unique_ptr<uint8_t[]> buf;
 	try
 	{
@@ -1383,6 +1422,12 @@ bool load_ktx_compressed_image(
 		return true;
 	}
 
+	if (!stream_holds_bytes(file, data_len))
+	{
+		print_error("ERROR: Image header corrupt '%s'\n", filename);
+		return true;
+	}
+
 	// Read the data
 	img.data.resize(data_len);
 	file.read(reinterpret_cast<char*>(img.data.data()), data_len);
@@ -2078,6 +2123,12 @@ static astcenc_image_ptr load_dds_uncompressed_image(
 		return nullptr;
 	}
 
+	if (!stream_holds_bytes(file, bytes_per_image))
+	{
+		print_error("ERROR: Image header corrupt '%s'\n", filename);
+		return nullptr;
+	}
+
 	std::unique_ptr<uint8_t[]> buf;
 	try
 	{
@@ -2677,6 +2728,12 @@ int load_cimage(
 		return 1;
 	}
 
+	if (!stream_holds_bytes(file, data_size))
+	{
+		print_error("ERROR: Image header corrupt '%s'\n", filename);
+		return 1;
+	}
+
 	// Allocation may fail if image is suspiciously large
 	try
 	{
diff --git a/Test/astc_test_functional.py b/Test/astc_test_functional.py
index 063e0160..dbc7da59 100644
--- a/Test/astc_test_functional.py
+++ b/Test/astc_test_functional.py
@@ -2207,13 +2207,11 @@ def test_dl_corrupt_astc_magic(self) -> None:
 
     def test_dl_corrupt_astc_huge_size(self) -> None:
         '''
-        Test -dl with an astc file with a data size that will fail alloc.
+        Test -dl with an astc file declaring more data than the file holds.
 
-        Note this test will cause ASAN to error because ASAN itself cannot cope
-        with OOM issues. ASAN documentation says you can run with environment
-        ASAN_OPTIONS=allocator_may_return_null=1, but this doesn't seem to work
-        for allocations made with C++ new[]. This test will need to be disabled
-        when ASAN is used.
+        The loader compares the declared payload size against the bytes left in
+        the file before it allocates, so this is rejected as a corrupt header
+        and never reaches an allocation large enough to trouble ASAN.
         '''
         # Build a valid cmd with a bad image file
         cmd = [