Added exclusive locking.

googleprojectzero · Feb 27, 2017 · 74a11f6 · 74a11f6
1 parent 1a3bec5
commit 74a11f6
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 10 deletions.
diff --git a/BaitAndSwitch/BaitAndSwitch.cpp b/BaitAndSwitch/BaitAndSwitch.cpp
@@ -29,7 +29,12 @@ int _tmain(int argc, _TCHAR* argv[])
 {
 	if (argc < 3)
 	{
-		printf("Usage: BaitAndSwitch symlink target1 target2 [rwd]\n");
+		printf("Usage: BaitAndSwitch symlink target1 target2 [rwdx]\n");
+    printf("Share Mode:\n");
+    printf("r - FILE_SHARE_READ\n");
+    printf("w - FILE_SHARE_WRITE\n");
+    printf("d - FILE_SHARE_DELETE\n");
+    printf("x - Exclusive lock\n");
 		return 1;
 	}
 	else

diff --git a/CommonUtils/FileOpLock.cpp b/CommonUtils/FileOpLock.cpp
@@ -52,7 +52,7 @@ FileOpLock::~FileOpLock()
 	}
 }
 
-bool FileOpLock::BeginLock(const std::wstring& filename, DWORD dwShareMode)
+bool FileOpLock::BeginLock(const std::wstring& filename, DWORD dwShareMode, bool exclusive)
 {
 	g_hLockCompleted = CreateEvent(nullptr, TRUE, FALSE, nullptr);
 	g_o.hEvent = CreateEvent(nullptr, FALSE, FALSE, nullptr);
@@ -61,7 +61,7 @@ bool FileOpLock::BeginLock(const std::wstring& filename, DWORD dwShareMode)
 
 	if (GetFileAttributesW(filename.c_str()) & FILE_ATTRIBUTE_DIRECTORY)
 	{
-		flags |= FILE_FLAG_BACKUP_SEMANTICS;
+		flags |= FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OPEN_REPARSE_POINT;
 	}
 
 	g_hFile = CreateFileW(filename.c_str(), GENERIC_READ,
@@ -81,10 +81,25 @@ bool FileOpLock::BeginLock(const std::wstring& filename, DWORD dwShareMode)
 
 	SetThreadpoolWait(g_wait, g_o.hEvent, nullptr);
 
-	DeviceIoControl(g_hFile, FSCTL_REQUEST_OPLOCK,
-		&g_inputBuffer, sizeof(g_inputBuffer),
-		&g_outputBuffer, sizeof(g_outputBuffer),
-		nullptr, &g_o);
+	DWORD bytesReturned;
+
+  if (exclusive)
+  {
+    DeviceIoControl(g_hFile,
+      FSCTL_REQUEST_OPLOCK_LEVEL_1,
+      NULL, 0,
+      NULL, 0,
+      &bytesReturned,
+      &g_o);
+  }
+  else
+  {
+    DeviceIoControl(g_hFile, FSCTL_REQUEST_OPLOCK,
+      &g_inputBuffer, sizeof(g_inputBuffer),
+      &g_outputBuffer, sizeof(g_outputBuffer),
+      nullptr, &g_o);
+  }
+
 	DWORD err = GetLastError();
 	if (err != ERROR_IO_PENDING) {
 		DebugPrintf("Oplock Failed %d\n", err);
@@ -98,6 +113,7 @@ FileOpLock* FileOpLock::CreateLock(const std::wstring& name, const std::wstring&
 {
 	FileOpLock* ret = new FileOpLock(cb);
 	DWORD dwShareMode = 0;
+  bool exclusive = false;
 
 	if (share_mode.find('r') != std::wstring::npos)
 	{
@@ -114,7 +130,12 @@ FileOpLock* FileOpLock::CreateLock(const std::wstring& name, const std::wstring&
 		dwShareMode |= FILE_SHARE_DELETE;
 	}
 
-	if (ret->BeginLock(name, dwShareMode))
+  if (share_mode.find('x') != std::wstring::npos)
+  {
+    exclusive = true;
+  }
+
+	if (ret->BeginLock(name, dwShareMode, exclusive))
 	{
 		return ret;
 	}

diff --git a/CommonUtils/FileOpLock.h b/CommonUtils/FileOpLock.h
@@ -30,7 +30,7 @@ class FileOpLock
 
 	void DoWaitCallback();
 
-	bool BeginLock(const std::wstring& name, DWORD dwShareMode);
+	bool BeginLock(const std::wstring& name, DWORD dwShareMode, bool exclusive);
 
 };
 
diff --git a/SetOpLock/SetOpLock.cpp b/SetOpLock/SetOpLock.cpp
@@ -27,7 +27,12 @@ int _tmain(int argc, _TCHAR* argv[])
 {	
 	if (argc < 2)
 	{
-		printf("Usage: SetOpLock target [rwd]\n");
+		printf("Usage: SetOpLock target [rwdx]\n");
+    printf("Share Mode:\n");
+    printf("r - FILE_SHARE_READ\n");
+    printf("w - FILE_SHARE_WRITE\n");
+    printf("d - FILE_SHARE_DELETE\n");
+    printf("x - Exclusive lock\n");
 		return 1;
 	}
 	else