Implemented WorkQueue's as they are pretty common in WDK.

 * added calloc() symbol

Signed-off-by: Toni Uhlig <matzeton@googlemail.com>

4 files changed, 175 insertions, 3 deletions

diff --git a/CRT/DriverThread.cpp b/CRT/DriverThread.cpp
index efc4023..e2bb2ec 100644
--- a/CRT/DriverThread.cpp
+++ b/CRT/DriverThread.cpp
@@ -25,6 +25,11 @@ NTSTATUS DriverThread::Thread::Start(threadRoutine_t routine, PVOID threadContex
     NTSTATUS status;
 
     LockGuard lock(m_mutex);
+    if (m_threadObject != nullptr)
+    {
+        return STATUS_UNSUCCESSFUL;
+    }
+
     m_routine = routine;
     m_threadContext = threadContext;
     status = PsCreateSystemThread(&threadHandle, (ACCESS_MASK)0, NULL, (HANDLE)0, NULL, InterceptorThreadRoutine, this);
@@ -34,7 +39,14 @@ NTSTATUS DriverThread::Thread::Start(threadRoutine_t routine, PVOID threadContex
         return status;
     }
 
-    ObReferenceObjectByHandle(threadHandle, THREAD_ALL_ACCESS, NULL, KernelMode, (PVOID *)&m_threadObject, NULL);
+    status =
+        ObReferenceObjectByHandle(threadHandle, THREAD_ALL_ACCESS, NULL, KernelMode, (PVOID *)&m_threadObject, NULL);
+
+    if (!NT_SUCCESS(status))
+    {
+        return status;
+    }
+
     return ZwClose(threadHandle);
 }
 
@@ -44,6 +56,7 @@ NTSTATUS DriverThread::Thread::WaitForTermination(LONGLONG timeout)
     {
         return STATUS_UNSUCCESSFUL;
     }
+
     LockGuard lock(m_mutex);
     if (m_threadObject == nullptr)
     {
@@ -135,3 +148,107 @@ DriverThread::LockGuard::~LockGuard(void)
 {
     m_Lock.Unlock();
 }
+
+// WorkQueue
+
+DriverThread::WorkQueue::WorkQueue(void) : m_worker()
+{
+    InitializeSListHead(&m_work);
+    KeInitializeEvent(&m_wakeEvent, SynchronizationEvent, FALSE);
+    m_stopWorker = FALSE;
+}
+
+DriverThread::WorkQueue::~WorkQueue(void)
+{
+    Stop();
+}
+
+NTSTATUS DriverThread::WorkQueue::Start(workerRoutine_t workerRoutine)
+{
+    NTSTATUS status;
+
+    {
+        LockGuard lock(m_mutex);
+        m_workerRoutine = workerRoutine;
+        status = m_worker.Start(WorkerInterceptorRoutine, this);
+    }
+
+    if (!NT_SUCCESS(status) && status != STATUS_UNSUCCESSFUL)
+    {
+        Stop();
+    }
+
+    return status;
+}
+
+void DriverThread::WorkQueue::Stop(void)
+{
+    LockGuard lock(m_mutex);
+    if (m_stopWorker == TRUE)
+    {
+        return;
+    }
+    m_stopWorker = TRUE;
+    KeSetEvent(&m_wakeEvent, 0, FALSE);
+}
+
+void DriverThread::WorkQueue::Enqueue(PSLIST_ENTRY workItem)
+{
+    if (InterlockedPushEntrySList(&m_work, workItem) == NULL)
+    {
+        // Work queue was empty. So, signal the work queue event in case the
+        // worker thread is waiting on the event for more operations.
+        KeSetEvent(&m_wakeEvent, 0, FALSE);
+    }
+}
+
+NTSTATUS DriverThread::WorkQueue::WorkerInterceptorRoutine(PVOID workerContext)
+{
+    DriverThread::WorkQueue * wq = (DriverThread::WorkQueue *)workerContext;
+    PSLIST_ENTRY listEntryRev, listEntry, next;
+
+    PAGED_CODE();
+
+    for (;;)
+    {
+        // Flush all the queued operations into a local list
+        listEntryRev = InterlockedFlushSList(&wq->m_work);
+
+        if (listEntryRev == NULL)
+        {
+
+            // There's no work to do. If we are allowed to stop, then stop.
+            if (wq->m_stopWorker == TRUE)
+            {
+                break;
+            }
+
+            // Otherwise, wait for more operations to be enqueued.
+            KeWaitForSingleObject(&wq->m_wakeEvent, Executive, KernelMode, FALSE, 0);
+            continue;
+        }
+
+        // Need to reverse the flushed list in order to preserve the FIFO order
+        listEntry = NULL;
+        while (listEntryRev != NULL)
+        {
+            next = listEntryRev->Next;
+            listEntryRev->Next = listEntry;
+            listEntry = listEntryRev;
+            listEntryRev = next;
+        }
+
+        // Now process the correctly ordered list of operations one by one
+        while (listEntry)
+        {
+            PSLIST_ENTRY arg = listEntry;
+            listEntry = listEntry->Next;
+            if (wq->m_workerRoutine(arg) != STATUS_SUCCESS)
+            {
+                wq->m_stopWorker = TRUE;
+            }
+        }
+    }
+
+    return STATUS_SUCCESS;
+}
diff --git a/CRT/DriverThread.hpp b/CRT/DriverThread.hpp
index d00db1b..6ae9d2c 100644
--- a/CRT/DriverThread.hpp
+++ b/CRT/DriverThread.hpp
@@ -6,6 +6,7 @@
 extern "C" void InterceptorThreadRoutine(PVOID threadContext);
 
 typedef NTSTATUS (*threadRoutine_t)(PVOID);
+typedef NTSTATUS (*workerRoutine_t)(PSLIST_ENTRY);
 
 namespace DriverThread
 {
@@ -78,6 +79,26 @@ private:
     KSEMAPHORE m_semaphore;
 };
 
+class WorkQueue
+{
+public:
+    WorkQueue(void);
+    ~WorkQueue(void);
+    NTSTATUS Start(workerRoutine_t workerRoutine);
+    void Stop(void);
+    void Enqueue(PSLIST_ENTRY workItem);
+
+private:
+    Mutex m_mutex;
+    SLIST_HEADER m_work;
+    KEVENT m_wakeEvent;
+    BOOLEAN m_stopWorker; // Work LIST must be empty and StopWorker TRUE to be able to stop!
+    Thread m_worker;
+    workerRoutine_t m_workerRoutine;
+
+    static NTSTATUS WorkerInterceptorRoutine(PVOID workerContext);
+};
+
 }; // namespace DriverThread
 
 #endif
