From c78042065812cb58ca51f331159d46415e66db8d Mon Sep 17 00:00:00 2001
From: mayp <mayanping@ncti-gba.cn>
Date: Thu, 18 Aug 2022 16:49:57 +0800
Subject: [PATCH] Fix CVE-2021-20304

---
 IlmImf/ImfHuf.cpp      |   9 +++
 IlmImfTest/testHuf.cpp | 141 +++++++++++++++++++++++------------------
 2 files changed, 89 insertions(+), 61 deletions(-)

diff --git a/IlmImf/ImfHuf.cpp b/IlmImf/ImfHuf.cpp
index aa708a8..82af799 100644
--- a/IlmImf/ImfHuf.cpp
+++ b/IlmImf/ImfHuf.cpp
@@ -897,6 +897,11 @@ hufDecode
 		//
 
 		lc -= pl.len;
+
+		if ( lc < 0 )
+		{
+			invalidCode(); // code length too long
+		}
 		getCode (pl.lit, rlc, c, lc, in, out, outb, oe);
 	    }
 	    else
@@ -954,6 +959,10 @@ hufDecode
 	if (pl.len)
 	{
 	    lc -= pl.len;
+            if ( lc < 0 )
+            {
+   	        invalidCode(); // code length too long
+            }
 	    getCode (pl.lit, rlc, c, lc, in, out, outb, oe);
 	}
 	else
diff --git a/IlmImfTest/testHuf.cpp b/IlmImfTest/testHuf.cpp
index d2728fb..10d3906 100644
--- a/IlmImfTest/testHuf.cpp
+++ b/IlmImfTest/testHuf.cpp
@@ -180,67 +180,86 @@ testHuf (const std::string&)
 
 	IMATH_NAMESPACE::Rand48 rand48 (0);
 
-	const int N = 1000000;
-	Array <unsigned short> raw (N);
-
-	fill1 (raw, N, 1, rand48);	  // test various symbol distributions
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
-	fill1 (raw, N, 10, rand48);
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
-	fill1 (raw, N, 100, rand48);
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
-	fill1 (raw, N, 1000, rand48);
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
-
-	fill2 (raw, N, 1, rand48);
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
-	fill2 (raw, N, 10, rand48);
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
-	fill2 (raw, N, 100, rand48);
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
-	fill2 (raw, N, 1000, rand48);
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
-
-	fill3 (raw, N, 0);
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
-	fill3 (raw, N, 1);
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
-	fill3 (raw, N, USHRT_MAX - 1);
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
-	fill3 (raw, N, USHRT_MAX);
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
-
-	fill4 (raw, USHRT_MAX + 1);
-	compressUncompress (raw, USHRT_MAX + 1);
-	compressUncompressSubset (raw, USHRT_MAX + 1);
-	fill4 (raw, N);
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
-
-	fill4 (raw, 0);
-	compressUncompress (raw, 0);	// test small input data sets
-	fill4 (raw, 1);
-	compressUncompress (raw, 1);
-	fill4 (raw, 2);
-	compressUncompress (raw, 2);
-	fill4 (raw, 3);
-	compressUncompress (raw, 3);
-
-	fill5 (raw, N);			// test run-length coding of code table
-	compressUncompress (raw, N);
-	compressUncompressSubset (raw, N);
+        //
+        // FastHufDecoder is used for more than 128 bits, so first test with fewer than 128 bits,
+        // then test FastHufDecoder
+        //
+        for (int pass = 0 ; pass < 2 ; ++pass)
+        {
+
+            int N = pass==0 ? 12 : 1000000;
+            Array <unsigned short> raw (N);
+
+            fill1 (raw, N, 1, rand48);	  // test various symbol distributions
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+            fill1 (raw, N, 10, rand48);
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+            fill1 (raw, N, 100, rand48);
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+            fill1 (raw, N, 1000, rand48);
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+
+            fill2 (raw, N, 1, rand48);
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+            fill2 (raw, N, 10, rand48);
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+            fill2 (raw, N, 100, rand48);
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+            fill2 (raw, N, 1000, rand48);
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+
+            fill3 (raw, N, 0);
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+            fill3 (raw, N, 1);
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+            fill3 (raw, N, USHRT_MAX - 1);
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+            fill3 (raw, N, USHRT_MAX);
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+
+            if (pass==1)
+            {
+                fill4 (raw, USHRT_MAX + 1);
+                compressVerify(raw, USHRT_MAX + 1, HUF_COMPRESS_DEK_HASH_FOR_FILL4_USHRT_MAX_PLUS_ONE);
+
+                compressUncompress (raw, USHRT_MAX + 1);
+                compressUncompressSubset (raw, USHRT_MAX + 1);
+                fill4 (raw, N);
+                compressVerify(raw, N, HUF_COMPRESS_DEK_HASH_FOR_FILL4_N);
+            }
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+
+            fill4 (raw, 0);
+            compressUncompress (raw, 0);	// test small input data sets
+            fill4 (raw, 1);
+            compressUncompress (raw, 1);
+            fill4 (raw, 2);
+            compressUncompress (raw, 2);
+            fill4 (raw, 3);
+            compressUncompress (raw, 3);
+
+            fill5 (raw, N);			// test run-length coding of code table
+            if (pass==1)
+            {
+                compressVerify(raw, N, HUF_COMPRESS_DEK_HASH_FOR_FILL5_N);
+            }
+            compressUncompress (raw, N);
+            compressUncompressSubset (raw, N);
+
+        }
 
 	cout << "ok\n" << endl;
     }
-- 
2.33.0