Driver: Improve speed limiting by deferring ACK packets.

src/driver/fortdrv.c

@@ -436,28 +436,24 @@ fort_callout_transport_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,
                                     FWPS_CLASSIFY_OUT0 *classifyOut,
                                     int ipProtoField, BOOL inbound)
 {
+  PNET_BUFFER netBuf;
   const IPPROTO ip_proto = (IPPROTO) inFixedValues->incomingValue[
     ipProtoField].value.uint8;
   const BOOL is_udp = (ip_proto == IPPROTO_UDP);
 
   if (!(is_udp || FWPS_IS_METADATA_FIELD_PRESENT(inMetaValues,
         FWPS_METADATA_FIELD_ALE_CLASSIFY_REQUIRED))
-      && netBufList != NULL) {
+      && netBufList != NULL
-
+      && (netBuf = NET_BUFFER_LIST_FIRST_NB(netBufList)) != NULL) {
     PFORT_STAT_FLOW flow = (PFORT_STAT_FLOW) flowContext;
 
     const BOOL ignore_tcp_rst = inbound && g_device->conf_flags.ignore_tcp_rst;
-    const BOOL defer_flow = flow->opt.speed_limit && !g_device->power_off
-      && (inbound ? flow->opt.defer_in : flow->opt.defer_out);
 
-    if (ignore_tcp_rst || defer_flow) {
+    const UINT32 defer_flag = inbound
-      const PNET_BUFFER netBuf = NET_BUFFER_LIST_FIRST_NB(netBufList);
+      ? FORT_STAT_FLOW_DEFER_IN : FORT_STAT_FLOW_DEFER_OUT;
-      TCP_HEADER buf;
+    const UINT32 flow_flags = FORT_STAT_FLOW_SPEED_LIMIT | defer_flag;
-      PTCP_HEADER tcpHeader;
+    const BOOL defer_flow = (flow->opt.flags & flow_flags) == flow_flags
-      UINT32 tcpFlags;
+      && !g_device->power_off;
-
-      if (netBuf == NULL)
-        goto permit;
 
     /* Position in the packet data:
      * FWPS_LAYER_INBOUND_TRANSPORT_V4: The beginning of the data.
@@ -465,6 +461,11 @@ fort_callout_transport_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,
      */
     const UINT32 headerOffset = inbound ? 0 : sizeof(TCP_HEADER);
 
+    if (ignore_tcp_rst) {
+      TCP_HEADER buf;
+      PTCP_HEADER tcpHeader;
+      UINT32 tcpFlags;
+
       if (headerOffset != 0) {
         NdisRetreatNetBufferDataStart(netBuf, headerOffset, 0, NULL);
       }
@@ -479,22 +480,20 @@ fort_callout_transport_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,
       if (tcpHeader == NULL)
         goto permit;
 
-      if (ignore_tcp_rst && (tcpFlags & TCP_FLAG_RST))
+      if (tcpFlags & TCP_FLAG_RST)
         goto block;
-
-      if (defer_flow //&& (tcpFlags & TCP_FLAG_ACK)
-          && NET_BUFFER_DATA_LENGTH(netBuf) == headerOffset) {
-        NTSTATUS status;
-
-        status = fort_defer_add(&g_device->defer, inFixedValues, inMetaValues,
-          netBufList, inbound);
-
-        if (!NT_SUCCESS(status)) {
-          DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
-                     "FORT: Transport Classify: Defer error: %x\n", status);
     }
 
+    if (defer_flow && NET_BUFFER_DATA_LENGTH(netBuf) == headerOffset) {
+      const NTSTATUS status = fort_defer_add(&g_device->defer,
+        inFixedValues, inMetaValues, netBufList, inbound);
+
+      if (NT_SUCCESS(status))
         goto block;
+
+      if (status == STATUS_CANT_TERMINATE_SELF) {
+        /* Clear ACK deferring */
+        flow->opt.flags &= ~defer_flag;
       }
     }
   }

src/driver/fortpkt.c

@@ -120,6 +120,16 @@ fort_defer_add (PFORT_DEFER defer,
   if (defer->injection4_id == INVALID_HANDLE_VALUE)
     return STATUS_FWP_TCPIP_NOT_READY;
 
+  /* Skip self injected packet */
+  {
+    const FWPS_PACKET_INJECTION_STATE state = FwpsQueryPacketInjectionState0(
+      defer->injection4_id, netBufList, NULL);
+
+    if (state == FWPS_PACKET_INJECTED_BY_SELF
+        || state == FWPS_PACKET_PREVIOUSLY_INJECTED_BY_SELF)
+      return STATUS_CANT_TERMINATE_SELF;
+  }
+
   /* Skip IpSec protected packet */
   if (inbound) {
     FWPS_PACKET_LIST_INFORMATION info;
@@ -156,6 +166,7 @@ fort_defer_add (PFORT_DEFER defer,
     defer->packet_head = defer->packet_tail = pkt;
   } else {
     defer->packet_tail->next = pkt;
+    defer->packet_tail = pkt;
   }
 
   pkt->inbound = inbound;

src/driver/fortstat.c

@@ -28,10 +28,12 @@ typedef struct fort_stat_proc {
   };
 } FORT_STAT_PROC, *PFORT_STAT_PROC;
 
+#define FORT_STAT_FLOW_SPEED_LIMIT	0x01
+#define FORT_STAT_FLOW_DEFER_IN		0x02
+#define FORT_STAT_FLOW_DEFER_OUT	0x04
+
 typedef struct fort_stat_flow_opt {
-  UCHAR speed_limit	: 1;
+  UCHAR volatile flags;
-  UCHAR defer_in	: 1;
-  UCHAR defer_out	: 1;
   UCHAR group_index;
   UINT16 proc_index;
 } FORT_STAT_FLOW_OPT, *PFORT_STAT_FLOW_OPT;
@@ -45,7 +47,7 @@ typedef struct fort_stat_flow {
 #if defined(_WIN64)
     UINT64 flow_id;
 #else
-    FORT_STAT_FLOW_OPT volatile opt;
+    FORT_STAT_FLOW_OPT opt;
 #endif
     void *data;
   };
@@ -346,8 +348,7 @@ fort_stat_flow_add (PFORT_STAT stat, UINT64 flow_id,
     fort_stat_proc_inc(stat, proc_index);
   }
 
-  flow->opt.speed_limit = (UCHAR) speed_limit;
+  flow->opt.flags = (speed_limit ? FORT_STAT_FLOW_SPEED_LIMIT : 0);
-  flow->opt.defer_in = flow->opt.defer_out = FALSE;
   flow->opt.group_index = group_index;
   flow->opt.proc_index = proc_index;
 
@@ -518,7 +519,7 @@ fort_stat_flow_classify (PFORT_STAT stat, UINT64 flowContext,
     /* Add traffic to process */
     *proc_bytes += data_len;
 
-    if (flow->opt.speed_limit) {
+    if (flow->opt.flags & FORT_STAT_FLOW_SPEED_LIMIT) {
       const UCHAR group_index = flow->opt.group_index;
       const PFORT_CONF_LIMIT group_limit = &stat->limits[group_index];
       const UINT32 limit_bytes = inbound ? group_limit->in_bytes
@@ -528,7 +529,7 @@ fort_stat_flow_classify (PFORT_STAT stat, UINT64 flowContext,
         PFORT_STAT_GROUP group = &stat->groups[group_index];
         UINT32 *group_bytes = inbound ? &group->traf.in_bytes
           : &group->traf.out_bytes;
-        UCHAR defer_flow = TRUE;
+        BOOL defer_flow = TRUE;
 
         if (*group_bytes < limit_bytes) {
           /* Add traffic to app. group */
@@ -540,10 +541,14 @@ fort_stat_flow_classify (PFORT_STAT stat, UINT64 flowContext,
         }
 
         /* Defer ACK */
-        if (inbound) {
+        {
-          flow->opt.defer_out = defer_flow;
+          const UINT32 defer_flag = inbound
-        } else {
+            ? FORT_STAT_FLOW_DEFER_OUT : FORT_STAT_FLOW_DEFER_IN;
-          flow->opt.defer_in = defer_flow;
+
+          if (defer_flow)
+            flow->opt.flags |= defer_flag;
+          else
+            flow->opt.flags &= ~defer_flag;
         }
       }
     }