Library - Move UserContext under OpenCount lock

Close is called when all events are done and we get  a kernel close request.
The kernel can issue an event (GetFileInfo) but actually do not wait for its completion and directly send a close. We will have a race condition in ReleaseDokanOpenInfo where the DokanFileInfo might have the GetFileInfo UserContext and not the Cleanup UserContext that is expected to happen before Close.

This change enforce to get the latest UserContext we can get from DokanOpenInfo. Yes, it does not prevent the GetFileInfo to be the latest to set UserContext in EventCompletion but that's the best we can do.

Author: Liryna

dokan/close.c

@@ -35,13 +35,5 @@ VOID DispatchClose(PDOKAN_IO_EVENT IoEvent) {
   // Driver has simply notifying us of the Close request which he has
   // already completed at this stage. Driver is not expecting us
   // to reply from this so there is no need to send an EVENT_INFORMATION.
-
-  if (IoEvent->DokanOpenInfo) {
-    EnterCriticalSection(&IoEvent->DokanOpenInfo->CriticalSection);
-    IoEvent->DokanOpenInfo->FileName =
-        _wcsdup(IoEvent->EventContext->Operation.Close.FileName);
-    IoEvent->DokanOpenInfo->OpenCount--;
-    LeaveCriticalSection(&IoEvent->DokanOpenInfo->CriticalSection);
-    ReleaseDokanOpenInfo(IoEvent);
-  }
+  ReleaseDokanOpenInfo(IoEvent);
 }

dokan/dokan.c

