From 53f60a698a862bb0439e5dda51e277dd468cf8c8 Mon Sep 17 00:00:00 2001
From: Bill Zissimopoulos <billziss@navimatics.com>
Date: Sat, 19 Nov 2016 22:58:40 -0800
Subject: [PATCH] tst: winfsp-tests: add oplock option

---
 tst/winfsp-tests/hooks.c        | 74 +++++++++++++++++++++++++++++++++
 tst/winfsp-tests/winfsp-tests.c | 11 +++++
 tst/winfsp-tests/winfsp-tests.h |  1 +
 3 files changed, 86 insertions(+)

diff --git a/tst/winfsp-tests/hooks.c b/tst/winfsp-tests/hooks.c
index 9e29238d..743aec54 100644
--- a/tst/winfsp-tests/hooks.c
+++ b/tst/winfsp-tests/hooks.c
@@ -103,6 +103,78 @@ static VOID PrepareFileName(PCWSTR FileName, PWSTR FileNameBuf)
     }
 }
 
+typedef struct
+{
+    HANDLE File;
+    HANDLE Wait;
+    OVERLAPPED Overlapped;
+} OPLOCK_BREAK_WAIT_DATA;
+
+static VOID CALLBACK OplockBreakWait(PVOID Context, BOOLEAN Timeout)
+{
+    OPLOCK_BREAK_WAIT_DATA *Data = Context;
+
+    UnregisterWaitEx(Data->Wait, 0);
+    CloseHandle(Data->File);
+    HeapFree(GetProcessHeap(), 0, Data);
+}
+
+static VOID MaybeRequestOplock(PCWSTR FileName)
+{
+    HANDLE File;
+    OPLOCK_BREAK_WAIT_DATA *Data;
+    DWORD FsControlCode;
+    BOOL Success;
+    DWORD BytesTransferred;
+
+    if (0 == OptOplock)
+        return;
+
+    File = CreateFileW(
+        FileName,
+        FILE_READ_ATTRIBUTES | SYNCHRONIZE,
+        FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
+        0,
+        OPEN_EXISTING,
+        FILE_FLAG_OVERLAPPED,
+        0);
+    if (INVALID_HANDLE_VALUE == File)
+        return;
+
+    Data = HeapAlloc(GetProcessHeap(), 0, sizeof *Data);
+    if (0 == Data)
+        ABORT("cannot malloc filter oplock data");
+    memset(Data, 0, sizeof *Data);
+
+    switch (OptOplock)
+    {
+    case 'B':
+        FsControlCode = FSCTL_REQUEST_BATCH_OPLOCK;
+        break;
+    case 'F':
+        FsControlCode = FSCTL_REQUEST_FILTER_OPLOCK;
+        break;
+    default:
+        FsControlCode = FSCTL_REQUEST_FILTER_OPLOCK;
+        break;
+    }
+
+    Success = DeviceIoControl(File, FsControlCode, 0, 0, 0, 0, &BytesTransferred,
+        &Data->Overlapped);
+    if (!Success && ERROR_IO_PENDING == GetLastError())
+    {
+        Data->File = File;
+        if (!RegisterWaitForSingleObject(&Data->Wait, File,
+            OplockBreakWait, Data, INFINITE, WT_EXECUTEONLYONCE))
+            ABORT("cannot register wait for filter oplock");
+    }
+    else
+    {
+        CloseHandle(File);
+        HeapFree(GetProcessHeap(), 0, Data);
+    }
+}
+
 static VOID MaybeAdjustTraversePrivilege(BOOL Enable)
 {
     if (OptNoTraverseToken)
@@ -134,6 +206,8 @@ HANDLE WINAPI HookCreateFileW(
 
     PrepareFileName(lpFileName, FileNameBuf);
 
+    MaybeRequestOplock(lpFileName);
+
     MaybeAdjustTraversePrivilege(FALSE);
     Handle = (OptResilient ? ResilientCreateFileW : CreateFileW)(
         FileNameBuf,
diff --git a/tst/winfsp-tests/winfsp-tests.c b/tst/winfsp-tests/winfsp-tests.c
index b8d186cb..4695c101 100644
--- a/tst/winfsp-tests/winfsp-tests.c
+++ b/tst/winfsp-tests/winfsp-tests.c
@@ -31,6 +31,7 @@ int WinFspNetTests = 1;
 BOOLEAN OptResilient = FALSE;
 BOOLEAN OptCaseInsensitive = FALSE;
 BOOLEAN OptCaseRandomize = FALSE;
+WCHAR OptOplock = 0;
 WCHAR OptMountPointBuf[MAX_PATH], *OptMountPoint;
 WCHAR OptShareNameBuf[MAX_PATH], *OptShareName, *OptShareTarget;
     WCHAR OptShareComputer[MAX_PATH] = L"\\\\localhost\\";
@@ -226,6 +227,16 @@ int main(int argc, char *argv[])
                 OptCaseInsensitive = TRUE;
                 rmarg(argv, argc, argi);
             }
+            else if (0 == strcmp("--oplock=batch", a))
+            {
+                OptOplock = 'B';
+                rmarg(argv, argc, argi);
+            }
+            else if (0 == strcmp("--oplock=filter", a))
+            {
+                OptOplock = 'F';
+                rmarg(argv, argc, argi);
+            }
             else if (0 == strncmp("--mountpoint=", a, sizeof "--mountpoint=" - 1))
             {
                 if (0 != MultiByteToWideChar(CP_UTF8, 0,
diff --git a/tst/winfsp-tests/winfsp-tests.h b/tst/winfsp-tests/winfsp-tests.h
index a813f2e1..e312babf 100644
--- a/tst/winfsp-tests/winfsp-tests.h
+++ b/tst/winfsp-tests/winfsp-tests.h
@@ -140,6 +140,7 @@ extern int WinFspNetTests;
 extern BOOLEAN OptResilient;
 extern BOOLEAN OptCaseInsensitive;
 extern BOOLEAN OptCaseRandomize;
+extern WCHAR OptOplock;
 extern WCHAR OptMountPointBuf[], *OptMountPoint;
 extern WCHAR OptShareNameBuf[], *OptShareName, *OptShareTarget;
     extern WCHAR OptShareComputer[];