[pull] master from ladislav-zezula:master

src/SBaseCommon.cpp

@@ -931,16 +931,23 @@ TMPQFile * CreateFileHandle(TMPQArchive * ha, TFileEntry * pFileEntry)
         hf->pStream = NULL;
         hf->ha = ha;
 
-        // If the called entered a file entry, we also copy informations from the file entry
-        if(ha != NULL && pFileEntry != NULL)
+        // If the caller entered a file entry, we also copy informations from the file entry
+        if(ha != NULL)
         {
-            // Set the raw position and MPQ position
-            hf->RawFilePos = FileOffsetFromMpqOffset(ha, pFileEntry->ByteOffset);
-            hf->MpqFilePos = pFileEntry->ByteOffset;
+            // Increment number of open files in the archive handle
+            ha->dwFileCount++;
 
-            // Set the data size
-            hf->dwDataSize = pFileEntry->dwFileSize;
-            hf->pFileEntry = pFileEntry;
+            // Add the file entry to the handle
+            if(pFileEntry != NULL)
+            {
+                // Set the raw position and MPQ position
+                hf->RawFilePos = FileOffsetFromMpqOffset(ha, pFileEntry->ByteOffset);
+                hf->MpqFilePos = pFileEntry->ByteOffset;
+
+                // Set the data size
+                hf->dwDataSize = pFileEntry->dwFileSize;
+                hf->pFileEntry = pFileEntry;
+            }
         }
     }
 
@@ -1691,75 +1698,6 @@ DWORD WriteMpqDataMD5(
     return dwErrCode;
 }
 
