Driver: Add Heartbeat log for testing.

src/common/fortlog.c

@@ -7,6 +7,7 @@
 #define FORT_LOG_FLAG_BLOCKED		0x01000000
 #define FORT_LOG_FLAG_PROC_NEW		0x02000000
 #define FORT_LOG_FLAG_STAT_TRAF		0x04000000
+#define FORT_LOG_FLAG_HEARTBEAT		0x08000000
 #define FORT_LOG_FLAG_TYPE_MASK		0xFF000000
 
 #define FORT_LOG_BLOCKED_HEADER_SIZE	(4 * sizeof(UINT32))
@@ -34,6 +35,8 @@
 #define FORT_LOG_STAT_BUFFER_PROC_COUNT \
   ((FORT_BUFFER_SIZE - FORT_LOG_STAT_HEADER_SIZE) / FORT_LOG_STAT_TRAF_SIZE(1))
 
+#define FORT_LOG_HEARTBEAT_SIZE		sizeof(UINT32)
+
 #define FORT_LOG_SIZE_MAX	FORT_LOG_BLOCKED_SIZE_MAX
 
 #define fort_log_type(p)	(*((UINT32 *) (p)) & FORT_LOG_FLAG_TYPE_MASK)
@@ -125,3 +128,19 @@ fort_log_stat_traf_header_read (const char *p, UINT16 *proc_count)
 
   *proc_count = (UINT16) *up;
 }
+
+static void
+fort_log_heartbeat_write (char *p, UINT16 tick)
+{
+  UINT32 *up = (UINT32 *) p;
+
+  *up = FORT_LOG_FLAG_HEARTBEAT | tick;
+}
+
+static void
+fort_log_heartbeat_read (const char *p, UINT16 *tick)
+{
+  const UINT32 *up = (const UINT32 *) p;
+
+  *tick = (UINT16) *up;
+}

src/common/version.h

@@ -14,6 +14,6 @@
 #define APP_UPDATES_URL		"https://github.com/tnodir/fort/releases"
 #define APP_UPDATES_API_URL	"https://api.github.com/repos/tnodir/fort/releases/latest"
 
-#define DRIVER_VERSION		15
+#define DRIVER_VERSION		16
 
 #endif // VERSION_H

src/driver/fortdrv.c

@@ -41,6 +41,10 @@ typedef struct fort_conf_ref {
 typedef struct fort_device {
   UCHAR volatile flags;
 
+#ifdef LOG_HEARTBEAT
+  UINT16 volatile heartbeat_tick;
+#endif
+
   UINT32 connect4_id;
   UINT32 accept4_id;
 
@@ -59,6 +63,9 @@ typedef struct fort_device {
   FORT_DEFER defer;
   FORT_TIMER log_timer;
   FORT_TIMER app_timer;
+#ifdef LOG_HEARTBEAT
+  FORT_TIMER heartbeat_timer;
+#endif
   FORT_WORKER worker;
 } FORT_DEVICE, *PFORT_DEVICE;
 
@@ -924,7 +931,7 @@ fort_callout_force_reauth (const FORT_CONF_FLAGS old_conf_flags,
       fort_device_flag_set(FORT_DEVICE_FILTER_TRANSPORT, filter_transport);
 
       if (old_conf_flags.log_stat) {
-        fort_prov_flow_unregister(engine);
+      fort_prov_flow_unregister(engine);
       }
 
  stat_prov:
@@ -1047,6 +1054,39 @@ fort_app_period_timer (void)
   }
 }
 
+#ifdef LOG_HEARTBEAT
+static void
+fort_heartbeat_timer (void)
+{
+  PFORT_BUFFER buf = &g_device->buffer;
+  KLOCK_QUEUE_HANDLE buf_lock_queue;
+
+  /* Lock buffer */
+  KeAcquireInStackQueuedSpinLock(&buf->lock, &buf_lock_queue);
+
+  /* Log heartbeat */
+  {
+    const UINT16 tick = InterlockedIncrement16(&g_device->heartbeat_tick);
+    const UINT32 len = FORT_LOG_HEARTBEAT_SIZE;
+    PCHAR out;
+    NTSTATUS status;
+
+    status = fort_buffer_prepare(buf, len, &out, NULL, NULL);
+    if (!NT_SUCCESS(status)) {
+      DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
+                 "FORT: Heartbeat Timer: Error: %x\n", status);
+      goto end;
+    }
+
+    fort_log_heartbeat_write(out, tick);
+  }
+
+ end:
+  /* Unlock buffer */
+  KeReleaseInStackQueuedSpinLock(&buf_lock_queue);
+}
+#endif
+
 static NTSTATUS
 fort_device_create (PDEVICE_OBJECT device, PIRP irp)
 {
@@ -1060,6 +1100,17 @@ fort_device_create (PDEVICE_OBJECT device, PIRP irp)
     status = STATUS_SHARING_VIOLATION;  /* Only one client may connect */
   }
 
