Merge pull request #326 from ladislav-zezula/LZ_NewProtectedMPQ

Support for new MPQ protection

3 files changed, 45 insertions, 11 deletions

diff --git a/src/SBaseFileTable.cpp b/src/SBaseFileTable.cpp
index fa388fc..a222ac5 100644
--- a/src/SBaseFileTable.cpp
+++ b/src/SBaseFileTable.cpp
@@ -479,8 +479,16 @@ DWORD ConvertMpqHeaderToFormat4(
     {

         case MPQ_FORMAT_VERSION_1:

 

-            // Check for malformed MPQ header version 1.0

+            // Make sure that the MPQ Header is properly swapped

             BSWAP_TMPQHEADER(pHeader, MPQ_FORMAT_VERSION_1);

+

+            // Check for blatantly wrong MPQ header by the hash table position

+            if(((ByteOffset + pHeader->dwHashTablePos) & 0xFFFFFFFF) > FileSize)

+                return ERROR_FAKE_MPQ_HEADER;

+            if(((ByteOffset + pHeader->dwBlockTablePos) & 0xFFFFFFFF) > FileSize)

+                return ERROR_FAKE_MPQ_HEADER;

+

+            // Check for malformed MPQ header version 1.0

             if(pHeader->wFormatVersion != MPQ_FORMAT_VERSION_1 || pHeader->dwHeaderSize != MPQ_HEADER_SIZE_V1)

             {

                 pHeader->wFormatVersion = MPQ_FORMAT_VERSION_1;

diff --git a/src/SFileOpenArchive.cpp b/src/SFileOpenArchive.cpp
index 4b2d3b4..a97ecea 100644
--- a/src/SFileOpenArchive.cpp
+++ b/src/SFileOpenArchive.cpp
@@ -228,7 +228,7 @@ bool WINAPI SFileOpenArchive(
     DWORD dwFlags,
     HANDLE * phMpq)
 {
-    TMPQUserData * pUserData;
+    TMPQUserData * pUserData = NULL;
     TFileStream * pStream = NULL;       // Open file stream
     TMPQArchive * ha = NULL;            // Archive handle
     TFileEntry * pFileEntry;
@@ -354,18 +354,25 @@ bool WINAPI SFileOpenArchive(
                 {
                     if(ha->pUserData == NULL && dwHeaderID == ID_MPQ_USERDATA)
                     {
+                        // Copy the eventual user data to the separate buffer
+                        memcpy(&ha->UserData, ha->HeaderData, sizeof(TMPQUserData));
+
                         // Verify if this looks like a valid user data
-                        pUserData = IsValidMpqUserData(ByteOffset, FileSize, ha->HeaderData);
+                        pUserData = IsValidMpqUserData(ByteOffset, FileSize, &ha->UserData);
                         if(pUserData != NULL)
                         {
-                            // Fill the user data header
-                            ha->UserDataPos = ByteOffset;
-                            ha->pUserData = &ha->UserData;
-                            memcpy(ha->pUserData, pUserData, sizeof(TMPQUserData));
-
-                            // Continue searching from that position
-                            ByteOffset += ha->pUserData->dwHeaderOffs;
-                            break;
+                            // Set the byte offset to the loaded user data
+                            ULONGLONG TempByteOffset = ByteOffset + pUserData->dwHeaderOffs;
+
+                            // Read the eventual MPQ header from the position where the user data points
+                            if(!FileStream_Read(ha->pStream, &TempByteOffset, ha->HeaderData, sizeof(ha->HeaderData)))
+                            {
+                                dwErrCode = GetLastError();
+                                break;
+                            }
+
+                            // Re-initialize the header ID
+                            dwHeaderID = BSWAP_INT32_UNSIGNED(ha->HeaderData[0]);
                         }
                     }
                 }
@@ -405,12 +412,24 @@ bool WINAPI SFileOpenArchive(
 
                 // Move the pointers
                 ByteOffset += 0x200;
+                pUserData = NULL;
             }
         }
 
         // Did we identify one of the supported headers?
         if(dwErrCode == ERROR_SUCCESS)
         {
+            // If we retrieved the offset from the user data offset, initialize the user data
+            if(pUserData != NULL)
+            {
+                // Fill the user data header
+                ha->pUserData = &ha->UserData;
+                ha->UserDataPos = ByteOffset;
+
+                // Set the real byte offset
+                ByteOffset = ByteOffset + pUserData->dwHeaderOffs;
+            }
+
             // Set the user data position to the MPQ header, if none
             if(ha->pUserData == NULL)
                 ha->UserDataPos = ByteOffset;
diff --git a/src/SFileReadFile.cpp b/src/SFileReadFile.cpp
index b0698c1..8ab5f48 100644
--- a/src/SFileReadFile.cpp
+++ b/src/SFileReadFile.cpp
@@ -194,6 +194,13 @@ static DWORD ReadMpqSectors(TMPQFile * hf, LPBYTE pbBuffer, DWORD dwByteOffset,
                         dwErrCode = ERROR_FILE_CORRUPT;
                         break;
                     }
+
+                    // Special case (MPQ_2024_v1_300TK2.09p.w3x, file File00010254.blp):
+                    // Extracted less than required. Fill the rest with zeros
+                    if((DWORD)(cbOutSector) < dwBytesInThisSector)
+                    {
+                        memset(pbOutSector + cbOutSector, 0, dwBytesInThisSector - cbOutSector);
+                    }
                 }
                 else
                 {