-// Frees the structure for MPQ file
-void FreeFileHandle(TMPQFile *& hf)
-{
-    if(hf != NULL)
-    {
-        // If we have patch file attached to this one, free it first
-        if(hf->hfPatch != NULL)
-            FreeFileHandle(hf->hfPatch);
-
-        // Then free all buffers allocated in the file structure
-        if(hf->pbFileData != NULL)
-            STORM_FREE(hf->pbFileData);
-        if(hf->pPatchInfo != NULL)
-            STORM_FREE(hf->pPatchInfo);
-        if(hf->SectorOffsets != NULL)
-            STORM_FREE(hf->SectorOffsets);
-        if(hf->SectorChksums != NULL)
-            STORM_FREE(hf->SectorChksums);
-        if(hf->hctx != NULL)
-            STORM_FREE(hf->hctx);
-        if(hf->pbFileSector != NULL)
-            STORM_FREE(hf->pbFileSector);
-        if(hf->pStream != NULL)
-            FileStream_Close(hf->pStream);
-        STORM_FREE(hf);
-        hf = NULL;
-    }
-}
-
-// Frees the MPQ archive
-void FreeArchiveHandle(TMPQArchive *& ha)
-{
-    if(ha != NULL)
-    {
-        // First of all, free the patch archive, if any
-        if(ha->haPatch != NULL)
-            FreeArchiveHandle(ha->haPatch);
-
-        // Free the patch prefix, if any
-        if(ha->pPatchPrefix != NULL)
-            STORM_FREE(ha->pPatchPrefix);
-
-        // Close the file stream
-        FileStream_Close(ha->pStream);
-        ha->pStream = NULL;
-
-        // Free the file names from the file table
-        if(ha->pFileTable != NULL)
-        {
-            for(DWORD i = 0; i < ha->dwFileTableSize; i++)
-            {
-                if(ha->pFileTable[i].szFileName != NULL)
-                    STORM_FREE(ha->pFileTable[i].szFileName);
-                ha->pFileTable[i].szFileName = NULL;
-            }
-
-            // Then free all buffers allocated in the archive structure
-            STORM_FREE(ha->pFileTable);
-        }
-
-        if(ha->pHashTable != NULL)
-            STORM_FREE(ha->pHashTable);
-        if(ha->pHetTable != NULL)
-            FreeHetTable(ha->pHetTable);
-        STORM_FREE(ha);
-        ha = NULL;
-    }
-}
-
 bool IsInternalMpqFileName(const char * szFileName)
 {
     if(szFileName != NULL && szFileName[0] == '(')
@@ -1859,6 +1797,123 @@ void CalculateDataBlockHash(void * pvDataBlock, DWORD cbDataBlock, LPBYTE md5_ha
     md5_done(&md5_state, md5_hash);
 }
 
+//-----------------------------------------------------------------------------
+// Free the handle structures
+
+static void DeleteArchiveHandle(TMPQArchive * ha)
+{
+    // Sanity check
+    assert(ha->dwFileCount == 0 && ha->dwRefCount == 0);
+
+    // Invalidate the add file callback so it won't be called
+    // when saving (listfile) and (attributes)
+    ha->pfnAddFileCB = NULL;
+    ha->pvAddFileUserData = NULL;
+
+    // Flush all unsaved data to the storage
+    SFileFlushArchive((HANDLE)(ha));
+    assert(ha->dwFileCount == 0 && ha->dwRefCount == 0);
+
+    // First of all, free the patch archive, if any
+    if(ha->haPatch != NULL)
+        DereferenceArchive(ha->haPatch);
+    ha->haPatch = NULL;
+
+    // Free the patch prefix, if any
+    if(ha->pPatchPrefix != NULL)
+        STORM_FREE(ha->pPatchPrefix);
+
+    // Close the file stream
+    FileStream_Close(ha->pStream);
+    ha->pStream = NULL;
+
+    // Free the file names from the file table
+    if(ha->pFileTable != NULL)
+    {
+        for(DWORD i = 0; i < ha->dwFileTableSize; i++)
+        {
+            if(ha->pFileTable[i].szFileName != NULL)
+                STORM_FREE(ha->pFileTable[i].szFileName);
+            ha->pFileTable[i].szFileName = NULL;
+        }
+
+        // Then free all buffers allocated in the archive structure
+        STORM_FREE(ha->pFileTable);
+    }
+
+    if(ha->pHashTable != NULL)
+        STORM_FREE(ha->pHashTable);
+    if(ha->pHetTable != NULL)
+        FreeHetTable(ha->pHetTable);
+    STORM_FREE(ha);
+}
+
+void FreeFileHandle(TMPQFile *& hf)
+{
+    TMPQArchive * ha;
+
+    if(hf != NULL)
+    {
+        // If we have patch file attached to this one, free it first
+        if(hf->hfPatch != NULL)
+            FreeFileHandle(hf->hfPatch);
+
+        // Then free all buffers allocated in the file structure
+        if(hf->pbFileData != NULL)
+            STORM_FREE(hf->pbFileData);
+        if(hf->pPatchInfo != NULL)
+            STORM_FREE(hf->pPatchInfo);
+        if(hf->SectorOffsets != NULL)
+            STORM_FREE(hf->SectorOffsets);
+        if(hf->SectorChksums != NULL)
+            STORM_FREE(hf->SectorChksums);
+        if(hf->hctx != NULL)
+            STORM_FREE(hf->hctx);
+        if(hf->pbFileSector != NULL)
+            STORM_FREE(hf->pbFileSector);
+        if(hf->pStream != NULL)
+            FileStream_Close(hf->pStream);
+
+        // Dereference file count in the archive handle
+        if((ha = hf->ha) != NULL)
+            DereferenceArchiveFiles(ha);
+
+        // Free the file handle
+        STORM_FREE(hf);
+        hf = NULL;
+    }
+}
+
+bool DereferenceArchiveFiles(TMPQArchive * ha)
+{
+    // There must be at least one reference
+    if(ha == NULL || ha->dwFileCount == 0)
+        return false;
+
+    // Decrement the file count
+    ha->dwFileCount--;
+
+    // If we reached zero, free the archive
+    if(ha->dwRefCount == 0 && ha->dwFileCount == 0)
+        DeleteArchiveHandle(ha);
+    return true;
+}
+
+bool DereferenceArchive(TMPQArchive * ha)
+{
+    // There must be at least one reference
+    if(ha == NULL || ha->dwRefCount == 0)
+        return false;
+
+    // Decrement the file count
+    ha->dwRefCount--;
+
+    // If we reached zero, free the archive
+    if(ha->dwRefCount == 0 && ha->dwFileCount == 0)
+        DeleteArchiveHandle(ha);
+    return true;
+}
+
 //-----------------------------------------------------------------------------
 // Swapping functions
 

src/SFileCompactArchive.cpp

@@ -462,20 +462,24 @@ DWORD WINAPI SFileGetMaxFileCount(HANDLE hMpq)
 
 bool WINAPI SFileSetMaxFileCount(HANDLE hMpq, DWORD dwMaxFileCount)
 {
-    TMPQArchive * ha = (TMPQArchive *)hMpq;
+    TMPQArchive * ha;
     DWORD dwErrCode = ERROR_SUCCESS;
 
     // Calculate the hash table size for the new file limit
     DWORD dwNewHashTableSize = GetNearestPowerOfTwo(dwMaxFileCount);
 
     // Test the valid parameters
-    if(!IsValidMpqHandle(hMpq))
+    if((ha = IsValidMpqHandle(hMpq)) == NULL)
         dwErrCode = ERROR_INVALID_HANDLE;
     if(ha->dwFlags & MPQ_FLAG_READ_ONLY)
         dwErrCode = ERROR_ACCESS_DENIED;
     if(dwNewHashTableSize < ha->dwFileTableSize)
         dwErrCode = ERROR_DISK_FULL;
 
+    // Are there some files open?
+    if(ha->dwFileCount)
+        dwErrCode = ERROR_ACCESS_DENIED;
+
     // ALL file names must be known in order to be able to rebuild hash table
     if(dwErrCode == ERROR_SUCCESS && ha->pHashTable != NULL)
     {
@@ -538,6 +542,10 @@ bool WINAPI SFileCompactArchive(HANDLE hMpq, const TCHAR * szListFile, bool /* b
     if(ha->dwFlags & MPQ_FLAG_READ_ONLY)
         dwErrCode = ERROR_ACCESS_DENIED;
 
+    // Are there some files open?
+    if(ha->dwFileCount)
+        dwErrCode = ERROR_ACCESS_DENIED;
+
     // If the MPQ is changed at this moment, we have to flush the archive
     if(dwErrCode == ERROR_SUCCESS && (ha->dwFlags & MPQ_FLAG_CHANGED))
     {

src/SFileCreateArchive.cpp

@@ -224,6 +224,7 @@ bool WINAPI SFileCreateArchive2(const TCHAR * szMpqName, PSFILE_CREATE_MPQ pCrea
         ha->dwFileFlags3     = pCreateInfo->dwFileFlags3 ? MPQ_FILE_EXISTS : 0;
         ha->dwAttrFlags      = pCreateInfo->dwAttrFlags;
         ha->dwFlags          = dwMpqFlags | MPQ_FLAG_CHANGED;
+        ha->dwRefCount       = 1;
         pStream = NULL;
 
         // Fill the MPQ header
@@ -274,7 +275,7 @@ bool WINAPI SFileCreateArchive2(const TCHAR * szMpqName, PSFILE_CREATE_MPQ pCrea
     if(dwErrCode != ERROR_SUCCESS)
     {
         FileStream_Close(pStream);
-        FreeArchiveHandle(ha);
+        DereferenceArchive(ha);
         SErrSetLastError(dwErrCode);
         ha = NULL;
     }

src/SFileOpenArchive.cpp

@@ -354,6 +354,9 @@ bool WINAPI SFileOpenArchive(
         // Also remember if this MPQ is a patch
         ha->dwFlags |= (dwFlags & MPQ_OPEN_PATCH) ? MPQ_FLAG_PATCH : 0;
 
+        // Set the reference count to 1
+        ha->dwRefCount = 1;
+
         // Limit the header searching to about 130 MB of data
         if(EndOfSearch > 0x08000000)
             EndOfSearch = 0x08000000;
@@ -626,7 +629,7 @@ bool WINAPI SFileOpenArchive(
     if(dwErrCode != ERROR_SUCCESS)
     {
         FileStream_Close(pStream);
-        FreeArchiveHandle(ha);
+        DereferenceArchive(ha);
         SErrSetLastError(dwErrCode);
         ha = NULL;
     }
@@ -695,20 +698,26 @@ bool WINAPI SFileFlushArchive(HANDLE hMpq)
         // Note that the (signature) file is usually before (listfile) in the file table
         //
 
+        // Up the number of references to prevent it from going away during archive cleanup
+        ha->dwRefCount++;
+
+        // Add the (signature) file
         if(ha->dwFlags & MPQ_FLAG_SIGNATURE_NEW)
         {
             dwErrCode = SSignFileCreate(ha);
             if(dwErrCode != ERROR_SUCCESS)
                 dwResultError = dwErrCode;
         }
 
+        // Add the (listfile) file
         if(ha->dwFlags & (MPQ_FLAG_LISTFILE_NEW | MPQ_FLAG_LISTFILE_FORCE))
         {
             dwErrCode = SListFileSaveToMpq(ha);
             if(dwErrCode != ERROR_SUCCESS)
                 dwResultError = dwErrCode;
         }
 
+        // Add the (attributes) file
         if(ha->dwFlags & MPQ_FLAG_ATTRIBUTES_NEW)
         {
             dwErrCode = SAttrFileSaveToMpq(ha);
@@ -737,6 +746,10 @@ bool WINAPI SFileFlushArchive(HANDLE hMpq)
             }
         }
 
+        // Drop the reference count. Do NOT call DereferenceHandle() to prevent recursion
+        assert(ha->dwRefCount > 0);
+        ha->dwRefCount--;
+
         // We are no longer saving internal MPQ structures
         ha->dwFlags &= ~MPQ_FLAG_SAVING_TABLES;
     }
@@ -754,7 +767,6 @@ bool WINAPI SFileFlushArchive(HANDLE hMpq)
 bool WINAPI SFileCloseArchive(HANDLE hMpq)
 {
     TMPQArchive * ha = IsValidMpqHandle(hMpq);
-    bool bResult = false;
 
     // Only if the handle is valid
     if(ha == NULL)
@@ -763,15 +775,12 @@ bool WINAPI SFileCloseArchive(HANDLE hMpq)
         return false;
     }
 
-    // Invalidate the add file callback so it won't be called
-    // when saving (listfile) and (attributes)
-    ha->pfnAddFileCB = NULL;
-    ha->pvAddFileUserData = NULL;
-
-    // Flush all unsaved data to the storage
-    bResult = SFileFlushArchive(hMpq);
+    // Dereference the archive
+    if(!DereferenceArchive(ha))
+    {
+        SErrSetLastError(ERROR_ACCESS_DENIED);
+        return false;
+    }
 
-    // Free all memory used by MPQ archive
-    FreeArchiveHandle(ha);
-    return bResult;
+    return true;
 }

src/StormCommon.h

@@ -410,8 +410,10 @@ DWORD WriteSectorOffsets(TMPQFile * hf);
 DWORD WriteSectorChecksums(TMPQFile * hf);
 DWORD WriteMemDataMD5(TFileStream * pStream, ULONGLONG RawDataOffs, void * pvRawData, DWORD dwRawDataSize, DWORD dwChunkSize, LPDWORD pcbTotalSize);
 DWORD WriteMpqDataMD5(TFileStream * pStream, ULONGLONG RawDataOffs, DWORD dwRawDataSize, DWORD dwChunkSize);
+
+bool DereferenceArchiveFiles(TMPQArchive * ha);
+bool DereferenceArchive(TMPQArchive * ha);
 void FreeFileHandle(TMPQFile *& hf);
-void FreeArchiveHandle(TMPQArchive *& ha);
 
 //-----------------------------------------------------------------------------
 // Patch functions

src/StormLib.h

@@ -864,6 +864,9 @@ typedef struct _TMPQArchive
     DWORD          dwFlags;                     // See MPQ_FLAG_XXXXX
     DWORD          dwSubType;                   // See MPQ_SUBTYPE_XXX
 
+    DWORD          dwFileCount;                 // Number of open files
+    DWORD          dwRefCount;                  // Number of references
+
     SFILE_ADDFILE_CALLBACK pfnAddFileCB;        // Callback function for adding files
     void         * pvAddFileUserData;           // User data thats passed to the callback