diff --git a/CRT/kcrt.c b/CRT/kcrt.c
index fff767d..84b7f25 100644
--- a/CRT/kcrt.c
+++ b/CRT/kcrt.c
@@ -112,6 +112,11 @@ void * __cdecl malloc(size_t size)
     return NULL;
 }
 
+void * __cdecl calloc(size_t nmemb, size_t size)
+{
+    return malloc(nmemb * size);
+}
+
 void * __cdecl realloc(void * ptr, size_t new_size)
 {
     if (!ptr)
@@ -181,7 +186,8 @@ void __cdecl __cxa_pure_virtual(void)
 
 float __cdecl ceilf(float x)
 {
-    union {
+    union
+    {
         float f;
         UINT32 i;
     } u = {x};
@@ -215,7 +221,7 @@ float __cdecl ceilf(float x)
 
 static void __cdecl __ctors(void)
 {
-    unsigned long long int const * const * const l = (unsigned long long int const * const * const)&__CTOR_LIST__;
+    unsigned long long int const * const * const l = (unsigned long long int const * const * const) & __CTOR_LIST__;
     unsigned long long int i = (unsigned long long int)*l;
     init_and_deinit_fn const * p;
 
diff --git a/examples/dpp-example-cplusplus.cpp b/examples/dpp-example-cplusplus.cpp
index c1d9b29..0d04d37 100644
--- a/examples/dpp-example-cplusplus.cpp
+++ b/examples/dpp-example-cplusplus.cpp
@@ -56,6 +56,13 @@ private:
     unsigned int some_value = 0;
 };
 
+struct WorkItem
+{
+    SLIST_ENTRY QueueEntry;
+
+    UINT32 counter;
+};
+static DriverThread::WorkQueue work_queue;
 static DerivedWithCDtor some_static(0xDEADC0DE);
 
 struct threadContext
@@ -79,6 +86,21 @@ static NTSTATUS threadRoutine(PVOID threadContext)
     return STATUS_SUCCESS;
 }
 
+static NTSTATUS test_worker(PSLIST_ENTRY workItem)
+{
+    struct WorkItem * wi = CONTAINING_RECORD(workItem, struct WorkItem, QueueEntry);
+
+    while (wi->counter-- > 0)
+    {
+        DbgPrint("WorkItem Counter: %u\n", wi->counter);
+    }
+    DbgPrint("Worker finished.\n");
+
+    free(wi);
+
+    return STATUS_SUCCESS;
+}
+
 static void test_cplusplus(void)
 {
     TestSmth t;
@@ -94,6 +116,12 @@ static void test_cplusplus(void)
     ctx.dth.WaitForTermination();
     DbgPrint("MainThread EOF\n");
 
+    struct WorkItem * wi = (struct WorkItem *)calloc(1, sizeof(*wi));
+    wi->counter = 3;
+    work_queue.Enqueue(&wi->QueueEntry);
+    work_queue.Start(test_worker);
+    work_queue.Stop();
+
     some_static.doSmth();
 }