+  if (NT_SUCCESS(status)) {
+    /* Clear buffer */
+    fort_buffer_clear(&g_device->buffer);
+
+#ifdef LOG_HEARTBEAT
+    InterlockedAnd16(&g_device->heartbeat_tick, 0);
+    fort_heartbeat_timer();
+    fort_timer_update(&g_device->heartbeat_timer, TRUE);
+#endif
+  }
+
   fort_request_complete(irp, status);
 
   return status;
@@ -1088,6 +1139,11 @@ fort_device_cleanup (PDEVICE_OBJECT device, PIRP irp)
     fort_callout_force_reauth(old_conf_flags, FORT_DEFER_FLUSH_ALL);
   }
 
+#ifdef LOG_HEARTBEAT
+  fort_timer_update(&g_device->heartbeat_timer, FALSE);
+  fort_heartbeat_timer();
+#endif
+
   /* Clear buffer */
   fort_buffer_clear(&g_device->buffer);
 
@@ -1315,6 +1371,9 @@ fort_driver_unload (PDRIVER_OBJECT driver)
   if (g_device != NULL) {
     fort_callout_defer_flush();
 
+#ifdef LOG_HEARTBEAT
+    fort_timer_close(&g_device->heartbeat_timer);
+#endif
     fort_timer_close(&g_device->app_timer);
     fort_timer_close(&g_device->log_timer);
     fort_defer_close(&g_device->defer);
@@ -1404,6 +1463,9 @@ DriverEntry (PDRIVER_OBJECT driver, PUNICODE_STRING reg_path)
       fort_defer_open(&g_device->defer);
       fort_timer_open(&g_device->log_timer, 500, FALSE, &fort_callout_timer);
       fort_timer_open(&g_device->app_timer, 60000, TRUE, &fort_app_period_timer);
+#ifdef LOG_HEARTBEAT
+      fort_timer_open(&g_device->heartbeat_timer, 1000, TRUE, &fort_heartbeat_timer);
+#endif
 
       KeInitializeSpinLock(&g_device->conf_lock);
 

src/driver/fortdrv.h

@@ -5,6 +5,7 @@
 #define DOS_DEVICE_NAME		L"\\DosDevices\\fortfw"
 
 #define FORT_DRIVER
+//#define LOG_HEARTBEAT
 
 #define fort_request_complete_info(irp, status, info) \
   do { \

src/ui/FortFirewall.pro

@@ -30,6 +30,7 @@ SOURCES += \
     log/logbuffer.cpp \
     log/logentry.cpp \
     log/logentryblocked.cpp \
+    log/logentryheartbeat.cpp \
     log/logentryprocnew.cpp \
     log/logentrystattraf.cpp \
     log/logmanager.cpp \
@@ -100,6 +101,7 @@ HEADERS += \
     log/logbuffer.h \
     log/logentry.h \
     log/logentryblocked.h \
+    log/logentryheartbeat.h \
     log/logentryprocnew.h \
     log/logentrystattraf.h \
     log/logmanager.h \

src/ui/db/databasemanager.cpp

@@ -344,8 +344,9 @@ void DatabaseManager::logStatTraf(quint16 procCount, const quint32 *procTrafByte
         }
 
         const int procIndex = m_appIndexes.value(pid, INVALID_APP_INDEX);
-        if (procIndex == INVALID_APP_INDEX) {
-            qCritical(CLOG_DATABASE_MANAGER()) << "UI & Driver's states mismatch.";
+        if (Q_UNLIKELY(procIndex == INVALID_APP_INDEX)) {
+            qCritical(CLOG_DATABASE_MANAGER()) << "UI & Driver's states mismatch! Expected process count:"
+                                               << procCount << "Got:" << m_appIndexes.size();
             abort();
         }
 

src/ui/fortcommon.cpp

@@ -88,6 +88,11 @@ quint32 FortCommon::logStatSize(quint16 procCount)
     return FORT_LOG_STAT_SIZE(procCount);
 }
 
+quint32 FortCommon::logHeartbeatSize()
+{
+    return FORT_LOG_HEARTBEAT_SIZE;
+}
+
 quint32 FortCommon::logType(const char *input)
 {
     return fort_log_type(input);
@@ -127,6 +132,11 @@ void FortCommon::logStatTrafHeaderRead(const char *input,
     fort_log_stat_traf_header_read(input, procCount);
 }
 
+void FortCommon::logHeartbeatRead(const char *input, quint16 *tick)
+{
+    fort_log_heartbeat_read(input, tick);
+}
+
 void FortCommon::confAppPermsMaskInit(void *drvConf)
 {
     PFORT_CONF conf = (PFORT_CONF) drvConf;

src/ui/fortcommon.h

@@ -32,6 +32,8 @@ public:
     static quint32 logStatTrafSize(quint16 procCount);
     static quint32 logStatSize(quint16 procCount);
 
+    static quint32 logHeartbeatSize();
+
     static quint32 logType(const char *input);
 
     static void logBlockedHeaderWrite(char *output,
@@ -49,6 +51,8 @@ public:
     static void logStatTrafHeaderRead(const char *input,
                                       quint16 *procCount);
 
+    static void logHeartbeatRead(const char *input, quint16 *tick);
+
     static void confAppPermsMaskInit(void *drvConf);
     static bool confIpInRange(const void *drvConf, quint32 ip,
                               bool included = false, int addrGroupIndex = 0);

src/ui/fortmanager.cpp

@@ -176,6 +176,7 @@ bool FortManager::setupDriver()
 void FortManager::closeDriver()
 {
     updateLogManager(false);
+    updateDatabaseManager(nullptr);
 
     m_driverManager->closeDevice();
 }
@@ -533,10 +534,11 @@ bool FortManager::updateDriverConf(FirewallConf *conf, bool onlyFlags)
 
     if (res) {
         updateDatabaseManager(conf);
+        updateLogManager(true);
+    } else {
+        closeDriver();
     }
 
-    updateLogManager(true);
-
     return res;
 }
 

src/ui/log/logbuffer.cpp

@@ -2,6 +2,7 @@
 
 #include "fortcommon.h"
 #include "logentryblocked.h"
+#include "logentryheartbeat.h"
 #include "logentryprocnew.h"
 #include "logentrystattraf.h"
 
@@ -163,3 +164,18 @@ void LogBuffer::readEntryStatTraf(LogEntryStatTraf *logEntry)
     const int entrySize = int(FortCommon::logStatTrafSize(procCount));
     m_offset += entrySize;
 }
+
+void LogBuffer::readEntryHeartbeat(LogEntryHeartbeat *logEntry)
+{
+    Q_ASSERT(m_offset < m_top);
+
+    const char *input = this->input();
+
+    quint16 tick;
+    FortCommon::logHeartbeatRead(input, &tick);
+
+    logEntry->setTick(tick);
+
+    const int entrySize = int(FortCommon::logHeartbeatSize());
+    m_offset += entrySize;
+}

src/ui/log/logbuffer.h

@@ -7,6 +7,7 @@
 #include "logentry.h"
 
 QT_FORWARD_DECLARE_CLASS(LogEntryBlocked)
+QT_FORWARD_DECLARE_CLASS(LogEntryHeartbeat)
 QT_FORWARD_DECLARE_CLASS(LogEntryProcNew)
 QT_FORWARD_DECLARE_CLASS(LogEntryStatTraf)
 
@@ -33,6 +34,8 @@ public:
 
     void readEntryStatTraf(LogEntryStatTraf *logEntry);
 
+    void readEntryHeartbeat(LogEntryHeartbeat *logEntry);
+
 signals:
 
 public slots:

src/ui/log/logentry.h

@@ -13,7 +13,8 @@ public:
         // synchronize with FORT_LOG_FLAG_*
         AppBlocked  = 0x01000000,
         ProcNew     = 0x02000000,
-        StatTraf    = 0x04000000
+        StatTraf    = 0x04000000,
+        Heartbeat   = 0x08000000
     };
 
     explicit LogEntry() = default;

src/ui/log/logentryheartbeat.cpp

@@ -0,0 +1,11 @@
+#include "logentryheartbeat.h"
+
+LogEntryHeartbeat::LogEntryHeartbeat(quint16 tick) :
+    m_tick(tick)
+{
+}
+
+void LogEntryHeartbeat::setTick(quint16 tick)
+{
+    m_tick = tick;
+}

src/ui/log/logentryheartbeat.h

@@ -0,0 +1,20 @@
+#ifndef LOGENTRYHEARBEAT_H
+#define LOGENTRYHEARBEAT_H
+
+#include "logentry.h"
+
+class LogEntryHeartbeat : public LogEntry
+{
+public:
+    explicit LogEntryHeartbeat(quint16 tick = 0);
+
+    LogEntry::LogType type() const override { return Heartbeat; }
+
+    quint16 tick() const { return m_tick; }
+    void setTick(quint16 tick);
+
+private:
+    quint16 m_tick;
+};
+
+#endif // LOGENTRYHEARBEAT_H

src/ui/log/logmanager.cpp

@@ -1,11 +1,13 @@
 #include "logmanager.h"
 
 #include <QCoreApplication>
+#include <QDebug>
 
 #include "../driver/driverworker.h"
 #include "../fortcommon.h"
 #include "logbuffer.h"
 #include "logentryblocked.h"
+#include "logentryheartbeat.h"
 #include "logentryprocnew.h"
 #include "logentrystattraf.h"
 #include "model/appblockedmodel.h"
@@ -20,7 +22,6 @@ LogManager::LogManager(DatabaseManager *databaseManager,
     m_appBlockedModel(new AppBlockedModel(this)),
     m_appStatModel(new AppStatModel(databaseManager, this))
 {
-    setupDriverWorker();
 }
 
 void LogManager::setActive(bool active)
@@ -49,6 +50,9 @@ void LogManager::setErrorMessage(const QString &errorMessage)
 void LogManager::initialize()
 {
     m_appStatModel->initialize();
+
+    connect(m_driverWorker, &DriverWorker::readLogResult,
+            this, &LogManager::processLogBuffer, Qt::QueuedConnection);
 }
 
 void LogManager::close()
@@ -58,12 +62,6 @@ void LogManager::close()
     disconnect(m_driverWorker);
 }
 
-void LogManager::setupDriverWorker()
-{
-    connect(m_driverWorker, &DriverWorker::readLogResult,
-            this, &LogManager::processLogBuffer, Qt::QueuedConnection);
-}
-
 void LogManager::readLogAsync()
 {
     LogBuffer *logBuffer = getFreeBuffer();
@@ -110,27 +108,36 @@ void LogManager::processLogBuffer(LogBuffer *logBuffer, bool success,
 
 void LogManager::readLogEntries(LogBuffer *logBuffer)
 {
-    LogEntryBlocked blockedEntry;
-    LogEntryProcNew procNewEntry;
-    LogEntryStatTraf statTrafEntry;
-
     forever {
         switch (logBuffer->peekEntryType()) {
         case LogEntry::AppBlocked: {
+            LogEntryBlocked blockedEntry;
             logBuffer->readEntryBlocked(&blockedEntry);
             m_appBlockedModel->addLogEntry(blockedEntry);
             break;
         }
         case LogEntry::ProcNew: {
+            LogEntryProcNew procNewEntry;
             logBuffer->readEntryProcNew(&procNewEntry);
             m_appStatModel->handleProcNew(procNewEntry);
             break;
         }
         case LogEntry::StatTraf: {
+            LogEntryStatTraf statTrafEntry;
             logBuffer->readEntryStatTraf(&statTrafEntry);
             m_appStatModel->handleStatTraf(statTrafEntry);
             break;
         }
+        case LogEntry::Heartbeat: {
+            LogEntryHeartbeat heartbeatEntry;
+            logBuffer->readEntryHeartbeat(&heartbeatEntry);
+            if (++m_heartbeatTick != heartbeatEntry.tick()) {
+                qCritical() << "Heartbeat ticks mismatch! Expected:"
+                            << heartbeatEntry.tick() << "Got:" << m_heartbeatTick;
+                abort();
+            }
+            break;
+        }
         default:
             Q_ASSERT(logBuffer->offset() != 0);
             return;

src/ui/log/logmanager.h

@@ -45,8 +45,6 @@ private slots:
 private:
     void setErrorMessage(const QString &errorMessage);
 
-    void setupDriverWorker();
-
     void readLogAsync();
     void cancelAsyncIo();
 
@@ -58,6 +56,8 @@ private:
 private:
     bool m_active;
 
+    quint16 m_heartbeatTick;
+
     DriverWorker *m_driverWorker;
     QList<LogBuffer *> m_freeBuffers;