DriverLoader: Refactor initialization

src/driver/fortcb.c

@@ -17,11 +17,7 @@ FORT_API FortCallbackFunc fort_callback(int id, FortCallbackFunc func)
 
 FORT_API void fort_callback_setup(PFORT_PROXYCB_INFO cb_info)
 {
-    if (cb_info == NULL) {
+    fort_proxycb_dst_setup(cb_info);
-        g_callbackInfo.src = NULL;
-    } else {
-        fort_proxycb_dst_setup(cb_info);
 
-        g_callbackInfo = *cb_info;
+    g_callbackInfo = *cb_info;
-    }
 }

src/driver/fortdrv.c

@@ -63,6 +63,8 @@ static NTSTATUS fort_driver_load(PDRIVER_OBJECT driver, PUNICODE_STRING reg_path
     if (!NT_SUCCESS(status))
         return status;
 
+    device_obj->Flags |= DO_BUFFERED_IO;
+
     UNICODE_STRING device_link;
 
     RtlInitUnicodeString(&device_link, FORT_DOS_DEVICE_NAME);
@@ -76,29 +78,27 @@ static NTSTATUS fort_driver_load(PDRIVER_OBJECT driver, PUNICODE_STRING reg_path
     driver->MajorFunction[IRP_MJ_CLEANUP] = fort_device_cleanup;
     driver->MajorFunction[IRP_MJ_DEVICE_CONTROL] = fort_device_control;
 
-    device_obj->Flags |= DO_BUFFERED_IO;
-
     return fort_device_load(device_obj);
 }
 
-NTSTATUS __declspec(dllexport) DriverCallbackEntry(
+NTSTATUS __declspec(dllexport) DriverCallbacksSetup(PFORT_PROXYCB_INFO cb_info)
-        PDRIVER_OBJECT driver, PUNICODE_STRING reg_path, PFORT_PROXYCB_INFO cb_info)
+{
+    fort_callback_setup(cb_info);
+
+    return STATUS_SUCCESS;
+}
+
+NTSTATUS __declspec(dllexport) DriverEntry(PDRIVER_OBJECT driver, PUNICODE_STRING reg_path)
 {
     NTSTATUS status;
 
-    fort_callback_setup(cb_info);
-
     status = fort_driver_load(driver, reg_path);
 
     if (!NT_SUCCESS(status)) {
         DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL, "FORT: Entry: Error: %x\n", status);
+
         fort_driver_unload(driver);
     }
 
     return status;
 }
-
-NTSTATUS DriverEntry(PDRIVER_OBJECT driver, PUNICODE_STRING reg_path)
-{
-    return DriverCallbackEntry(driver, reg_path, NULL);
-}

src/driver/loader/fortdl.c

@@ -39,19 +39,24 @@ static NTSTATUS fort_loader_entry(PDRIVER_OBJECT driver, PUNICODE_STRING regPath
     FORT_PROXYCB_INFO cbInfo;
     fort_proxycb_src_prepare(&cbInfo);
 
-    /* Run the module entry function */
+    /* Setup the module's callbacks */
-    status = CallModuleEntry(&g_loader.module, driver, regPath, &cbInfo);
+    status = SetupModuleCallbacks(&g_loader.module, &cbInfo);
     if (!NT_SUCCESS(status))
         return status;
 
     /* Setup the proxy callbacks */
     fort_proxycb_src_setup(&cbInfo);
 
-    /* Proxy the driver major functions */
+    /* Run the module's entry function */
+    status = CallModuleEntry(&g_loader.module, driver, regPath);
+    if (!NT_SUCCESS(status))
+        return status;
+
+    /* Proxy the driver's unload function */
     g_loader.DriverUnload = driver->DriverUnload;
     driver->DriverUnload = fort_loader_unload;
 
-    /* Setup the major functions */
+    /* Proxy the driver's major functions */
     fort_proxycb_drv_setup(driver->MajorFunction);
 
     return status;

src/driver/loader/fortmm.c

@@ -21,8 +21,9 @@
 #define fort_nt_headers(pImage)                                                                    \
     ((PIMAGE_NT_HEADERS) & ((PUCHAR) (pImage))[((PIMAGE_DOS_HEADER) pImage)->e_lfanew])
 
-typedef NTSTATUS(WINAPI *DriverCallbackEntryProc)(
+typedef NTSTATUS(WINAPI *DriverCallbacksSetupProc)(PFORT_PROXYCB_INFO cbInfo);
-        PDRIVER_OBJECT driver, PUNICODE_STRING regPath, PFORT_PROXYCB_INFO cbInfo);
+
+typedef NTSTATUS(WINAPI *DriverEntryProc)(PDRIVER_OBJECT driver, PUNICODE_STRING regPath);
 
 static NTSTATUS GetModuleInfo(PLOADEDMODULE pModule, LPCSTR name,
         const PAUX_MODULE_EXTENDED_INFO modules, DWORD modulesCount)
@@ -419,19 +420,32 @@ FORT_API void UnloadModule(PLOADEDMODULE pModule)
     }
 }
 
