From b72cf2577a672fe80632572cf7e317d8d51d7ad8 Mon Sep 17 00:00:00 2001
From: Vladimir Stoilov <vladimir@safing.io>
Date: Tue, 4 Jun 2024 14:47:56 +0300
Subject: [PATCH] [windows_kext] Fix ALE layer TCP conn pend

---
 windows_kext/driver/src/ale_callouts.rs          | 16 +---------------
 windows_kext/driver/src/id_cache.rs              |  2 +-
 .../wdk/src/filter_engine/callout_data.rs        |  9 +++++++--
 windows_kext/wdk/src/filter_engine/packet.rs     |  8 +++++---
 4 files changed, 14 insertions(+), 21 deletions(-)

diff --git a/windows_kext/driver/src/ale_callouts.rs b/windows_kext/driver/src/ale_callouts.rs
index c45b7f4c..4d4e8a7e 100644
--- a/windows_kext/driver/src/ale_callouts.rs
+++ b/windows_kext/driver/src/ale_callouts.rs
@@ -287,21 +287,7 @@ fn save_packet(
     ale_data: &AleLayerData,
     pend: bool,
 ) -> Result<Packet, alloc::string::String> {
-    let mut packet_list = None;
-    let mut save_packet_list = true;
-    match ale_data.protocol {
-        IpProtocol::Tcp => {
-            if let Direction::Outbound = ale_data.direction {
-                // Only time a packet data is missing is during connect state of outbound TCP connection.
-                // Don't save packet list only if connection is outbound, reauthorize is false and the protocol is TCP.
-                save_packet_list = ale_data.reauthorize;
-            }
-        }
-        _ => {}
-    };
-    if save_packet_list {
-        packet_list = create_packet_list(device, callout_data, ale_data);
-    }
+    let packet_list = create_packet_list(device, callout_data, ale_data);
     if pend && matches!(ale_data.protocol, IpProtocol::Tcp | IpProtocol::Udp) {
         match callout_data.pend_operation(packet_list) {
             Ok(classify_defer) => Ok(Packet::AleLayer(classify_defer)),
diff --git a/windows_kext/driver/src/id_cache.rs b/windows_kext/driver/src/id_cache.rs
index e8474538..3b791ff9 100644
--- a/windows_kext/driver/src/id_cache.rs
+++ b/windows_kext/driver/src/id_cache.rs
@@ -65,7 +65,7 @@ fn get_payload<'a>(packet: &'a Packet) -> Option<&'a [u8]> {
                 wdk::filter_engine::callout_data::ClassifyDefer::Reauthorization(_, p) => p,
             };
             if let Some(tpl) = p {
-                tpl.net_buffer_list_queue.get_data()
+                tpl.net_buffer_list.get_data()
             } else {
                 None
             }
diff --git a/windows_kext/wdk/src/filter_engine/callout_data.rs b/windows_kext/wdk/src/filter_engine/callout_data.rs
index c09be368..82039e47 100644
--- a/windows_kext/wdk/src/filter_engine/callout_data.rs
+++ b/windows_kext/wdk/src/filter_engine/callout_data.rs
@@ -32,8 +32,13 @@ impl ClassifyDefer {
         unsafe {
             match self {
                 ClassifyDefer::Initial(context, packet_list) => {
-                    FwpsCompleteOperation0(context, core::ptr::null_mut());
-                    return Ok(packet_list);
+                    if let Some(packet) = packet_list {
+                        FwpsCompleteOperation0(context, packet.net_buffer_list.nbl as _);
+                        return Ok(Some(packet));
+                    } else {
+                        FwpsCompleteOperation0(context, core::ptr::null_mut());
+                    }
+                    return Ok(None);
                 }
                 ClassifyDefer::Reauthorization(_callout_id, packet_list) => {
                     // There is no way to reset single filter. If another request for filter reset is trigger at the same time it will fail.
diff --git a/windows_kext/wdk/src/filter_engine/packet.rs b/windows_kext/wdk/src/filter_engine/packet.rs
index f6b9018f..85e26006 100644
--- a/windows_kext/wdk/src/filter_engine/packet.rs
+++ b/windows_kext/wdk/src/filter_engine/packet.rs
@@ -24,7 +24,7 @@ use super::{callout_data::CalloutData, net_buffer::NetBufferList};
 
 pub struct TransportPacketList {
     ipv6: bool,
-    pub net_buffer_list_queue: NetBufferList,
+    pub net_buffer_list: NetBufferList,
     remote_ip: [u8; 16],
     endpoint_handle: u64,
     remote_scope_id: SCOPE_ID,
@@ -112,7 +112,7 @@ impl Injector {
 
         TransportPacketList {
             ipv6,
-            net_buffer_list_queue: net_buffer_list,
+            net_buffer_list,
             remote_ip,
             endpoint_handle: callout_data.get_transport_endpoint_handle().unwrap_or(0),
             remote_scope_id: callout_data
@@ -153,7 +153,7 @@ impl Injector {
             };
             let address_family = if packet_list.ipv6 { AF_INET6 } else { AF_INET };
 
-            let net_buffer_list = packet_list.net_buffer_list_queue;
+            let net_buffer_list = packet_list.net_buffer_list;
             // Escape the stack. Packet buffer should be valid until the packet is injected.
             let boxed_nbl = Box::new(net_buffer_list);
             let raw_nbl = boxed_nbl.nbl;
@@ -338,6 +338,8 @@ unsafe extern "C" fn free_packet(
     if let Some(nbl) = net_buffer_list.as_ref() {
         if let Err(err) = check_ntstatus(nbl.Status) {
             crate::err!("inject status: {}", err);
+        } else {
+            crate::dbg!("inject status: Ok");
         }
     }
     _ = Box::from_raw(context as *mut NetBufferList);