Fixed heap overflow in handling of file patch

1 files changed, 23 insertions, 17 deletions

diff --git a/src/SBaseCommon.cpp b/src/SBaseCommon.cpp
index b0029e0..b5b880e 100644
--- a/src/SBaseCommon.cpp
+++ b/src/SBaseCommon.cpp
@@ -1154,6 +1154,7 @@ DWORD AllocateSectorBuffer(TMPQFile * hf)
 DWORD AllocatePatchInfo(TMPQFile * hf, bool bLoadFromFile)

 {

     TMPQArchive * ha = hf->ha;

+    TPatchInfo * pPatchInfo;

     DWORD dwLength = sizeof(TPatchInfo);

 

     // The following conditions must be true

@@ -1164,35 +1165,39 @@ __AllocateAndLoadPatchInfo:
 

     // Allocate space for patch header. Start with default size,

     // and if its size if bigger, then we reload them

-    hf->pPatchInfo = STORM_ALLOC(TPatchInfo, 1);

-    if(hf->pPatchInfo == NULL)

+    pPatchInfo = (TPatchInfo *)(STORM_ALLOC(BYTE, dwLength));

+    if(pPatchInfo == NULL)

         return ERROR_NOT_ENOUGH_MEMORY;

 

     // Do we have to load the patch header from the file ?

     if(bLoadFromFile)

     {

         // Load the patch header

-        if(!FileStream_Read(ha->pStream, &hf->RawFilePos, hf->pPatchInfo, dwLength))

+        if(!FileStream_Read(ha->pStream, &hf->RawFilePos, pPatchInfo, dwLength))

         {

-            // Free the patch info

-            STORM_FREE(hf->pPatchInfo);

-            hf->pPatchInfo = NULL;

+            STORM_FREE(pPatchInfo);

             return GetLastError();

         }

 

         // Perform necessary swapping

-        hf->pPatchInfo->dwLength = BSWAP_INT32_UNSIGNED(hf->pPatchInfo->dwLength);

-        hf->pPatchInfo->dwFlags = BSWAP_INT32_UNSIGNED(hf->pPatchInfo->dwFlags);

-        hf->pPatchInfo->dwDataSize = BSWAP_INT32_UNSIGNED(hf->pPatchInfo->dwDataSize);

+        pPatchInfo->dwLength = BSWAP_INT32_UNSIGNED(pPatchInfo->dwLength);

+        pPatchInfo->dwFlags = BSWAP_INT32_UNSIGNED(pPatchInfo->dwFlags);

+        pPatchInfo->dwDataSize = BSWAP_INT32_UNSIGNED(pPatchInfo->dwDataSize);

+

+        // Do nothing if the patch info is not valid

+        if(!(pPatchInfo->dwFlags & MPQ_PATCH_INFO_VALID))

+        {

+            STORM_FREE(pPatchInfo);

+            return ERROR_FILE_CORRUPT;

+        }

 

         // Verify the size of the patch header

         // If it's not default size, we have to reload them

-        if(hf->pPatchInfo->dwLength > dwLength)

+        if(pPatchInfo->dwLength > dwLength)

         {

             // Free the patch info

-            dwLength = hf->pPatchInfo->dwLength;

-            STORM_FREE(hf->pPatchInfo);

-            hf->pPatchInfo = NULL;

+            dwLength = pPatchInfo->dwLength;

+            STORM_FREE(pPatchInfo);

 

             // If the length is out of all possible ranges, fail the operation

             if(dwLength > 0x400)

@@ -1201,16 +1206,17 @@ __AllocateAndLoadPatchInfo:
         }

 

         // Patch file data size according to the patch header

-        hf->dwDataSize = hf->pPatchInfo->dwDataSize;

+        hf->dwDataSize = pPatchInfo->dwDataSize;

     }

     else

     {

-        memset(hf->pPatchInfo, 0, dwLength);

+        memset(pPatchInfo, 0, dwLength);

+        pPatchInfo->dwLength = dwLength;

+        pPatchInfo->dwFlags = MPQ_PATCH_INFO_VALID;

     }

 

     // Save the final length to the patch header

-    hf->pPatchInfo->dwLength = dwLength;

-    hf->pPatchInfo->dwFlags  = 0x80000000;

+    hf->pPatchInfo = pPatchInfo;

     return ERROR_SUCCESS;

 }