From ecf822dedd24ec3e54d90451782df7d84b2b951c Mon Sep 17 00:00:00 2001
From: Nodir Temirkhodjaev <nodir.temir@gmail.com>
Date: Fri, 24 Feb 2023 12:40:35 +0300
Subject: [PATCH] Driver: fortcout: Simplify fort_callout_install()

---
 src/driver/fortcout.c | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/driver/fortcout.c b/src/driver/fortcout.c
index 2f55b7a0..f57cd7cb 100644
--- a/src/driver/fortcout.c
+++ b/src/driver/fortcout.c
@@ -613,7 +613,7 @@ inline static NTSTATUS fort_callout_register(
     return STATUS_SUCCESS;
 }
 
-FORT_API NTSTATUS fort_callout_install(PDEVICE_OBJECT device)
+static NTSTATUS fort_callout_install_ale(PDEVICE_OBJECT device)
 {
     const FWPS_CALLOUT0 callouts[] = {
         /* IPv4 connect callout */
@@ -640,6 +640,15 @@ FORT_API NTSTATUS fort_callout_install(PDEVICE_OBJECT device)
                 .classifyFn = &fort_callout_accept_v6,
                 .notifyFn = &fort_callout_notify,
         },
+    };
+
+    return fort_callout_register(
+            device, callouts, /*count=*/sizeof(callouts) / sizeof(callouts[0]));
+}
+
+static NTSTATUS fort_callout_install_stream(PDEVICE_OBJECT device)
+{
+    const FWPS_CALLOUT0 callouts[] = {
         /* IPv4 stream callout */
         {
                 .calloutKey = FORT_GUID_CALLOUT_STREAM_V4,
@@ -672,6 +681,15 @@ FORT_API NTSTATUS fort_callout_install(PDEVICE_OBJECT device)
                 .flowDeleteFn = &fort_callout_flow_delete,
                 .flags = FWP_CALLOUT_FLAG_CONDITIONAL_ON_FLOW,
         },
+    };
+
+    return fort_callout_register(
+            device, callouts, /*count=*/sizeof(callouts) / sizeof(callouts[0]));
+}
+
+static NTSTATUS fort_callout_install_transport(PDEVICE_OBJECT device)
+{
+    const FWPS_CALLOUT0 callouts[] = {
         /* IPv4 inbound transport callout */
         {
                 .calloutKey = FORT_GUID_CALLOUT_IN_TRANSPORT_V4,
@@ -710,6 +728,18 @@ FORT_API NTSTATUS fort_callout_install(PDEVICE_OBJECT device)
             device, callouts, /*count=*/sizeof(callouts) / sizeof(callouts[0]));
 }
 
+FORT_API NTSTATUS fort_callout_install(PDEVICE_OBJECT device)
+{
+    NTSTATUS status;
+    if (!NT_SUCCESS(status = fort_callout_install_ale(device))
+            || !NT_SUCCESS(status = fort_callout_install_stream(device))
+            || !NT_SUCCESS(status = fort_callout_install_transport(device))) {
+        return status;
+    }
+
+    return STATUS_SUCCESS;
+}
+
 FORT_API void fort_callout_remove(void)
 {
     PFORT_STAT stat = &fort_device()->stat;