@@ -274,9 +274,8 @@ VOID SetupIOEventForProcessing(PDOKAN_IO_EVENT IoEvent) {
   if (IoEvent->DokanOpenInfo) {
     EnterCriticalSection(&IoEvent->DokanOpenInfo->CriticalSection);
     IoEvent->DokanOpenInfo->OpenCount++;
+    IoEvent->DokanFileInfo.Context = IoEvent->DokanOpenInfo->UserContext;
     LeaveCriticalSection(&IoEvent->DokanOpenInfo->CriticalSection);
-    IoEvent->DokanFileInfo.Context =
-        InterlockedAdd64(&IoEvent->DokanOpenInfo->UserContext, 0);
     IoEvent->DokanFileInfo.IsDirectory =
         (UCHAR)IoEvent->DokanOpenInfo->IsDirectory;
 
@@ -879,11 +878,7 @@ VOID ALIGN_ALLOCATION_SIZE(PLARGE_INTEGER size, PDOKAN_OPTIONS DokanOptions) {
 
 VOID EventCompletion(PDOKAN_IO_EVENT IoEvent) {
   assert(IoEvent->EventResult);
-  if (IoEvent->DokanOpenInfo) {
-    InterlockedExchange64(&IoEvent->DokanOpenInfo->UserContext,
-                          IoEvent->DokanFileInfo.Context);
-    ReleaseDokanOpenInfo(IoEvent);
-  }
+  ReleaseDokanOpenInfo(IoEvent);
 }
 
 VOID CheckFileName(LPWSTR FileName) {
@@ -963,26 +958,38 @@ VOID CreateDispatchCommon(PDOKAN_IO_EVENT IoEvent, ULONG SizeOfEventInfo, BOOL U
 }
 
 VOID ReleaseDokanOpenInfo(PDOKAN_IO_EVENT IoEvent) {
-  LPWSTR fileNameForClose = NULL;
-
+  if (!IoEvent->DokanOpenInfo) {
+    return;
+  }
   EnterCriticalSection(&IoEvent->DokanOpenInfo->CriticalSection);
+  IoEvent->DokanOpenInfo->UserContext = IoEvent->DokanFileInfo.Context;
   IoEvent->DokanOpenInfo->OpenCount--;
-  if (IoEvent->DokanOpenInfo->OpenCount < 1) {
-    if (IoEvent->DokanOpenInfo->FileName) {
-      fileNameForClose = IoEvent->DokanOpenInfo->FileName;
-      IoEvent->DokanOpenInfo->FileName = NULL;
-    }
-    LeaveCriticalSection(&IoEvent->DokanOpenInfo->CriticalSection);
-    PushFileOpenInfo(IoEvent->DokanOpenInfo);
-    IoEvent->DokanOpenInfo = NULL;
-    if (IoEvent->EventResult) {
-      // Reset the Kernel UserContext if we can. Close events do not have one.
-      IoEvent->EventResult->Context = 0;
-    }
-  } else {
+  if (IoEvent->EventContext->MajorFunction == IRP_MJ_CLOSE) {
+    IoEvent->DokanOpenInfo->CloseFileName =
+        _wcsdup(IoEvent->EventContext->Operation.Close.FileName);
+    IoEvent->DokanOpenInfo->CloseUserContext = IoEvent->DokanFileInfo.Context;
+    IoEvent->DokanOpenInfo->OpenCount--;
+  }
+  if (IoEvent->DokanOpenInfo->OpenCount > 0) {
+    // We are still waiting for the Close event or there is another event running. We delay the Close event.
     LeaveCriticalSection(&IoEvent->DokanOpenInfo->CriticalSection);
+    return;
   }
 
+  // Process close event as OpenCount is now 0
+  LPWSTR fileNameForClose = NULL;
+  if (IoEvent->DokanOpenInfo->CloseFileName) {
+    fileNameForClose = IoEvent->DokanOpenInfo->CloseFileName;
+    IoEvent->DokanOpenInfo->CloseFileName = NULL;
+  }
+  IoEvent->DokanFileInfo.Context = IoEvent->DokanOpenInfo->CloseUserContext;
+  LeaveCriticalSection(&IoEvent->DokanOpenInfo->CriticalSection);
+  PushFileOpenInfo(IoEvent->DokanOpenInfo);
+  IoEvent->DokanOpenInfo = NULL;
+  if (IoEvent->EventResult) {
+    // Reset the Kernel UserContext if we can. Close events do not have one.
+    IoEvent->EventResult->Context = 0;
+  }
   if (fileNameForClose) {
     if (IoEvent->DokanInstance->DokanOperations->CloseFile) {
       IoEvent->DokanInstance->DokanOperations->CloseFile(

dokan/dokan_pool.c

@@ -560,13 +560,17 @@ PDOKAN_OPEN_INFO PopFileOpenInfo() {
     RtlZeroMemory(fileInfo, sizeof(DOKAN_OPEN_INFO));
     InitializeCriticalSection(&fileInfo->CriticalSection);
   }
-
   if (fileInfo) {
     fileInfo->DokanInstance = NULL;
     fileInfo->DirList = NULL;
-    InterlockedExchange64(&fileInfo->UserContext, 0);
+    fileInfo->DirListSearchPattern= NULL;
+    fileInfo->UserContext = 0;
     fileInfo->EventId = 0;
     fileInfo->IsDirectory = FALSE;
+    fileInfo->OpenCount = 0;
+    fileInfo->CloseFileName = NULL;
+    fileInfo->CloseUserContext = 0;
+    fileInfo->EventContext = NULL;
   }
   return fileInfo;
 }

dokan/dokani.h

@@ -107,15 +107,16 @@ typedef struct _DOKAN_OPEN_INFO {
   PDOKAN_VECTOR DirList;
   PWCHAR DirListSearchPattern;
   /** User Context see DOKAN_FILE_INFO.Context */
-  volatile LONG64 UserContext;
+  LONG64 UserContext;
   /** Event Id */
   ULONG EventId;
   /** DOKAN_OPTIONS linked to the mount */
   BOOL IsDirectory;
   /** Open count on the file */
   ULONG OpenCount;
   /** Used when dispatching the close once the OpenCount drops to 0 **/
-  LPWSTR FileName;
+  LPWSTR CloseFileName;
+  LONG64 CloseUserContext;
   /** Event context */
   PEVENT_CONTEXT EventContext;
 } DOKAN_OPEN_INFO, *PDOKAN_OPEN_INFO;