-FORT_API NTSTATUS CallModuleEntry(PLOADEDMODULE pModule, PDRIVER_OBJECT driver,
+NTSTATUS SetupModuleCallbacks(PLOADEDMODULE pModule, PFORT_PROXYCB_INFO cbInfo)
-        PUNICODE_STRING regPath, PFORT_PROXYCB_INFO cbInfo)
 {
-    DriverCallbackEntryProc driverEntry =
+    DriverCallbacksSetupProc cbSetup =
-            (DriverCallbackEntryProc) ModuleGetProcAddress(pModule, "DriverCallbackEntry");
+            (DriverCallbacksSetupProc) ModuleGetProcAddress(pModule, "DriverCallbacksSetup");
+    if (cbSetup == NULL)
+        return STATUS_PROCEDURE_NOT_FOUND;
+
+    DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
+            "FORT: Loader Module: Setup Callbacks: %p data=%x\n", cbSetup,
+            *(PDWORD) (PVOID) &cbSetup);
+
+    return cbSetup(cbInfo);
+}
+
+FORT_API NTSTATUS CallModuleEntry(
+        PLOADEDMODULE pModule, PDRIVER_OBJECT driver, PUNICODE_STRING regPath)
+{
+    DriverEntryProc driverEntry = (DriverEntryProc) ModuleGetProcAddress(pModule, "DriverEntry");
     if (driverEntry == NULL)
         return STATUS_PROCEDURE_NOT_FOUND;
 
     DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
-            "FORT: Loader Module: Entry Proc: %p data=%x\n", driverEntry,
+            "FORT: Loader Module: Driver Entry: %p data=%x\n", driverEntry,
             *(PDWORD) (PVOID) &driverEntry);
 
-    return driverEntry(driver, regPath, cbInfo);
+    return driverEntry(driver, regPath);
 }
 
 /* Retrieve address of an exported function from the loaded module. */

src/driver/loader/fortmm.h

@@ -18,8 +18,10 @@ FORT_API NTSTATUS LoadModuleFromMemory(PLOADEDMODULE pModule, PUCHAR lpData, DWO
 
 FORT_API void UnloadModule(PLOADEDMODULE pModule);
 
-FORT_API NTSTATUS CallModuleEntry(PLOADEDMODULE pModule, PDRIVER_OBJECT driver,
+FORT_API NTSTATUS SetupModuleCallbacks(PLOADEDMODULE pModule, PFORT_PROXYCB_INFO cbInfo);
-        PUNICODE_STRING regPath, PFORT_PROXYCB_INFO cbInfo);
+
+FORT_API NTSTATUS CallModuleEntry(
+        PLOADEDMODULE pModule, PDRIVER_OBJECT driver, PUNICODE_STRING regPath);
 
 FORT_API FARPROC ModuleGetProcAddress(PLOADEDMODULE pModule, LPCSTR funcName);
 

src/driver/proxycb/fortpcb_def.h

@@ -7,7 +7,7 @@
 extern "C" {
 #endif
 
-typedef void (WINAPI *ProxyCallbackProc)(void);
+typedef void(WINAPI *ProxyCallbackProc)(void);
 
 typedef struct fort_proxycb_info
 {

src/driver/proxycb/fortpcb_drv.c

@@ -88,7 +88,7 @@ FORT_API void fort_proxycb_drv_setup(PDRIVER_DISPATCH *driver_major_funcs)
             driver_major_funcs[i] = g_proxyMajorCallbacks[i];
 
             DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
-                    "FORT: fort_proxycb_drv_setup: %d mf=%p\n", i, major_func);
+                    "FORT: Proxy Major: i=%d func=%p\n", i, major_func);
         }
     }
 }