+ Fixed crash when adding file after an empty archive has been flushed

5 files changed, 144 insertions, 119 deletions

diff --git a/src/SBaseFileTable.cpp b/src/SBaseFileTable.cpp
index 9e95d3e..a2eb5a8 100644
--- a/src/SBaseFileTable.cpp
+++ b/src/SBaseFileTable.cpp
@@ -1847,8 +1847,8 @@ TFileEntry * AllocateFileEntry(TMPQArchive * ha, const char * szFileName, LCID l
             if((ha->dwFileTableSize + ha->dwReservedFiles) >= ha->dwMaxFileCount)
                 return NULL;
 
-            // Invalidate the internal files. Then we need to search for a deleted entry again,
-            // because the previous call to InvalidateInternalFiles might have created some
+            // Invalidate the internal files so we free
+            // their file entries.
             InvalidateInternalFiles(ha);
 
             // Re-check for deleted entries
@@ -1856,7 +1856,7 @@ TFileEntry * AllocateFileEntry(TMPQArchive * ha, const char * szFileName, LCID l
             if(pFileEntry == NULL)
             {
                 // If there is still no deleted entry, we allocate an entry at the end of the file table
-                assert((ha->dwFileTableSize + ha->dwReservedFiles) < ha->dwMaxFileCount);
+                assert((ha->dwFileTableSize + ha->dwReservedFiles) <= ha->dwMaxFileCount);
                 pFileEntry = ha->pFileTable + ha->dwFileTableSize++;
             }
         }
@@ -1869,7 +1869,7 @@ TFileEntry * AllocateFileEntry(TMPQArchive * ha, const char * szFileName, LCID l
     else
     {
         // There should be at least one entry for that internal file
-        assert((ha->dwFileTableSize + ha->dwReservedFiles) < ha->dwMaxFileCount);
+        assert((ha->dwFileTableSize + ha->dwReservedFiles) <= ha->dwMaxFileCount);
         pFileEntry = ha->pFileTable + ha->dwFileTableSize++;
     }
 
@@ -1985,8 +1985,9 @@ void DeleteFileEntry(
 
 void InvalidateInternalFiles(TMPQArchive * ha)
 {
-    TFileEntry * pFileEntry;
-    DWORD dwReservedFiles = 0;
+    TFileEntry * pFileTableEnd;
+    TFileEntry * pFileEntry1 = NULL;
+    TFileEntry * pFileEntry2 = NULL;
 
     // Do nothing if we are in the middle of saving internal files
     if(!(ha->dwFlags & MPQ_FLAG_SAVING_TABLES))
@@ -1998,54 +1999,49 @@ void InvalidateInternalFiles(TMPQArchive * ha)
         //
 
         // Invalidate the (listfile), if not done yet
-        if(!(ha->dwFlags & MPQ_FLAG_LISTFILE_INVALID))
+        if(ha->dwFileFlags1 != 0 && (ha->dwFlags & MPQ_FLAG_LISTFILE_INVALID) == 0)
         {
-            pFileEntry = GetFileEntryExact(ha, LISTFILE_NAME, LANG_NEUTRAL);
-            if(pFileEntry != NULL)
-            {
-                DeleteFileEntry(ha, pFileEntry);
-                dwReservedFiles++;
-            }
+            // Delete the existing entry for (listfile)
+            pFileEntry1 = GetFileEntryExact(ha, LISTFILE_NAME, LANG_NEUTRAL);
+            if(pFileEntry1 != NULL)
+                DeleteFileEntry(ha, pFileEntry1);
 
+            // Reserve one entry for (listfile)
             ha->dwFlags |= MPQ_FLAG_LISTFILE_INVALID;
+            ha->dwReservedFiles++;
         }
 
         // Invalidate the (attributes), if not done yet
-        if(!(ha->dwFlags & MPQ_FLAG_ATTRIBUTES_INVALID))
+        if(ha->dwFileFlags2 != 0 && (ha->dwFlags & MPQ_FLAG_ATTRIBUTES_INVALID) == 0)
         {
-            pFileEntry = GetFileEntryExact(ha, ATTRIBUTES_NAME, LANG_NEUTRAL);
-            if(pFileEntry != NULL)
-            {
-                DeleteFileEntry(ha, pFileEntry);
-                dwReservedFiles++;
-            }
+            // Delete the existing entry for (attributes)
+            pFileEntry2 = GetFileEntryExact(ha, ATTRIBUTES_NAME, LANG_NEUTRAL);
+            if(pFileEntry2 != NULL)
+                DeleteFileEntry(ha, pFileEntry2);
 
+            // Reserve one entry for (attributes)
             ha->dwFlags |= MPQ_FLAG_ATTRIBUTES_INVALID;
+            ha->dwReservedFiles++;
         }
 
         // If the internal files are at the end of the file table (they usually are),
         // we want to free these 2 entries, so when new files are added, they get
         // added to the freed entries and the internal files get added after that
-        if(ha->dwFileTableSize > 0 && dwReservedFiles != 0)
+        if(ha->dwFileTableSize > 0)
         {
-            // Go backwards while there are free entries
-            for(pFileEntry = ha->pFileTable + ha->dwFileTableSize - 1; pFileEntry >= ha->pFileTable; pFileEntry--)
-            {
-                // Stop searching if a file is present
-                if(pFileEntry->dwFlags & MPQ_FILE_EXISTS)
-                {
-                    pFileEntry++;
-                    break;
-                }
-            }
+            pFileTableEnd = ha->pFileTable + ha->dwFileTableSize;
+
+            // Is one of the entries the last one?
+            if(pFileEntry1 == pFileTableEnd - 1 || pFileEntry2 == pFileTableEnd - 1)
+                pFileTableEnd--;
+            if(pFileEntry1 == pFileTableEnd - 1 || pFileEntry2 == pFileTableEnd - 1)
+                pFileTableEnd--;
 
             // Calculate the new file table size
-            ha->dwFileTableSize = (DWORD)(pFileEntry - ha->pFileTable);
-            ha->dwReservedFiles = dwReservedFiles;
+            ha->dwFileTableSize = (DWORD)(pFileTableEnd - ha->pFileTable);
         }
 
-        // Remember that the MPQ has been changed and it will be necessary
-        // to update the tables
+        // Remember that the MPQ has been changed
         ha->dwFlags |= MPQ_FLAG_CHANGED;
     }
 }
diff --git a/src/SFileAddFile.cpp b/src/SFileAddFile.cpp
index 040ed57..f941141 100644
--- a/src/SFileAddFile.cpp
+++ b/src/SFileAddFile.cpp
@@ -434,6 +434,11 @@ int SFileAddFile_Init(
             // If the caller didn't set MPQ_FILE_REPLACEEXISTING, fail it
             if((dwFlags & MPQ_FILE_REPLACEEXISTING) == 0)
                 nError = ERROR_ALREADY_EXISTS;
+
+            // When replacing an existing file,
+            // we still need to invalidate the (attributes) file
+            if(nError == ERROR_SUCCESS)
+                InvalidateInternalFiles(ha);
         }
     }
 
diff --git a/src/SFileAttributes.cpp b/src/SFileAttributes.cpp
index aafd4c3..ca865a0 100644
--- a/src/SFileAttributes.cpp
+++ b/src/SFileAttributes.cpp
@@ -246,7 +246,7 @@ static LPBYTE CreateAttributesFile(TMPQArchive * ha, DWORD * pcbAttrFile)
     LPBYTE pbAttrFile;
     LPBYTE pbAttrPtr;
     size_t cbAttrFile;
-    DWORD dwFinalEntries = ha->dwFileTableSize + ha->dwReservedFiles + 1;
+    DWORD dwFinalEntries = ha->dwFileTableSize + ha->dwReservedFiles;
 
     // Check if we need patch bits in the (attributes) file
     for(pFileEntry = ha->pFileTable; pFileEntry < pFileTableEnd; pFileEntry++)
@@ -283,7 +283,7 @@ static LPBYTE CreateAttributesFile(TMPQArchive * ha, DWORD * pcbAttrFile)
                 *pArrayCRC32++ = BSWAP_INT32_UNSIGNED(pFileEntry->dwCrc32);
 
             // Skip the reserved entries
-            pbAttrPtr = (LPBYTE)(pArrayCRC32 + ha->dwReservedFiles + 1);
+            pbAttrPtr = (LPBYTE)(pArrayCRC32 + ha->dwReservedFiles);
         }
 
         // Write the array of file time
