Fix bugs in process attach hook

Signed-off-by: Alan Jowett (from Dev Box) <[email protected]>
microsoft · Mar 1, 2024 · e5fd497 · e5fd497
1 parent b273179
commit e5fd497
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 5 deletions.
diff --git a/netebpfext/net_ebpf_ext_process.c b/netebpfext/net_ebpf_ext_process.c
@@ -87,8 +87,13 @@ _net_ebpf_extension_process_on_client_attach(
 
     if (!_ebpf_process_hook_provider_registered) {
         // Register the process create notify routine.
-        if (PsSetCreateProcessNotifyRoutineEx(_ebpf_process_create_process_notify_routine_ex, FALSE) !=
-            STATUS_SUCCESS) {
+        NTSTATUS status = PsSetCreateProcessNotifyRoutineEx(_ebpf_process_create_process_notify_routine_ex, FALSE);
+        if (!NT_SUCCESS(status)) {
+            NET_EBPF_EXT_LOG_MESSAGE_NTSTATUS(
+                NET_EBPF_EXT_TRACELOG_LEVEL_ERROR,
+                NET_EBPF_EXT_TRACELOG_KEYWORD_PROCESS,
+                "PsSetCreateProcessNotifyRoutineEx failed",
+                status);
             result = EBPF_OPERATION_NOT_SUPPORTED;
             goto Exit;
         }
@@ -120,7 +125,13 @@ _net_ebpf_extension_process_on_client_detach(_In_ const net_ebpf_extension_hook_
     _ebpf_process_hook_provider_registration_count--;
 
     if (_ebpf_process_hook_provider_registered && _ebpf_process_hook_provider_registration_count == 0) {
-        if (PsSetCreateProcessNotifyRoutineEx(_ebpf_process_create_process_notify_routine_ex, TRUE) != STATUS_SUCCESS) {
+        NTSTATUS status = PsSetCreateProcessNotifyRoutineEx(_ebpf_process_create_process_notify_routine_ex, TRUE);
+        if (!NT_SUCCESS(status)) {
+            NET_EBPF_EXT_LOG_MESSAGE_NTSTATUS(
+                NET_EBPF_EXT_TRACELOG_LEVEL_ERROR,
+                NET_EBPF_EXT_TRACELOG_KEYWORD_PROCESS,
+                "PsSetCreateProcessNotifyRoutineEx failed",
+                status);
             result = EBPF_OPERATION_NOT_SUPPORTED;
         }
         _ebpf_process_hook_provider_registered = FALSE;
@@ -337,7 +348,9 @@ _ebpf_process_create_process_notify_routine_ex(
                     NET_EBPF_EXT_TRACELOG_KEYWORD_PROCESS,
                     "net_ebpf_extension_hook_invoke_program failed");
             } else {
-                create_info->CreationStatus = process_notify_context.process_md.creation_status;
+                if (create_info != NULL) {
+                    create_info->CreationStatus = process_notify_context.process_md.creation_status;
+                }
             }
             net_ebpf_extension_hook_client_leave_rundown(client_context);
         } else {
@@ -347,7 +360,7 @@ _ebpf_process_create_process_notify_routine_ex(
                 "net_ebpf_extension_hook_client_enter_rundown failed");
         }
         // If the client returns a non-zero value, stop calling the other clients.
-        if (create_info->CreationStatus != STATUS_SUCCESS) {
+        if (create_info && create_info->CreationStatus != STATUS_SUCCESS) {
             break;
         }
 

diff --git a/netebpfext/sys/netebpfext.vcxproj b/netebpfext/sys/netebpfext.vcxproj
@@ -132,6 +132,7 @@
     <Link>
       <AdditionalDependencies>%(AdditionalDependencies);$(DDK_LIB_PATH)\ntoskrnl.lib;$(DDK_LIB_PATH)\ndis.lib;$(DDK_LIB_PATH)\wdmsec.lib;$(DDK_LIB_PATH)\fwpkclnt.lib;$(SDK_LIB_PATH)\uuid.lib;$(DDK_LIB_PATH)\netio.lib</AdditionalDependencies>
       <AdditionalLibraryDirectories>$(OutDir);$(SolutionDir)$(Platform)\$(ConfigurationName)\;%(AdditionalLibraryDirectories)</AdditionalLibraryDirectories>
+      <AdditionalOptions>/INTEGRITYCHECK %(AdditionalOptions)</AdditionalOptions>
     </Link>
     <DriverSign>
       <FileDigestAlgorithm>SHA256</FileDigestAlgorithm>
@@ -154,6 +155,7 @@
     <Link>
       <AdditionalDependencies>%(AdditionalDependencies);$(DDK_LIB_PATH)\ntoskrnl.lib;$(DDK_LIB_PATH)\ndis.lib;$(DDK_LIB_PATH)\wdmsec.lib;$(DDK_LIB_PATH)\fwpkclnt.lib;$(SDK_LIB_PATH)\uuid.lib;$(DDK_LIB_PATH)\netio.lib</AdditionalDependencies>
       <AdditionalLibraryDirectories>$(OutDir);$(SolutionDir)$(Platform)\$(ConfigurationName)\;%(AdditionalLibraryDirectories)</AdditionalLibraryDirectories>
+      <AdditionalOptions>/INTEGRITYCHECK %(AdditionalOptions)</AdditionalOptions>
     </Link>
     <DriverSign>
       <FileDigestAlgorithm>SHA256</FileDigestAlgorithm>
@@ -175,6 +177,7 @@
     <Link>
       <AdditionalDependencies>%(AdditionalDependencies);$(DDK_LIB_PATH)\ntoskrnl.lib;$(DDK_LIB_PATH)\ndis.lib;$(DDK_LIB_PATH)\wdmsec.lib;$(DDK_LIB_PATH)\fwpkclnt.lib;$(SDK_LIB_PATH)\uuid.lib;$(DDK_LIB_PATH)\netio.lib</AdditionalDependencies>
       <AdditionalLibraryDirectories>$(OutDir);$(SolutionDir)$(Platform)\$(ConfigurationName)\;%(AdditionalLibraryDirectories)</AdditionalLibraryDirectories>
+      <AdditionalOptions>/INTEGRITYCHECK %(AdditionalOptions)</AdditionalOptions>
     </Link>
     <DriverSign>
       <FileDigestAlgorithm>SHA256</FileDigestAlgorithm>
@@ -196,6 +199,7 @@
     <Link>
       <AdditionalDependencies>%(AdditionalDependencies);$(DDK_LIB_PATH)\ntoskrnl.lib;$(DDK_LIB_PATH)\ndis.lib;$(DDK_LIB_PATH)\wdmsec.lib;$(DDK_LIB_PATH)\fwpkclnt.lib;$(SDK_LIB_PATH)\uuid.lib;$(DDK_LIB_PATH)\netio.lib</AdditionalDependencies>
       <AdditionalLibraryDirectories>$(OutDir);$(SolutionDir)$(Platform)\$(ConfigurationName)\;%(AdditionalLibraryDirectories)</AdditionalLibraryDirectories>
+      <AdditionalOptions>/INTEGRITYCHECK %(AdditionalOptions)</AdditionalOptions>
     </Link>
     <DriverSign>
       <FileDigestAlgorithm>SHA256</FileDigestAlgorithm>