From 19a8f83554dd0d0f184f6336814f533fc0a85fa2 Mon Sep 17 00:00:00 2001
From: Ladislav Zezula <zezula@volny.cz>
Date: Sun, 21 Apr 2024 14:33:20 +0200
Subject: Fixed stack overflow in
 https://github.com/ladislav-zezula/StormLib/issues/328

---
 src/SBaseCommon.cpp    | 16 ++++++----------
 src/SBaseFileTable.cpp | 34 ++++++++++++++++++----------------
 2 files changed, 24 insertions(+), 26 deletions(-)

diff --git a/src/SBaseCommon.cpp b/src/SBaseCommon.cpp
index 587efe4..0de7864 100644
--- a/src/SBaseCommon.cpp
+++ b/src/SBaseCommon.cpp
@@ -1027,20 +1027,16 @@ void * LoadMpqTable(
         if(ByteOffset == SFILE_INVALID_POS)
             FileStream_GetPos(ha->pStream, &ByteOffset);
 
-        // On archives v 1.0, hash table and block table can go beyond EOF.
+        // The hash table and block table can go beyond EOF.
         // Storm.dll reads as much as possible, then fills the missing part with zeros.
         // Abused by Spazzler map protector which sets hash table size to 0x00100000
         // Abused by NP_Protect in MPQs v4 as well
-        if(ha->pHeader->wFormatVersion == MPQ_FORMAT_VERSION_1)
+        FileStream_GetSize(ha->pStream, &FileSize);
+        if((ByteOffset + dwBytesToRead) > FileSize)
         {
-            // Cut the table size
-            FileStream_GetSize(ha->pStream, &FileSize);
-            if((ByteOffset + dwBytesToRead) > FileSize)
-            {
-                // Fill the extra data with zeros
-                dwBytesToRead = (DWORD)(FileSize - ByteOffset);
-                memset(pbMpqTable + dwBytesToRead, 0, (dwTableSize - dwBytesToRead));
-            }
+            // Fill the extra data with zeros
+            dwBytesToRead = (DWORD)(FileSize - ByteOffset);
+            memset(pbMpqTable + dwBytesToRead, 0, (dwTableSize - dwBytesToRead));
         }
 
         // Give the caller information that the table was cut
diff --git a/src/SBaseFileTable.cpp b/src/SBaseFileTable.cpp
index 074501e..274129d 100644
--- a/src/SBaseFileTable.cpp
+++ b/src/SBaseFileTable.cpp
@@ -63,8 +63,8 @@ struct TMPQBits
 {
     static TMPQBits * Create(DWORD NumberOfBits, BYTE FillValue);
 
-    DWORD GetBits(unsigned int nBitPosition, unsigned int nBitLength, void * pvBuffer, int nResultSize);
-    DWORD SetBits(unsigned int nBitPosition, unsigned int nBitLength, void * pvBuffer, int nResultSize);
+    DWORD GetBits(unsigned int nBitPosition, unsigned int nBitLength, void * pvBuffer, unsigned int nResultSize);
+    DWORD SetBits(unsigned int nBitPosition, unsigned int nBitLength, void * pvBuffer, unsigned int nResultSize);
 
     static const USHORT SetBitsMask[];          // Bit mask for each number of bits (0-8)
 
@@ -98,7 +98,7 @@ DWORD TMPQBits::GetBits(
     unsigned int nBitPosition,
     unsigned int nBitLength,
     void * pvBuffer,
-    int nResultByteSize)
+    unsigned int nResultByteSize)
 {
     unsigned char * pbBuffer = (unsigned char *)pvBuffer;
     unsigned int nBytePosition0 = (nBitPosition / 8);
@@ -107,18 +107,17 @@ DWORD TMPQBits::GetBits(
     unsigned int nBitOffset = (nBitPosition & 0x07);
     unsigned char BitBuffer;
 
-    // Keep compilers happy for platforms where nResultByteSize is not used
-    STORMLIB_UNUSED(nResultByteSize);
-
     // Check for bit overflow
     if(nBitPosition + nBitLength < nBitPosition)
         return ERROR_BUFFER_OVERFLOW;
     if(nBitPosition + nBitLength > NumberOfBits)
         return ERROR_BUFFER_OVERFLOW;
+    if(nByteLength > nResultByteSize)
+        return ERROR_BUFFER_OVERFLOW;
 
 #ifdef _DEBUG
     // Check if the target is properly zeroed
-    for(int i = 0; i < nResultByteSize; i++)
+    for(unsigned int i = 0; i < nResultByteSize; i++)
         assert(pbBuffer[i] == 0);
 #endif
 
@@ -170,7 +169,7 @@ DWORD TMPQBits::SetBits(
     unsigned int nBitPosition,
     unsigned int nBitLength,
     void * pvBuffer,
-    int nResultByteSize)
+    unsigned int nResultByteSize)
 {
     unsigned char * pbBuffer = (unsigned char *)pvBuffer;
     unsigned int nBytePosition = (nBitPosition / 8);
@@ -187,6 +186,8 @@ DWORD TMPQBits::SetBits(
         return ERROR_BUFFER_OVERFLOW;
     if(nBitPosition + nBitLength > NumberOfBits)
         return ERROR_BUFFER_OVERFLOW;
+    if(nBitLength / 8 > nResultByteSize)
+        return ERROR_BUFFER_OVERFLOW;
 
 #ifndef STORMLIB_LITTLE_ENDIAN
     // Adjust the buffer pointer for big endian platforms
@@ -1618,15 +1619,16 @@ static DWORD GetFileIndex_Het(TMPQArchive * ha, const char * szFileName)
             DWORD dwFileIndex = 0;
 
             // Get the file index
-            pHetTable->pBetIndexes->GetBits(pHetTable->dwIndexSizeTotal * Index,
-                                            pHetTable->dwIndexSize,
-                                           &dwFileIndex,
-                                            sizeof(DWORD));
-
-            // Verify the FileNameHash against the entry in the table of name hashes
-            if(dwFileIndex <= ha->dwFileTableSize && ha->pFileTable[dwFileIndex].FileNameHash == FileNameHash)
+            if(pHetTable->pBetIndexes->GetBits(pHetTable->dwIndexSizeTotal * Index,
+                                               pHetTable->dwIndexSize,
+                                              &dwFileIndex,
+                                               sizeof(DWORD)) == ERROR_SUCCESS)
             {
-                return dwFileIndex;
+                // Verify the FileNameHash against the entry in the table of name hashes
+                if(dwFileIndex <= ha->dwFileTableSize && ha->pFileTable[dwFileIndex].FileNameHash == FileNameHash)
+                {
+                    return dwFileIndex;
+                }
             }
         }
 
-- 
cgit v1.2.3