From 1b0f2c251aa9c8a1fdf756842e3b287cdf10bcac Mon Sep 17 00:00:00 2001
From: Nodir Temirkhodjaev <nodir.temir@gmail.com>
Date: Sun, 19 Nov 2017 08:19:34 +0500
Subject: [PATCH] Driver: Collect logs in the pending buffer.

---
 src/common/fortlog.c |  2 ++
 src/driver/fortbuf.c | 62 ++++++++++++++++++++++++++++----------------
 src/driver/fortdrv.c | 27 +++++++++++--------
 3 files changed, 57 insertions(+), 34 deletions(-)

diff --git a/src/common/fortlog.c b/src/common/fortlog.c
index 1bd3f9c6..ec0a1f52 100644
--- a/src/common/fortlog.c
+++ b/src/common/fortlog.c
@@ -14,6 +14,8 @@
   ((FORT_LOG_BLOCKED_HEADER_SIZE + (path_len) \
     + (FORT_LOG_ALIGN - 1)) & ~(FORT_LOG_ALIGN - 1))
 
+#define FORT_LOG_BLOCKED_SIZE_MAX	FORT_LOG_BLOCKED_SIZE(FORT_LOG_PATH_MAX)
+
 #define fort_log_type(p)	(*((UINT32 *) (p)) & FORT_LOG_FLAG_TYPE_MASK)
 
 
diff --git a/src/driver/fortbuf.c b/src/driver/fortbuf.c
index 1537162f..c884f01e 100644
--- a/src/driver/fortbuf.c
+++ b/src/driver/fortbuf.c
@@ -41,7 +41,7 @@ fort_buffer_blocked_write (PFORT_BUFFER buf, UINT32 remote_ip, UINT32 pid,
                            UINT32 path_len, const PVOID path,
                            PIRP *irp, NTSTATUS *irp_status, ULONG_PTR *info)
 {
-  UINT32 len;
+  UINT32 len, new_top;
   PCHAR out;
   KIRQL irq;
   NTSTATUS status = STATUS_SUCCESS;
@@ -54,34 +54,36 @@ fort_buffer_blocked_write (PFORT_BUFFER buf, UINT32 remote_ip, UINT32 pid,
 
   KeAcquireSpinLock(&buf->lock, &irq);
 
+  new_top = buf->top + len;
+
   /* Try to directly write to pending client */
-  if (buf->irp) {
-    *irp = buf->irp;
-    buf->irp = NULL;
+  if (buf->irp != NULL) {
+    if (FORT_LOG_BLOCKED_SIZE_MAX > buf->out_len - new_top) {
+      *irp = buf->irp;
+      buf->irp = NULL;
 
-    *info = len;
-
-    if (buf->out_len < len) {
-      *irp_status = STATUS_BUFFER_TOO_SMALL;
-      /* fallback to buffer */
-    } else {
-      out = buf->out;
       *irp_status = STATUS_SUCCESS;
-      goto log_blocked;
+
+      *info = new_top;
+      new_top = 0;
     }
+
+    out = buf->out;
+  } else {
+    if (new_top > FORT_BUFFER_SIZE) {
+      status = STATUS_BUFFER_TOO_SMALL;
+      goto end;  /* drop on buffer overflow */
+    }
+
+    out = buf->data;
   }
 
-  if (len > FORT_BUFFER_SIZE - buf->top) {
-    status = STATUS_BUFFER_TOO_SMALL;
-    goto end;  /* drop on buffer overflow */
-  }
+  out += buf->top;
 
-  out = buf->data + buf->top;
-  buf->top += len;
-
- log_blocked:
   fort_log_blocked_write(out, remote_ip, pid, path_len, path);
 
+  buf->top = new_top;
+
  end:
   KeReleaseSpinLock(&buf->lock, irq);
 
@@ -100,8 +102,10 @@ fort_buffer_xmove (PFORT_BUFFER buf, PIRP irp, PVOID out, ULONG out_len,
   *info = buf->top;
 
   if (!buf->top) {
-    if (buf->irp) {
+    if (buf->irp != NULL) {
       status = STATUS_INSUFFICIENT_RESOURCES;
+    } else if (out_len < FORT_LOG_BLOCKED_SIZE_MAX) {
+      status = STATUS_BUFFER_TOO_SMALL;
     } else {
       buf->irp = irp;
       buf->out = out;
@@ -125,14 +129,26 @@ fort_buffer_xmove (PFORT_BUFFER buf, PIRP irp, PVOID out, ULONG out_len,
   return status;
 }
 
-static void
-fort_buffer_cancel_pending (PFORT_BUFFER buf, PIRP irp)
+static NTSTATUS
+fort_buffer_cancel_pending (PFORT_BUFFER buf, PIRP irp, ULONG_PTR *info)
 {
+  NTSTATUS status = STATUS_CANCELLED;
   KIRQL irq;
 
+  *info = 0;
+
   KeAcquireSpinLock(&buf->lock, &irq);
   if (irp == buf->irp) {
     buf->irp = NULL;
+
+    if (buf->top) {
+      *info = buf->top;
+      buf->top = 0;
+
+      status = STATUS_SUCCESS;
+    }
   }
   KeReleaseSpinLock(&buf->lock, irq);
+
+  return status;
 }
diff --git a/src/driver/fortdrv.c b/src/driver/fortdrv.c
index 98b41b68..cee872e5 100644
--- a/src/driver/fortdrv.c
+++ b/src/driver/fortdrv.c
@@ -212,20 +212,22 @@ fort_callout_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,
 
   if (conf_flags.log_blocked) {
     PIRP irp = NULL;
-    NTSTATUS irp_status;
     ULONG_PTR info;
+    NTSTATUS irp_status, status;
 
-    if (!NT_SUCCESS(fort_buffer_blocked_write(
+    status = fort_buffer_blocked_write(
         &g_device->buffer, remote_ip,
         (UINT32) inMetaValues->processId, path_len, path,
-        &irp, &irp_status, &info))) {
-      DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
-                 "FORT: Classify v4: Log buffer underflow: %d\n",
-                 FORT_LOG_BLOCKED_SIZE(path_len));
-    }
+        &irp, &irp_status, &info);
 
-    if (irp != NULL) {
-      fort_request_complete_info(irp, irp_status, info);
+    if (NT_SUCCESS(status)) {
+      if (irp != NULL) {
+        fort_request_complete_info(irp, irp_status, info);
+      }
+    } else {
+      DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
+                 "FORT: Classify v4: Log buffer error: %d (path len: %d)\n",
+                 status, path_len);
     }
   }
 
@@ -468,13 +470,16 @@ fort_device_cleanup (PDEVICE_OBJECT device, PIRP irp)
 static void
 fort_device_cancel_pending (PDEVICE_OBJECT device, PIRP irp)
 {
+  ULONG_PTR info;
+  NTSTATUS irp_status;
+
   UNUSED(device);
 
-  fort_buffer_cancel_pending(&g_device->buffer, irp);
+  irp_status = fort_buffer_cancel_pending(&g_device->buffer, irp, &info);
 
   IoReleaseCancelSpinLock(irp->CancelIrql);  /* before IoCompleteRequest()! */
 
-  fort_request_complete(irp, STATUS_CANCELLED);
+  fort_request_complete_info(irp, irp_status, info);
 }
 
 static NTSTATUS