+ Support for MPQs locked by the Spazzler protector

4 files changed, 87 insertions, 38 deletions

diff --git a/src/SBaseFileTable.cpp b/src/SBaseFileTable.cpp
index 417a44b..564287c 100644
--- a/src/SBaseFileTable.cpp
+++ b/src/SBaseFileTable.cpp
@@ -250,7 +250,7 @@ static DWORD GetMaxFileOffset32(TMPQArchive * ha)
     return dwMaxFileOffset;
 }
 
-static ULONGLONG DetermineArchiveSize_V1_V2(
+static ULONGLONG DetermineArchiveSize_V1(
     TMPQArchive * ha,
     TMPQHeader * pHeader,
     ULONGLONG MpqOffset,
@@ -261,20 +261,20 @@ static ULONGLONG DetermineArchiveSize_V1_V2(
     DWORD SignatureHeader = 0;
     DWORD dwArchiveSize32;
 
+    // This could only be called for MPQs version 1.0
+    assert(pHeader->wFormatVersion == MPQ_FORMAT_VERSION_1);
+
     // Check if we can rely on the archive size in the header
-    if((FileSize >> 0x20) == 0)
+    if(pHeader->dwBlockTablePos < pHeader->dwArchiveSize)
     {
-        if(pHeader->dwBlockTablePos < pHeader->dwArchiveSize)
-        {
-            // If the block table end matches the archive size, we trust the archive size
-            if((pHeader->dwArchiveSize - pHeader->dwBlockTablePos) <= (pHeader->dwBlockTableSize * sizeof(TMPQBlock)))
-                return pHeader->dwArchiveSize;
-
-            // If the archive size in the header is less than real file size
-            dwArchiveSize32 = (DWORD)(FileSize - MpqOffset);
-            if(pHeader->dwArchiveSize <= dwArchiveSize32)
-                return pHeader->dwArchiveSize;
-        }
+        // The block table cannot be compressed, so the sizes must match
+        if((pHeader->dwArchiveSize - pHeader->dwBlockTablePos) == (pHeader->dwBlockTableSize * sizeof(TMPQBlock)))
+            return pHeader->dwArchiveSize;
+
+        // If the archive size in the header is less than real file size
+        dwArchiveSize32 = (DWORD)(FileSize - MpqOffset);
+        if(pHeader->dwArchiveSize == dwArchiveSize32)
+            return pHeader->dwArchiveSize;
     }
 
     // Check if there is a signature header
@@ -290,6 +290,38 @@ static ULONGLONG DetermineArchiveSize_V1_V2(
     return (EndOfMpq - MpqOffset);
 }
 
+static ULONGLONG DetermineArchiveSize_V2(
+    TMPQHeader * pHeader,
+    ULONGLONG MpqOffset,
+    ULONGLONG FileSize)
+{
+    ULONGLONG EndOfMpq = FileSize;
+    DWORD dwArchiveSize32;
+
+    // Check if we can rely on the archive size in the header
+    if((FileSize >> 0x20) == 0)
+    {
+        if(pHeader->dwBlockTablePos < pHeader->dwArchiveSize)
+        {
+            // MPQs version 2.0: The block table can be compressed,
+            // so block table size could be less than dwBlockTableSize * sizeof(TMPQBlock)
+            if(pHeader->wFormatVersion == MPQ_FORMAT_VERSION_2)
+            {
+                if((pHeader->dwArchiveSize - pHeader->dwBlockTablePos) <= (pHeader->dwBlockTableSize * sizeof(TMPQBlock)))
+                    return pHeader->dwArchiveSize;
+
+                // If the archive size in the header is less than real file size
+                dwArchiveSize32 = (DWORD)(FileSize - MpqOffset);
+                if(pHeader->dwArchiveSize <= dwArchiveSize32)
+                    return pHeader->dwArchiveSize;
+            }
+        }
+    }
+
+    // Return the calculated archive size
+    return (EndOfMpq - MpqOffset);
+}
+
 // This function converts the MPQ header so it always looks like version 4
 /*static ULONGLONG GetArchiveSize64(TMPQHeader * pHeader)
 {
@@ -350,6 +382,7 @@ int ConvertMpqHeaderToFormat4(
     ULONGLONG HashTablePos64 = 0;
     ULONGLONG ByteOffset;
     USHORT wFormatVersion = BSWAP_INT16_UNSIGNED(pHeader->wFormatVersion);
+    DWORD dwTableSize;
     int nError = ERROR_SUCCESS;
 
     // If version 1.0 is forced, then the format version is forced to be 1.0
@@ -393,8 +426,19 @@ int ConvertMpqHeaderToFormat4(
             // Determine the archive size on malformed MPQs
             if(ha->dwFlags & MPQ_FLAG_MALFORMED)
             {
-                pHeader->ArchiveSize64 = DetermineArchiveSize_V1_V2(ha, pHeader, MpqOffset, FileSize);
+                // Calculate the archive size
+                pHeader->ArchiveSize64 = DetermineArchiveSize_V1(ha, pHeader, MpqOffset, FileSize);
                 pHeader->dwArchiveSize = (DWORD)pHeader->ArchiveSize64;
+
+                // Handle case when either the MPQ is cut in the middle of the block table
+                // or the MPQ is malformed so that block table size is greater than should be
+                ByteOffset = MpqOffset + pHeader->dwBlockTablePos;
+                dwTableSize = pHeader->dwBlockTableSize * sizeof(TMPQBlock);
+                if((ByteOffset + dwTableSize) > FileSize)
+                {
+                    pHeader->BlockTableSize64 = FileSize - ByteOffset;
+                    pHeader->dwBlockTableSize = (DWORD)(pHeader->BlockTableSize64 / sizeof(TMPQBlock));
+                }
             }
             break;
 
@@ -404,6 +448,7 @@ int ConvertMpqHeaderToFormat4(
             BSWAP_TMPQHEADER(pHeader, MPQ_FORMAT_VERSION_2);
             if(pHeader->dwHeaderSize != MPQ_HEADER_SIZE_V2)
             {
+                pHeader->wFormatVersion = MPQ_FORMAT_VERSION_1;
                 pHeader->dwHeaderSize = MPQ_HEADER_SIZE_V1;
                 ha->dwFlags |= MPQ_FLAG_MALFORMED;
                 goto Label_ArchiveVersion1;
@@ -436,7 +481,7 @@ int ConvertMpqHeaderToFormat4(
                     assert(pHeader->BlockTableSize64 <= (pHeader->dwBlockTableSize * sizeof(TMPQBlock)));
 
                     // Determine real archive size
-                    pHeader->ArchiveSize64 = DetermineArchiveSize_V1_V2(ha, pHeader, MpqOffset, FileSize);
+                    pHeader->ArchiveSize64 = DetermineArchiveSize_V2(pHeader, MpqOffset, FileSize);
 
                     // Calculate the size of the hi-block table
                     pHeader->HiBlockTableSize64 = pHeader->ArchiveSize64 - pHeader->HiBlockTablePos64;
@@ -445,7 +490,7 @@ int ConvertMpqHeaderToFormat4(
                 else
                 {
                     // Determine real archive size
-                    pHeader->ArchiveSize64 = DetermineArchiveSize_V1_V2(ha, pHeader, MpqOffset, FileSize);
+                    pHeader->ArchiveSize64 = DetermineArchiveSize_V2(pHeader, MpqOffset, FileSize);
 
                     // Calculate size of the block table
                     pHeader->BlockTableSize64 = pHeader->ArchiveSize64 - BlockTablePos64;
@@ -2201,7 +2246,6 @@ TMPQBlock * LoadBlockTable(TMPQArchive * ha, bool bDontFixEntries)
     TMPQHeader * pHeader = ha->pHeader;
     TMPQBlock * pBlockTable = NULL;
     ULONGLONG ByteOffset;
-    ULONGLONG FileSize;
     DWORD dwTableSize;
     DWORD dwCmpSize;
 
@@ -2223,16 +2267,8 @@ TMPQBlock * LoadBlockTable(TMPQArchive * ha, bool bDontFixEntries)
             dwTableSize = pHeader->dwBlockTableSize * sizeof(TMPQBlock);
             dwCmpSize = (DWORD)pHeader->BlockTableSize64;
 
-            // Handle case when either the MPQ is cut in the middle of the block table
-            // or the MPQ is malformed so that block table size is greater than should be
-            FileStream_GetSize(ha->pStream, &FileSize);
-            if(dwTableSize == dwCmpSize && (ByteOffset + dwTableSize) > FileSize)
-            {
-                pHeader->BlockTableSize64 = FileSize - ByteOffset;
-                pHeader->dwBlockTableSize = (DWORD)(pHeader->BlockTableSize64 / sizeof(TMPQBlock));
-                dwTableSize = dwCmpSize = (DWORD)pHeader->BlockTableSize64;
-                ha->dwFlags |= MPQ_FLAG_MALFORMED;
-            }
+            // If the the block table size was invalid, it should be fixed by now
+            assert((ByteOffset + dwCmpSize) <= (ha->MpqPos + ha->pHeader->ArchiveSize64));
 
             //
             // One of the first cracked versions of Diablo I had block table unencrypted 
@@ -2540,19 +2576,23 @@ int BuildFileTable_HetBet(
 int BuildFileTable(TMPQArchive * ha)
 {
     TFileEntry * pFileTable;
+    DWORD dwFileTableSize;
     bool bFileTableCreated = false;
 
     // Sanity checks
     assert(ha->dwFileTableSize == 0);
     assert(ha->dwMaxFileCount != 0);
 
+    // Determine the allocation size for the file table
+    dwFileTableSize = STORMLIB_MAX(ha->pHeader->dwBlockTableSize, ha->dwMaxFileCount);
+
     // Allocate the file table with size determined before
-    pFileTable = STORM_ALLOC(TFileEntry, ha->dwMaxFileCount);
+    pFileTable = STORM_ALLOC(TFileEntry, dwFileTableSize);
     if(pFileTable == NULL)
         return ERROR_NOT_ENOUGH_MEMORY;
 
     // Fill the table with zeros
-    memset(pFileTable, 0, ha->dwMaxFileCount * sizeof(TFileEntry));
+    memset(pFileTable, 0, dwFileTableSize * sizeof(TFileEntry));
 
     // If we have HET table, we load file table from the BET table
     // Note: If BET table is corrupt or missing, we set the archive as read only
diff --git a/src/SFileOpenArchive.cpp b/src/SFileOpenArchive.cpp
index 27c9972..356aa59 100644
--- a/src/SFileOpenArchive.cpp
+++ b/src/SFileOpenArchive.cpp
@@ -21,12 +21,11 @@
 /* Local functions                                                           */
 /*****************************************************************************/
 
-static bool IsAviFile(void * pvFileBegin)
+static bool IsAviFile(DWORD * HeaderData)
 {
-    LPDWORD AviHeader = (DWORD *)pvFileBegin;
-    DWORD DwordValue0 = BSWAP_INT32_UNSIGNED(AviHeader[0]);
-    DWORD DwordValue2 = BSWAP_INT32_UNSIGNED(AviHeader[2]);
-    DWORD DwordValue3 = BSWAP_INT32_UNSIGNED(AviHeader[3]);
+    DWORD DwordValue0 = BSWAP_INT32_UNSIGNED(HeaderData[0]);
+    DWORD DwordValue2 = BSWAP_INT32_UNSIGNED(HeaderData[2]);
+    DWORD DwordValue3 = BSWAP_INT32_UNSIGNED(HeaderData[3]);
 
     // Test for 'RIFF', 'AVI ' or 'LIST'
     return (DwordValue0 == 0x46464952 && DwordValue2 == 0x20495641 && DwordValue3 == 0x5453494C);
@@ -192,6 +191,7 @@ bool WINAPI SFileOpenArchive(
     {
         ULONGLONG SearchOffset = 0;
         DWORD dwStreamFlags = 0;
+        DWORD dwHeaderSize;
         DWORD dwHeaderID;
 
         memset(ha, 0, sizeof(TMPQArchive));
@@ -231,7 +231,7 @@ bool WINAPI SFileOpenArchive(
             }
 
             // If there is the MPQ user data signature, process it
-            dwHeaderID = BSWAP_INT32_UNSIGNED(*(LPDWORD)ha->HeaderData);
+            dwHeaderID = BSWAP_INT32_UNSIGNED(ha->HeaderData[0]);
             if(dwHeaderID == ID_MPQ_USERDATA && ha->pUserData == NULL && (dwFlags & MPQ_OPEN_FORCE_MPQ_V1) == 0)
             {
                 // Verify if this looks like a valid user data
@@ -249,8 +249,12 @@ bool WINAPI SFileOpenArchive(
                 }
             }
 
-            // There must be MPQ header signature
-            if(dwHeaderID == ID_MPQ)
+            // There must be MPQ header signature. Note that STORM.dll from Warcraft III actually
+            // tests the MPQ header size. It must be at least 0x20 bytes in order to load it
+            // Abused by Spazzler Map protector. Note that the size check is not present
+            //  in Storm.dll v 1.00, so Diablo I code would load the MPQ anyway.
+            dwHeaderSize = BSWAP_INT32_UNSIGNED(ha->HeaderData[1]);
+            if(dwHeaderID == ID_MPQ && dwHeaderSize >= MPQ_HEADER_SIZE_V1)
             {
                 // Now convert the header to version 4
                 nError = ConvertMpqHeaderToFormat4(ha, SearchOffset, FileSize, dwFlags);
diff --git a/src/StormLib.h b/src/StormLib.h
index a0cfeb6..b1064d5 100644
--- a/src/StormLib.h
+++ b/src/StormLib.h
@@ -500,6 +500,7 @@ void SetBits(TBitArray * array, unsigned int nBitPosition, unsigned int nBitLeng
 #define MPQ_HEADER_SIZE_V2    0x2C
 #define MPQ_HEADER_SIZE_V3    0x44
 #define MPQ_HEADER_SIZE_V4    0xD0
+#define MPQ_HEADER_DWORDS     (MPQ_HEADER_SIZE_V4 / 0x04)
 
 typedef struct _TMPQUserData
 {
@@ -837,7 +838,7 @@ typedef struct _TMPQArchive
     HASH_STRING    pfnHashString;               // Hashing function that will convert the file name into hash
     
     TMPQUserData   UserData;                    // MPQ user data. Valid only when ID_MPQ_USERDATA has been found
-    BYTE           HeaderData[MPQ_HEADER_SIZE_V4];  // Storage for MPQ header
+    DWORD          HeaderData[MPQ_HEADER_DWORDS];  // Storage for MPQ header
 
     DWORD          dwHETBlockSize;
     DWORD          dwBETBlockSize;
diff --git a/test/Test.cpp b/test/Test.cpp
index 7c50b4b..6b76848 100644
--- a/test/Test.cpp
+++ b/test/Test.cpp
@@ -3534,6 +3534,10 @@ int main(int argc, char * argv[])
     if(nError == ERROR_SUCCESS)
         nError = TestOpenArchive("MPQ_2010_v2_HasUserData.s2ma");
 
+    // Open an Warcraft III map locked by the Spazzler protector
+    if(nError == ERROR_SUCCESS)
+        nError = TestOpenArchive("MPQ_2002_v1_ProtectedMap_Spazzler.w3x");
+
     // Open a MPQ archive v 3.0
     if(nError == ERROR_SUCCESS)
         nError = TestOpenArchive("MPQ_2010_v3_expansion-locale-frFR.MPQ");