@@ -296,7 +296,7 @@ static LPBYTE CreateAttributesFile(TMPQArchive * ha, DWORD * pcbAttrFile)
                 *pArrayFileTime++ = BSWAP_INT64_UNSIGNED(pFileEntry->FileTime);
 
             // Skip the reserved entries
-            pbAttrPtr = (LPBYTE)(pArrayFileTime + ha->dwReservedFiles + 1);
+            pbAttrPtr = (LPBYTE)(pArrayFileTime + ha->dwReservedFiles);
         }
 
         // Write the array of MD5s
@@ -312,7 +312,7 @@ static LPBYTE CreateAttributesFile(TMPQArchive * ha, DWORD * pcbAttrFile)
             }
 
             // Skip the reserved items
-            pbAttrPtr = pbArrayMD5 + ((ha->dwReservedFiles + 1) * MD5_DIGEST_SIZE);
+            pbAttrPtr = pbArrayMD5 + (ha->dwReservedFiles * MD5_DIGEST_SIZE);
         }
 
         // Write the array of patch bits
@@ -334,8 +334,8 @@ static LPBYTE CreateAttributesFile(TMPQArchive * ha, DWORD * pcbAttrFile)
                 dwBitMask = (dwBitMask << 0x07) | (dwBitMask >> 0x01);
             }
 
