Fix race in IRP_MJ_DEVICE_CONTROL. (#3)

This PR fixes the race described in #2: basically if more than one thread executes the IRP_MJ_DEVICE_CONTROL callback there are a lot of avenues that lead to memory corruption and more.
0vercl0k · Feb 14, 2021 · 703effa · 703effa
1 parent 6b24801
commit 703effa
Show file tree

Hide file tree

Showing 3 changed files with 87 additions and 12 deletions.
diff --git a/src/sic-drv/sic-drv.c b/src/sic-drv/sic-drv.c
@@ -17,14 +17,6 @@
 #include "sic-drv.h"
 #include "..\common\common.h"
 
-//
-// Declare a bunch of functions to satisfy the below pragmas.
-//
-
-DRIVER_INITIALIZE DriverEntry;
-DRIVER_UNLOAD SicDriverUnload;
-DRIVER_DISPATCH_PAGED SicDispatchDeviceControl;
-
 //
 // Our code doesn't need to be allocated in non-paged memory.
 // There is no functions running above APC_LEVEL as a result page faults
@@ -37,6 +29,17 @@ DRIVER_DISPATCH_PAGED SicDispatchDeviceControl;
 #    pragma alloc_text(INIT, DriverEntry)
 #    pragma alloc_text(PAGE, SicDriverUnload)
 #    pragma alloc_text(PAGE, SicDispatchDeviceControl)
+#    pragma alloc_text(PAGE, SicClearAvl)
+#    pragma alloc_text(PAGE, SicGetProcessName)
+#    pragma alloc_text(PAGE, SicGetProcessList)
+#    pragma alloc_text(PAGE, SicDumpVad)
+#    pragma alloc_text(PAGE, SicWalkVadTreeInOrder)
+#    pragma alloc_text(PAGE, SicAvlCompareRoutine)
+#    pragma alloc_text(PAGE, SicAvlAllocateRoutine)
+#    pragma alloc_text(PAGE, SicAvlFreeRoutine)
+#    pragma alloc_text(PAGE, SicSerializeShms)
+#    pragma alloc_text(PAGE, SicFindShms)
+#    pragma alloc_text(PAGE, SicCreateClose)
 #endif
 
 //
@@ -45,6 +48,7 @@ DRIVER_DISPATCH_PAGED SicDispatchDeviceControl;
 
 typedef struct _SIC_CONTEXT
 {
+    FAST_MUTEX Mutex;
     SIC_OFFSETS Offsets;
     RTL_AVL_TABLE ShmsTable;
 } SIC_CONTEXT, *PSIC_CONTEXT;
@@ -964,7 +968,7 @@ Return Value:
 
 _IRQL_requires_(PASSIVE_LEVEL)
 _IRQL_requires_same_ NTSTATUS
-SicFindShms(PDWORD64 OutputSize)
+SicFindShms(_In_ PDWORD64 OutputSize)
 
 /*++
 
@@ -1292,6 +1296,12 @@ Return Value:
 
     PAGED_CODE();
 
+    //
+    // Ensure that only one thread handles a request at a time.
+    //
+
+    ExAcquireFastMutex(&gSicCtx.Mutex);
+
     switch (IoControlCode)
     {
     case IOCTL_SIC_INIT_CONTEXT: {
@@ -1348,6 +1358,12 @@ Return Value:
     }
     }
 
+    //
+    // Release the mutex as we are done modifying global state.
+    //
+
+    ExReleaseFastMutex(&gSicCtx.Mutex);
+
     //
     // We are done with this IRP, so we fill in the IoStatus part.
     //
@@ -1365,7 +1381,7 @@ Return Value:
 
 _IRQL_requires_max_(PASSIVE_LEVEL)
 _IRQL_requires_same_ NTSTATUS
-SicCreateClose(PDEVICE_OBJECT DeviceObject, PIRP Irp)
+SicCreateClose(_In_ PDEVICE_OBJECT DeviceObject, _Inout_ PIRP Irp)
 
 /*++
 
@@ -1484,5 +1500,7 @@ Return Value:
     memset(&gSicCtx, 0, sizeof(gSicCtx));
     RtlInitializeGenericTableAvl(
         &gSicCtx.ShmsTable, SicAvlCompareRoutine, SicAvlAllocateRoutine, SicAvlFreeRoutine, NULL);
+
+    ExInitializeFastMutex(&gSicCtx.Mutex);
     return STATUS_SUCCESS;
 }
diff --git a/src/sic-drv/sic-drv.h b/src/sic-drv/sic-drv.h
@@ -62,3 +62,59 @@ typedef struct _SIC_PROCESS_TO_DEREF
     LIST_ENTRY List;
     PEPROCESS Process;
 } SIC_PROCESS_TO_DEREF, *PSIC_PROCESS_TO_DEREF;
+
+_IRQL_requires_(PASSIVE_LEVEL)
+_IRQL_requires_same_ NTSTATUS
+SicClearAvl(_In_ PRTL_AVL_TABLE);
+
+_Function_class_(DRIVER_UNLOAD) _IRQL_requires_(PASSIVE_LEVEL)
+_IRQL_requires_same_ VOID
+SicDriverUnload(_In_ PDRIVER_OBJECT);
+
+_IRQL_requires_(PASSIVE_LEVEL)
+_IRQL_requires_same_ NTSTATUS
+SicGetProcessName(_In_ const PEPROCESS, _Out_ PUNICODE_STRING *);
+
+_IRQL_requires_(PASSIVE_LEVEL)
+_IRQL_requires_same_ NTSTATUS
+SicGetProcessList(_Out_ PSYSTEM_PROCESS_INFORMATION *);
+
+_IRQL_requires_(PASSIVE_LEVEL)
+_IRQL_requires_same_ NTSTATUS
+SicDumpVad(_In_ const PMMVAD, _Inout_ const PVOID);
+
+_IRQL_requires_(PASSIVE_LEVEL)
+_IRQL_requires_same_ NTSTATUS
+SicWalkVadTreeInOrder(_In_ const PMMVAD, _In_ const SIC_WALK_VAD_ROUTINE, _Inout_opt_ const PVOID);
+
+_IRQL_requires_(PASSIVE_LEVEL)
+_IRQL_requires_same_ RTL_GENERIC_COMPARE_RESULTS
+SicAvlCompareRoutine(_In_ const PRTL_AVL_TABLE, _In_ const PVOID, _In_ const PVOID);
+
+_IRQL_requires_(PASSIVE_LEVEL)
+_IRQL_requires_same_ PVOID
+SicAvlAllocateRoutine(_In_ const PRTL_AVL_TABLE, _In_ const ULONG);
+
+_IRQL_requires_(PASSIVE_LEVEL)
+_IRQL_requires_same_ VOID
+SicAvlFreeRoutine(_In_ const PRTL_AVL_TABLE, _In_ PVOID);
+
+_IRQL_requires_(PASSIVE_LEVEL)
+_IRQL_requires_same_ NTSTATUS
+SicSerializeShms(_In_ const PVOID, _In_ const ULONG, _Out_opt_ PULONG_PTR);
+
+_IRQL_requires_(PASSIVE_LEVEL)
+_IRQL_requires_same_ NTSTATUS
+SicFindShms(_In_ PDWORD64);
+
+_Function_class_(DRIVER_DISPATCH) _IRQL_requires_max_(PASSIVE_LEVEL)
+_IRQL_requires_same_ NTSTATUS
+SicDispatchDeviceControl(_In_ PDEVICE_OBJECT, _Inout_ PIRP);
+
+_IRQL_requires_max_(PASSIVE_LEVEL)
+_IRQL_requires_same_ NTSTATUS
+SicCreateClose(_In_ PDEVICE_OBJECT, _Inout_ PIRP);
+
+_Function_class_(DRIVER_INITIALIZE) _IRQL_requires_same_ _IRQL_requires_(PASSIVE_LEVEL)
+NTSTATUS
+DriverEntry(_In_ PDRIVER_OBJECT, _In_ PUNICODE_STRING);
diff --git a/src/sic/install.cc b/src/sic/install.cc
@@ -77,13 +77,14 @@ bool StopDriver(const char *ServiceName) {
   BOOL Success =
       ControlService(Service, SERVICE_CONTROL_STOP, (LPSERVICE_STATUS)&Status);
   while (Success && Status.dwCurrentState != SERVICE_STOPPED) {
-    printf("Waiting for %u..", Status.dwWaitHint);
+    printf("Waiting for %u..\n", Status.dwWaitHint);
     Sleep(Status.dwWaitHint);
     DWORD BytesNeeded;
     Success =
         QueryServiceStatusEx(Service, SC_STATUS_PROCESS_INFO, (LPBYTE)&Status,
                              sizeof(Status), &BytesNeeded);
-    printf("Success: %d, dwCurrentState: %u..", Success, Status.dwCurrentState);
+    printf("Success: %d, dwCurrentState: %u..\n", Success,
+           Status.dwCurrentState);
   }
 
   return Success && Status.dwCurrentState == SERVICE_STOPPED;