Driver: Write statistic logs.

2024-11-15 09:17:28 +00:00 · 2017-11-25 12:29:24 +05:00 · 2017-11-25 12:29:24 +05:00 · 43018d29db
commit 43018d29db
parent d4d33a5c70
3 changed files with 112 additions and 52 deletions
--- a/src/driver/fortbuf.c
+++ b/src/driver/fortbuf.c
@ -106,7 +106,7 @@ fort_buffer_close (PFORT_BUFFER buf)

 static NTSTATUS
 fort_buffer_prepare (PFORT_BUFFER buf, UINT32 len, PCHAR *out,
-                     PIRP *irp, NTSTATUS *irp_status, ULONG_PTR *info)
+                     PIRP *irp, ULONG_PTR *info)
 {
  /* Check pending buffer */
  if (buf->out_len && buf->out_top < buf->out_len) {
@ -121,8 +121,6 @@ fort_buffer_prepare (PFORT_BUFFER buf, UINT32 len, PCHAR *out,
        *irp = buf->irp;
        buf->irp = NULL;

-        *irp_status = STATUS_SUCCESS;
-
        *info = new_top;
        new_top = 0;
      } else {
@ -152,7 +150,7 @@ fort_buffer_prepare (PFORT_BUFFER buf, UINT32 len, PCHAR *out,
 static NTSTATUS
 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)
+                           PIRP *irp, ULONG_PTR *info)
 {
  PCHAR out;
  UINT32 len;
@ -167,8 +165,7 @@ fort_buffer_blocked_write (PFORT_BUFFER buf, UINT32 remote_ip, UINT32 pid,

  KeAcquireInStackQueuedSpinLock(&buf->lock, &lock_queue);

-  status = fort_buffer_prepare(buf, len, &out,
-    irp, irp_status, info);
+  status = fort_buffer_prepare(buf, len, &out, irp, info);

  if (NT_SUCCESS(status)) {
    fort_log_blocked_write(out, remote_ip, pid, path_len, path);
@ -182,7 +179,7 @@ fort_buffer_blocked_write (PFORT_BUFFER buf, UINT32 remote_ip, UINT32 pid,
 static NTSTATUS
 fort_buffer_proc_new_write (PFORT_BUFFER buf, UINT32 pid,
                            UINT32 path_len, const PVOID path,
-                            PIRP *irp, NTSTATUS *irp_status, ULONG_PTR *info)
+                            PIRP *irp, ULONG_PTR *info)
 {
  PCHAR out;
  UINT32 len;
@ -197,8 +194,7 @@ fort_buffer_proc_new_write (PFORT_BUFFER buf, UINT32 pid,

  KeAcquireInStackQueuedSpinLock(&buf->lock, &lock_queue);

-  status = fort_buffer_prepare(buf, len, &out,
-    irp, irp_status, info);
+  status = fort_buffer_prepare(buf, len, &out, irp, info);

  if (NT_SUCCESS(status)) {
    fort_log_proc_new_write(out, pid, path_len, path);
@ -219,7 +215,7 @@ fort_buffer_proc_del_write (PFORT_BUFFER buf, UINT32 pid)

  KeAcquireInStackQueuedSpinLock(&buf->lock, &lock_queue);

-  status = fort_buffer_prepare(buf, len, &out, NULL, NULL, NULL);
+  status = fort_buffer_prepare(buf, len, &out, NULL, NULL);

  if (NT_SUCCESS(status)) {
    fort_log_proc_del_write(out, pid);
@ -298,14 +294,23 @@ fort_buffer_cancel_pending (PFORT_BUFFER buf, PIRP irp, ULONG_PTR *info)
  return status;
 }

-static BOOL
-fort_buffer_flush_pending (PFORT_BUFFER buf, PIRP *irp, ULONG_PTR *info)
+static void
+fort_buffer_dpc_begin (PFORT_BUFFER buf, PKLOCK_QUEUE_HANDLE lock_queue)
 {
-  KLOCK_QUEUE_HANDLE lock_queue;
-  UINT32 out_top;
+  KeAcquireInStackQueuedSpinLockAtDpcLevel(&buf->lock, lock_queue);
+}
+
+static void
+fort_buffer_dpc_end (PKLOCK_QUEUE_HANDLE lock_queue)
+{
+  KeReleaseInStackQueuedSpinLockFromDpcLevel(lock_queue);
+}
+
+static void
+fort_buffer_dpc_flush_pending (PFORT_BUFFER buf, PIRP *irp, ULONG_PTR *info)
+{
+  const UINT32 out_top = buf->out_top;

-  KeAcquireInStackQueuedSpinLockAtDpcLevel(&buf->lock, &lock_queue);
-  out_top = buf->out_top;
  if (out_top) {
    *info = out_top;

@ -315,7 +320,4 @@ fort_buffer_flush_pending (PFORT_BUFFER buf, PIRP *irp, ULONG_PTR *info)
    *irp = buf->irp;
    buf->irp = NULL;
  }
-  KeReleaseInStackQueuedSpinLockFromDpcLevel(&lock_queue);
-
-  return out_top != 0;
 }
--- a/src/driver/fortdrv.c
+++ b/src/driver/fortdrv.c
@ -179,7 +179,6 @@ fort_callout_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,

  PIRP irp = NULL;
  ULONG_PTR info;
-  NTSTATUS irp_status;

  UNUSED(layerData);
  UNUSED(flowContext);
@ -222,8 +221,7 @@ fort_callout_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,

      if (NT_SUCCESS(status)) {
        fort_buffer_proc_new_write(&g_device->buffer,
-          process_id, path_len, path,
-          &irp, &irp_status, &info);
+          process_id, path_len, path, &irp, &info);
      } else if (status != STATUS_OBJECT_NAME_EXISTS) {
        DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
                   "FORT: Classify v4: Flow assoc. error: %d\n", status);
@ -234,8 +232,7 @@ fort_callout_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,

  if (conf_flags.log_blocked) {
    fort_buffer_blocked_write(&g_device->buffer,
-      remote_ip, process_id, path_len, path,
-      &irp, &irp_status, &info);
+      remote_ip, process_id, path_len, path, &irp, &info);
  }

 block:
@ -251,7 +248,7 @@ fort_callout_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,

 end:
  if (irp != NULL) {
-    fort_request_complete_info(irp, irp_status, info);
+    fort_request_complete_info(irp, STATUS_SUCCESS, info);
  }
 }

@ -403,7 +400,7 @@ fort_callout_force_reauth (PDEVICE_OBJECT device,
 {
  NTSTATUS status;

-  // Check provider filters
+  /* Check provider filters */
  if (old_conf_flags.prov_boot != conf_flags.prov_boot) {
    fort_prov_unregister();

@ -416,7 +413,7 @@ fort_callout_force_reauth (PDEVICE_OBJECT device,
    goto end;
  }

-  // Check flow filter
+  /* Check flow filter */
  if (old_conf_flags.log_stat != conf_flags.log_stat) {
    if (old_conf_flags.log_stat) {
      fort_prov_flow_unregister();
@ -433,7 +430,7 @@ fort_callout_force_reauth (PDEVICE_OBJECT device,
    goto end;
  }

-  // Force reauth filter
+  /* Force reauth filter */
  status = fort_prov_reauth();
  if (!NT_SUCCESS(status)) {
    DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
@ -449,19 +446,46 @@ fort_callout_force_reauth (PDEVICE_OBJECT device,
 static void
 fort_callout_timer (void)
 {
-  // Flush traffic statistics
-  {
-  }
+  PFORT_BUFFER buf = &g_device->buffer;
+  PFORT_STAT stat = &g_device->stat;

-  // Flush pending buffer
-  {
-    PIRP irp;
+  KLOCK_QUEUE_HANDLE stat_lock_queue;
+  KLOCK_QUEUE_HANDLE buf_lock_queue;
+
+  PIRP irp = NULL;
  ULONG_PTR info;

-    if (fort_buffer_flush_pending(&g_device->buffer, &irp, &info)) {
-      fort_request_complete_info(irp, STATUS_SUCCESS, info);
+  /* Lock stat */
+  fort_stat_dpc_begin(stat, &stat_lock_queue);
+
+  /* Lock buffer */
+  fort_buffer_dpc_begin(buf, &buf_lock_queue);
+
+  /* Flush traffic statistics */
+  {
+    PCHAR out;
+    const UINT16 proc_count = stat->proc_count;
+    const UINT32 len = proc_count * sizeof(FORT_STAT_TRAF);
+
+    /* TODO: Write by chunks */
+    if (len < FORT_BUFFER_SIZE
+        && NT_SUCCESS(fort_buffer_prepare(buf, len, &out, &irp, &info))) {
+      fort_stat_dpc_traf_write(stat, 0, proc_count, out);
    }
  }
+
+  /* Unlock stat */
+  fort_stat_dpc_end(&stat_lock_queue);
+
+  /* Flush pending buffer */
+  fort_buffer_dpc_flush_pending(buf, &irp, &info);
+
+  /* Unlock buffer */
+  fort_buffer_dpc_end(&buf_lock_queue);
+
+  if (irp != NULL) {
+    fort_request_complete_info(irp, STATUS_SUCCESS, info);
+  }
 }

 static NTSTATUS
@ -529,15 +553,15 @@ static void
 fort_device_cancel_pending (PDEVICE_OBJECT device, PIRP irp)
 {
  ULONG_PTR info;
-  NTSTATUS irp_status;
+  NTSTATUS status;

  UNUSED(device);

-  irp_status = fort_buffer_cancel_pending(&g_device->buffer, irp, &info);
+  status = fort_buffer_cancel_pending(&g_device->buffer, irp, &info);

  IoReleaseCancelSpinLock(irp->CancelIrql);  /* before IoCompleteRequest()! */

-  fort_request_complete_info(irp, irp_status, info);
+  fort_request_complete_info(irp, status, info);
 }

 static NTSTATUS
@ -664,7 +688,7 @@ DriverEntry (PDRIVER_OBJECT driver, PUNICODE_STRING reg_path)

  UNUSED(reg_path);

-  // Wait for BFE to start
+  /* Wait for BFE to start */
  status = fort_bfe_wait();
  if (!NT_SUCCESS(status)) {
    DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
@ -701,17 +725,17 @@ DriverEntry (PDRIVER_OBJECT driver, PUNICODE_STRING reg_path)

      KeInitializeSpinLock(&g_device->conf_lock);

-      // Unegister old filters provider
+      /* Unegister old filters provider */
      {
        g_device->prov_boot = fort_prov_is_boot();

        fort_prov_unregister();
      }

-      // Install callouts
+      /* Install callouts */
      status = fort_callout_install(device);

-      // Register filters provider
+      /* Register filters provider */
      if (NT_SUCCESS(status)) {
        status = fort_prov_register(g_device->prov_boot);
      }
--- a/src/driver/fortstat.c
+++ b/src/driver/fortstat.c
@ -19,8 +19,9 @@ typedef struct fort_stat_proc {
 } FORT_STAT_PROC, *PFORT_STAT_PROC;

 typedef struct fort_stat {
-  UINT32 proc_top;
-  UINT32 proc_count;
+  UINT16 proc_top;
+  UINT16 proc_end;
+  UINT16 proc_count;

  int proc_free_index;

@ -70,6 +71,8 @@ fort_stat_proc_dec (PFORT_STAT stat, int proc_index)

    proc->refcount = stat->proc_free_index;
    stat->proc_free_index = proc_index;
+
+    stat->proc_count--;
  }

  return STATUS_SUCCESS;
@ -84,22 +87,22 @@ fort_stat_proc_del (PFORT_STAT_PROC procs)
 static BOOL
 fort_stat_proc_realloc (PFORT_STAT stat)
 {
-  const UINT32 count = stat->proc_count;
-  const UINT32 new_count = (count ? count : 16) * 3 / 2;
+  const UINT16 proc_end = stat->proc_end;
+  const UINT16 new_end = (proc_end ? proc_end : 16) * 3 / 2;
  PFORT_STAT_PROC new_procs = ExAllocatePoolWithTag(NonPagedPool,
-    new_count * sizeof(FORT_STAT_PROC), FORT_STAT_POOL_TAG);
+    new_end * sizeof(FORT_STAT_PROC), FORT_STAT_POOL_TAG);

  if (new_procs == NULL)
    return FALSE;

-  if (count) {
+  if (proc_end) {
    PFORT_STAT_PROC procs = stat->procs;

    RtlCopyMemory(new_procs, procs, stat->proc_top * sizeof(FORT_STAT_PROC));
    fort_stat_proc_del(procs);
  }

-  stat->proc_count = new_count;
+  stat->proc_end = new_end;
  stat->procs = new_procs;

  return TRUE;
@ -116,7 +119,7 @@ fort_stat_proc_add (PFORT_STAT stat, UINT32 process_id)

    stat->proc_free_index = proc->refcount;
  } else {
-    if (stat->proc_top >= stat->proc_count
+    if (stat->proc_top >= stat->proc_end
        && !fort_stat_proc_realloc(stat)) {
      return -1;
    }
@ -129,6 +132,8 @@ fort_stat_proc_add (PFORT_STAT stat, UINT32 process_id)
  proc->traf_all.QuadPart = 0;
  proc->refcount = 1;

+  stat->proc_count++;
+
  return proc_index;
 }

@ -229,5 +234,34 @@ fort_stat_flow_classify (PFORT_STAT stat, UINT64 flow_context,
  KeReleaseInStackQueuedSpinLock(&lock_queue);

  DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
-             "FORT: flow >: %d %d\n", (UINT32) flow_context, data_len);
+             "FORT: flow >: %d %d %d\n", (UINT32) flow_context, data_len, inbound);
+}
+
+static void
+fort_stat_dpc_begin (PFORT_STAT stat, PKLOCK_QUEUE_HANDLE lock_queue)
+{
+  KeAcquireInStackQueuedSpinLockAtDpcLevel(&stat->lock, lock_queue);
+}
+
+static void
+fort_stat_dpc_end (PKLOCK_QUEUE_HANDLE lock_queue)
+{
+  KeReleaseInStackQueuedSpinLockFromDpcLevel(lock_queue);
+}
+
+static void
+fort_stat_dpc_traf_write (PFORT_STAT stat, UINT32 index, UINT32 count,
+                          PCHAR out)
+{
+  PFORT_STAT_PROC proc = &stat->procs[index];
+  PFORT_STAT_TRAF out_traf = (PFORT_STAT_TRAF) out;
+
+  for (; index < count; ++proc) {
+    if (!proc->process_id)
+      continue;
+
+    *out_traf = proc->traf;
+
+    ++index;
+  }
 }