+ Support for MPQs locked by the Spazzler protector

1 files changed, 69 insertions, 29 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