-            // Note: Do not increment the array by the last bit that belongs to (attributes).
-            // This might create the array one byte less (if the number of files a multiplier of 8).
+            // Having incremented the bit array just by the number of items in the file table
+            // will create the bit array one byte less of the number of files is a multiplier of 8).
             // Blizzard MPQs have the same feature.
 
             // Move past the bit array
@@ -410,48 +410,57 @@ int SAttrFileSaveToMpq(TMPQArchive * ha)
     DWORD cbAttrFile = 0;
     int nError = ERROR_SUCCESS;
 
-    // If there are no file flags for (attributes) file, do nothing
-    if(ha->dwFileFlags2 == 0 || ha->dwMaxFileCount == 0)
-        return ERROR_SUCCESS;
+    // Only save the attributes if we should do so
+    if(ha->dwFileFlags2 != 0)
+    {
+        // At this point, we expect to have at least one reserved entry in the file table
+        assert(ha->dwFlags & MPQ_FLAG_ATTRIBUTES_INVALID);
+        assert(ha->dwReservedFiles >= 1);
+
+        // Create the raw data that is to be written to (attributes)
+        // Note: Blizzard MPQs have entries for (listfile) and (attributes),
+        // but they are filled empty
+        pbAttrFile = CreateAttributesFile(ha, &cbAttrFile);
+        if(pbAttrFile != NULL)
+        {
+            // We expect it to be nonzero size
+            assert(cbAttrFile != 0);
+
+            // Determine the real flags for (attributes)
+            if(ha->dwFileFlags2 == MPQ_FILE_EXISTS)
+                ha->dwFileFlags2 = GetDefaultSpecialFileFlags(cbAttrFile, ha->pHeader->wFormatVersion);
+
+            // Create the attributes file in the MPQ
+            nError = SFileAddFile_Init(ha, ATTRIBUTES_NAME,
+                                           0,
+                                           cbAttrFile,
+                                           LANG_NEUTRAL,
+                                           ha->dwFileFlags2 | MPQ_FILE_REPLACEEXISTING,
+                                          &hf);
+
+            // Write the attributes file raw data to it
+            if(nError == ERROR_SUCCESS)
+            {
+                // Write the content of the attributes file to the MPQ
+                nError = SFileAddFile_Write(hf, pbAttrFile, cbAttrFile, MPQ_COMPRESSION_ZLIB);
+                SFileAddFile_Finish(hf);
+            }
 
-    // We expect at least one reserved entry to be there
-    assert(ha->dwReservedFiles >= 1);
-    ha->dwReservedFiles--;
+            // Free the attributes buffer
+            STORM_FREE(pbAttrFile);
+        }
+        else
+        {
+            // If the list file is empty, we assume ERROR_SUCCESS
+            nError = (cbAttrFile == 0) ? ERROR_SUCCESS : ERROR_NOT_ENOUGH_MEMORY;
+        }
 
-    // Create the raw data that is to be written to (attributes)
-    // Note: Blizzard MPQs have entries for (listfile) and (attributes),
-    // but they are filled empty
-    pbAttrFile = CreateAttributesFile(ha, &cbAttrFile);
-    if(pbAttrFile != NULL)
-    {
-        // We expect it to be nonzero size
-        assert(cbAttrFile != 0);
-
-        // Determine the real flags for (attributes)
-        if(ha->dwFileFlags2 == MPQ_FILE_EXISTS)
-            ha->dwFileFlags2 = GetDefaultSpecialFileFlags(cbAttrFile, ha->pHeader->wFormatVersion);
-
-        // Create the attributes file in the MPQ
-        nError = SFileAddFile_Init(ha, ATTRIBUTES_NAME,
-                                       0,
-                                       cbAttrFile,
-                                       LANG_NEUTRAL,
-                                       ha->dwFileFlags2 | MPQ_FILE_REPLACEEXISTING,
-                                      &hf);
-
-        // Write the attributes file raw data to it
+        // If the save process succeeded, we clear the MPQ_FLAG_ATTRIBUTE_INVALID flag
         if(nError == ERROR_SUCCESS)
         {
-            // Write the content of the attributes file to the MPQ
-            nError = SFileAddFile_Write(hf, pbAttrFile, cbAttrFile, MPQ_COMPRESSION_ZLIB);
-            SFileAddFile_Finish(hf);
-
-            // Clear the invalidate flag
             ha->dwFlags &= ~MPQ_FLAG_ATTRIBUTES_INVALID;
+            ha->dwReservedFiles--;
         }
-
-        // Free the attributes buffer
-        STORM_FREE(pbAttrFile);
     }
     
     return nError;
diff --git a/src/SFileCreateArchive.cpp b/src/SFileCreateArchive.cpp
index a922e90..569b953 100644
--- a/src/SFileCreateArchive.cpp
+++ b/src/SFileCreateArchive.cpp
@@ -107,6 +107,7 @@ bool WINAPI SFileCreateArchive2(const TCHAR * szMpqName, PSFILE_CREATE_MPQ pCrea
     DWORD dwBlockTableSize = 0;             // Initial block table size
     DWORD dwHashTableSize = 0;
     DWORD dwReservedFiles = 0;              // Number of reserved file entries
+    DWORD dwMpqFlags = 0;
     int nError = ERROR_SUCCESS;
 
     // Check the parameters, if they are valid
@@ -154,11 +155,19 @@ bool WINAPI SFileCreateArchive2(const TCHAR * szMpqName, PSFILE_CREATE_MPQ pCrea
             return false;
     }
 
-    // Increment the maximum amount of files to have space for (listfile) and (attributes)
+    // Increment the maximum amount of files to have space for (listfile)
     if(pCreateInfo->dwMaxFileCount && pCreateInfo->dwFileFlags1)
+    {
+        dwMpqFlags |= MPQ_FLAG_LISTFILE_INVALID;
         dwReservedFiles++;
+    }
+
+    // Increment the maximum amount of files to have space for (attributes)
     if(pCreateInfo->dwMaxFileCount && pCreateInfo->dwFileFlags2 && pCreateInfo->dwAttrFlags)
+    {
+        dwMpqFlags |= MPQ_FLAG_ATTRIBUTES_INVALID;
         dwReservedFiles++;
+    }
 
     // If file count is not zero, initialize the hash table size
     dwHashTableSize = GetHashTableSizeForFileCount(pCreateInfo->dwMaxFileCount + dwReservedFiles);
@@ -197,7 +206,7 @@ bool WINAPI SFileCreateArchive2(const TCHAR * szMpqName, PSFILE_CREATE_MPQ pCrea
         ha->dwFileFlags1    = pCreateInfo->dwFileFlags1;
         ha->dwFileFlags2    = pCreateInfo->dwFileFlags2;
         ha->dwAttrFlags     = pCreateInfo->dwAttrFlags;
-        ha->dwFlags         = 0;
+        ha->dwFlags         = dwMpqFlags | MPQ_FLAG_CHANGED;
         pStream = NULL;
 
         // Fill the MPQ header
@@ -219,9 +228,6 @@ bool WINAPI SFileCreateArchive2(const TCHAR * szMpqName, PSFILE_CREATE_MPQ pCrea
 
         // Write the naked MPQ header
         nError = WriteNakedMPQHeader(ha);
-
-        // Remember that the (listfile) and (attributes) need to be saved
-        ha->dwFlags |= MPQ_FLAG_CHANGED | MPQ_FLAG_LISTFILE_INVALID | MPQ_FLAG_ATTRIBUTES_INVALID;
     }
 
     // Create initial HET table, if the caller required an MPQ format 3.0 or newer
diff --git a/src/SFileListFile.cpp b/src/SFileListFile.cpp
index b2db0c5..ed549d1 100644
--- a/src/SFileListFile.cpp
+++ b/src/SFileListFile.cpp
@@ -363,49 +363,58 @@ int SListFileSaveToMpq(TMPQArchive * ha)
     DWORD cbListFile = 0;
     int nError = ERROR_SUCCESS;
 
-    // Only save (listfile) if we should do so
-    if(ha->dwFileFlags1 == 0 || ha->dwMaxFileCount == 0)
-        return ERROR_SUCCESS;
+    // Only save the listfile if we should do so
+    if(ha->dwFileFlags1 != 0)
+    {
+        // At this point, we expect to have at least one reserved entry in the file table
+        assert(ha->dwFlags & MPQ_FLAG_LISTFILE_INVALID);
+        assert(ha->dwReservedFiles >= 1);
+
+        // Create the raw data that is to be written to (listfile)
+        // Note: Creating the raw data before the (listfile) has been created in the MPQ
+        // causes that the name of the listfile will not be included in the listfile itself.
+        // That is OK, because (listfile) in Blizzard MPQs does not contain it either.
+        pbListFile = CreateListFile(ha, &cbListFile);
+        if(pbListFile != NULL)
+        {
+            // We expect it to be nonzero size
+            assert(cbListFile != 0);
+
+            // Determine the real flags for (listfile)
+            if(ha->dwFileFlags1 == MPQ_FILE_EXISTS)
+                ha->dwFileFlags1 = GetDefaultSpecialFileFlags(cbListFile, ha->pHeader->wFormatVersion);
+
+            // Create the listfile in the MPQ
+            nError = SFileAddFile_Init(ha, LISTFILE_NAME,
+                                           0,
+                                           cbListFile,
+                                           LANG_NEUTRAL,
+                                           ha->dwFileFlags1 | MPQ_FILE_REPLACEEXISTING,
+                                          &hf);
+
+            // Write the listfile raw data to it
+            if(nError == ERROR_SUCCESS)
+            {
+                // Write the content of the listfile to the MPQ
+                nError = SFileAddFile_Write(hf, pbListFile, cbListFile, MPQ_COMPRESSION_ZLIB);
+                SFileAddFile_Finish(hf);
+            }
 
-    // At this point, we expect to have at least one reserved entry in the file table
-    assert(ha->dwReservedFiles >= 1);
-    ha->dwReservedFiles--;
+            // Free the listfile buffer
+            STORM_FREE(pbListFile);
+        }
+        else
+        {
+            // If the list file is empty, we assume ERROR_SUCCESS
+            nError = (cbListFile == 0) ? ERROR_SUCCESS : ERROR_NOT_ENOUGH_MEMORY;
+        }
 
-    // Create the raw data that is to be written to (listfile)
-    // Note: Creating the raw data before the (listfile) has been created in the MPQ
-    // causes that the name of the listfile will not be included in the listfile itself.
-    // That is OK, because (listfile) in Blizzard MPQs does not contain it either.
-    pbListFile = CreateListFile(ha, &cbListFile);
-    if(pbListFile != NULL)
-    {
-        // We expect it to be nonzero size
-        assert(cbListFile != 0);
-
-        // Determine the real flags for (listfile)
-        if(ha->dwFileFlags1 == MPQ_FILE_EXISTS)
-            ha->dwFileFlags1 = GetDefaultSpecialFileFlags(cbListFile, ha->pHeader->wFormatVersion);
-
-        // Create the listfile in the MPQ
-        nError = SFileAddFile_Init(ha, LISTFILE_NAME,
-                                       0,
-                                       cbListFile,
-                                       LANG_NEUTRAL,
-                                       ha->dwFileFlags1 | MPQ_FILE_REPLACEEXISTING,
-                                      &hf);
-
-        // Write the listfile raw data to it
+        // If the save process succeeded, we clear the MPQ_FLAG_LISTFILE_INVALID flag
         if(nError == ERROR_SUCCESS)
         {
-            // Write the content of the listfile to the MPQ
-            nError = SFileAddFile_Write(hf, pbListFile, cbListFile, MPQ_COMPRESSION_ZLIB);
-            SFileAddFile_Finish(hf);
-
-            // Clear the invalidate flag
             ha->dwFlags &= ~MPQ_FLAG_LISTFILE_INVALID;
+            ha->dwReservedFiles--;
         }
-
-        // Free the listfile buffer
-        STORM_FREE(pbListFile);
     }
 
     return nError;