From 65bf8c5319c6e1244f3a7f980a15e6c523fb357f Mon Sep 17 00:00:00 2001
From: Bill Zissimopoulos <billziss@navimatics.com>
Date: Tue, 12 Dec 2023 12:36:41 +0000
Subject: [PATCH] dll: use adaptive locks when starting/stopping FSD

---
 src/dll/fsctl.c   | 26 +++++++++++++----
 src/dll/library.h | 15 ++++++++++
 src/dll/util.c    | 71 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 107 insertions(+), 5 deletions(-)

diff --git a/src/dll/fsctl.c b/src/dll/fsctl.c
index 6a3e7d1a..5bbde15f 100644
--- a/src/dll/fsctl.c
+++ b/src/dll/fsctl.c
@@ -431,7 +431,23 @@ FSP_API NTSTATUS FspFsctlServiceVersion(PUINT32 PVersion)
     return 0 != FspFsctlServiceVersionValue ? STATUS_SUCCESS : STATUS_UNSUCCESSFUL;
 }
 
-static SRWLOCK FspFsctlStartStopServiceLock = SRWLOCK_INIT;
+static FSP_ADAPTIVE_LOCK FspFsctlStartStopServiceLock = FSP_ADAPTIVE_LOCK_INIT;
+static VOID FspFsctlStartStopServiceLockAcquire(VOID)
+{
+    extern HINSTANCE DllInstance;
+    WCHAR DllPath[MAX_PATH];
+    PWSTR FileName = 0;
+
+    if (0 != GetModuleFileNameW(DllInstance, DllPath, MAX_PATH))
+        FileName = DllPath;
+
+    FspAdaptiveLockAcquire(&FspFsctlStartStopServiceLock,
+        FileName, 0xfffffffffffffff0ull, (10 + 1) * 1000);
+}
+static VOID FspFsctlStartStopServiceLockRelease(VOID)
+{
+    FspAdaptiveLockRelease(&FspFsctlStartStopServiceLock);
+}
 
 static BOOLEAN FspFsctlRunningInContainer(VOID)
 {
@@ -455,7 +471,7 @@ static NTSTATUS FspFsctlStartServiceByName(PWSTR DriverName)
     DWORD LastError;
     NTSTATUS Result;
 
-    AcquireSRWLockExclusive(&FspFsctlStartStopServiceLock);
+    FspFsctlStartStopServiceLockAcquire();
 
     if (FspFsctlRunningInContainer())
     {
@@ -517,7 +533,7 @@ exit:
     if (0 != ScmHandle)
         CloseServiceHandle(ScmHandle);
 
-    ReleaseSRWLockExclusive(&FspFsctlStartStopServiceLock);
+    FspFsctlStartStopServiceLockRelease();
 
     return Result;
 }
@@ -587,7 +603,7 @@ FSP_API NTSTATUS FspFsctlStopService(VOID)
 
     FspSxsAppendSuffix(DriverName, sizeof DriverName, L"" FSP_FSCTL_DRIVER_NAME);
 
-    AcquireSRWLockExclusive(&FspFsctlStartStopServiceLock);
+    FspFsctlStartStopServiceLockAcquire();
 
     if (FspFsctlRunningInContainer())
     {
@@ -659,7 +675,7 @@ exit:
     if (0 != ProcessToken)
         CloseHandle(ProcessToken);
 
-    ReleaseSRWLockExclusive(&FspFsctlStartStopServiceLock);
+    FspFsctlStartStopServiceLockRelease();
 
     return Result;
 }
diff --git a/src/dll/library.h b/src/dll/library.h
index 28cbeb98..eb436f42 100644
--- a/src/dll/library.h
+++ b/src/dll/library.h
@@ -106,6 +106,21 @@ NTSTATUS FspGetModuleFileName(
     ULONG Size,
     PWSTR RelativePath);
 
+typedef struct
+{
+    SRWLOCK Lock;
+    HANDLE Handle;
+    UINT64 Offset;
+} FSP_ADAPTIVE_LOCK;
+#define FSP_ADAPTIVE_LOCK_INIT          { SRWLOCK_INIT, INVALID_HANDLE_VALUE, 0 }
+VOID FspAdaptiveLockAcquire(
+    FSP_ADAPTIVE_LOCK *Lock,
+    PWSTR FileName,
+    UINT64 Offset,
+    DWORD Timeout);
+VOID FspAdaptiveLockRelease(
+    FSP_ADAPTIVE_LOCK *Lock);
+
 #define FspFileSystemDirectoryBufferEntryInvalid ((ULONG)-1)
 VOID FspFileSystemPeekInDirectoryBuffer(PVOID *PDirBuffer,
     PUINT8 *PBuffer, PULONG *PIndex, PULONG PCount);
diff --git a/src/dll/util.c b/src/dll/util.c
index 78030aad..f810bf17 100644
--- a/src/dll/util.c
+++ b/src/dll/util.c
@@ -418,3 +418,74 @@ NTSTATUS FspGetModuleFileName(
 
     return STATUS_SUCCESS;
 }
+
+VOID FspAdaptiveLockAcquire(
+    FSP_ADAPTIVE_LOCK *Lock,
+    PWSTR FileName,
+    UINT64 Offset,
+    DWORD Timeout)
+{
+    AcquireSRWLockExclusive(&Lock->Lock);
+
+    if (0 != FileName)
+    {
+        HANDLE Handle;
+        DWORD BytesTransferred;
+        OVERLAPPED Overlapped;
+        BOOL Success;
+
+        Handle = CreateFileW(
+            FileName,
+            FILE_READ_DATA,
+            FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
+            0,
+            OPEN_EXISTING,
+            FILE_FLAG_OVERLAPPED,
+            0);
+        if (INVALID_HANDLE_VALUE != Handle)
+        {
+            memset(&Overlapped, 0, sizeof Overlapped);
+            Overlapped.Offset = ((PLARGE_INTEGER)&Offset)->LowPart;
+            Overlapped.OffsetHigh = ((PLARGE_INTEGER)&Offset)->HighPart;
+
+            Success = LockFileEx(
+                Handle,
+                LOCKFILE_EXCLUSIVE_LOCK, 0,
+                1, 0,
+                &Overlapped);
+            if (Success || ERROR_IO_PENDING == GetLastError())
+            {
+                Success = FALSE;
+                if (WAIT_OBJECT_0 == WaitForSingleObject(Handle, Timeout))
+                    Success = GetOverlappedResult(Handle, &Overlapped, &BytesTransferred, TRUE);
+            }
+
+            if (Success)
+            {
+                Lock->Handle = Handle;
+                Lock->Offset = Offset;
+            }
+            else
+                CloseHandle(Handle);
+        }
+    }
+}
+
+VOID FspAdaptiveLockRelease(
+    FSP_ADAPTIVE_LOCK *Lock)
+{
+    if (INVALID_HANDLE_VALUE != Lock->Handle)
+    {
+        HANDLE Handle = Lock->Handle;
+        LARGE_INTEGER LargeOffset = *(PLARGE_INTEGER)&Lock->Offset;
+
+        UnlockFile(Handle, LargeOffset.LowPart, LargeOffset.HighPart, 1, 0);
+
+        CloseHandle(Handle);
+
+        Lock->Handle = INVALID_HANDLE_VALUE;
+        Lock->Offset = 0;
+    }
+
+    ReleaseSRWLockExclusive(&Lock->Lock);
+}