From 5e179844a864e130918f045165a75fa268629b4a Mon Sep 17 00:00:00 2001
From: Ladislav Zezula <ladislav.zezula@avast.com>
Date: Sun, 5 Jun 2022 16:34:05 +0200
Subject: Added more strict check for corrupt HET table

---
 src/SBaseFileTable.cpp | 47 ++++++++++++++++++++++++-----------------------
 test/StormTest.cpp     |  3 +--
 2 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/src/SBaseFileTable.cpp b/src/SBaseFileTable.cpp
index 94e4469..7d7a401 100644
--- a/src/SBaseFileTable.cpp
+++ b/src/SBaseFileTable.cpp
@@ -1480,31 +1480,32 @@ static TMPQHetTable * TranslateHetTable(TMPQHetHeader * pHetHeader)
         if(pHetHeader->ExtHdr.dwDataSize >= pHetHeader->dwTableSize)
         {
             // The size of the HET table must be sum of header, hash and index table size
-            assert((sizeof(TMPQHetHeader) - sizeof(TMPQExtHeader) + pHetHeader->dwTotalCount + pHetHeader->dwIndexTableSize) == pHetHeader->dwTableSize);
-
-            // So far, all MPQs with HET Table have had total number of entries equal to 4/3 of file count
-            // Exception: "2010 - Starcraft II\!maps\Tya's Zerg Defense (unprotected).SC2Map"
-//          assert(((pHetHeader->dwEntryCount * 4) / 3) == pHetHeader->dwTotalCount);
+            if((sizeof(TMPQHetHeader) - sizeof(TMPQExtHeader) + pHetHeader->dwTotalCount + pHetHeader->dwIndexTableSize) == pHetHeader->dwTableSize)
+            {
+                // So far, all MPQs with HET Table have had total number of entries equal to 4/3 of file count
+                // Exception: "2010 - Starcraft II\!maps\Tya's Zerg Defense (unprotected).SC2Map"
+//              assert(((pHetHeader->dwEntryCount * 4) / 3) == pHetHeader->dwTotalCount);
 
-            // The size of one index is predictable as well
-            assert(GetNecessaryBitCount(pHetHeader->dwEntryCount) == pHetHeader->dwIndexSizeTotal);
+                // The size of one index is predictable as well
+                assert(GetNecessaryBitCount(pHetHeader->dwEntryCount) == pHetHeader->dwIndexSizeTotal);
 
-            // The size of index table (in entries) is expected
-            // to be the same like the hash table size (in bytes)
-            assert(((pHetHeader->dwTotalCount * pHetHeader->dwIndexSizeTotal) + 7) / 8 == pHetHeader->dwIndexTableSize);
+                // The size of index table (in entries) is expected
+                // to be the same like the hash table size (in bytes)
+                assert(((pHetHeader->dwTotalCount * pHetHeader->dwIndexSizeTotal) + 7) / 8 == pHetHeader->dwIndexTableSize);
 
-            // Create translated table
-            pHetTable = CreateHetTable(pHetHeader->dwEntryCount, pHetHeader->dwTotalCount, pHetHeader->dwNameHashBitSize, pbSrcData);
-            if(pHetTable != NULL)
-            {
-                // Now the sizes in the hash table should be already set
-                assert(pHetTable->dwEntryCount     == pHetHeader->dwEntryCount);
-                assert(pHetTable->dwTotalCount     == pHetHeader->dwTotalCount);
-                assert(pHetTable->dwIndexSizeTotal == pHetHeader->dwIndexSizeTotal);
-
-                // Copy the missing variables
-                pHetTable->dwIndexSizeExtra = pHetHeader->dwIndexSizeExtra;
-                pHetTable->dwIndexSize      = pHetHeader->dwIndexSize;
+                // Create translated table
+                pHetTable = CreateHetTable(pHetHeader->dwEntryCount, pHetHeader->dwTotalCount, pHetHeader->dwNameHashBitSize, pbSrcData);
+                if(pHetTable != NULL)
+                {
+                    // Now the sizes in the hash table should be already set
+                    assert(pHetTable->dwEntryCount     == pHetHeader->dwEntryCount);
+                    assert(pHetTable->dwTotalCount     == pHetHeader->dwTotalCount);
+                    assert(pHetTable->dwIndexSizeTotal == pHetHeader->dwIndexSizeTotal);
+
+                    // Copy the missing variables
+                    pHetTable->dwIndexSizeExtra = pHetHeader->dwIndexSizeExtra;
+                    pHetTable->dwIndexSize      = pHetHeader->dwIndexSize;
+                }
             }
         }
     }
@@ -2449,7 +2450,7 @@ TMPQHetTable * LoadHetTable(TMPQArchive * ha)
         pExtTable = LoadExtTable(ha, pHeader->HetTablePos64, (size_t)pHeader->HetTableSize64, HET_TABLE_SIGNATURE, MPQ_KEY_HASH_TABLE);
         if(pExtTable != NULL)
         {
-            // If loading HET table fails, we ignore the result.
+            // Translate the loaded table into HET table.
             pHetTable = TranslateHetTable((TMPQHetHeader *)pExtTable);
             STORM_FREE(pExtTable);
         }
diff --git a/test/StormTest.cpp b/test/StormTest.cpp
index f8b8f7d..78eb94b 100644
--- a/test/StormTest.cpp
+++ b/test/StormTest.cpp
@@ -4345,8 +4345,7 @@ int _tmain(int argc, TCHAR * argv[])
 
     for(int i = 2; i < argc; i++)
     {
-//      TestArchive(argv[i], Bliz, 0, "sound\\zerg\\advisor\\zaderr05.wav", NULL);
-//      TestArchive(argv[i], Bliz, 0, "music\\trdyroom.wav", NULL);
+        TestArchive(argv[i], Bliz, 0, "sound\\zerg\\advisor\\zaderr05.wav", NULL);
     }
 
     //
-- 
cgit v1.2.3