8damon-Blackbird-Platform/UserMode/hook/hooks/controller.h

#pragma once

#include <cstdint>
#include <vector>
#include <mutex>

#define WIN32_LEAN_AND_MEAN
#include <Windows.h>

#include "ws.h"
#include "nt.h"
#include "ki.h"
#include "module.h"

#include "../instrument/stacktrace.h"

struct WinsockCapturedEvent
{
    DWORD ThreadId;
    SOCKET Socket;
    WinsockOperation Operation;
    void *Caller;
    std::uint64_t Args[4];
    std::vector<std::uint8_t> Data;

    IC_STACKTRACE::Trace Stack;
};

struct NtCapturedEvent
{
    DWORD ThreadId;
    NtOperation Operation;
    const char *FunctionName;
    void *Caller;
    NTSTATUS Status;
    std::uint64_t Args[8];
    std::uint32_t DataSize;
    std::uint8_t DataSample[64];

    IC_STACKTRACE::Trace Stack;
};

struct KiCapturedEvent
{
    DWORD ThreadId;
    const char *StubName;
    void *Caller;
    void *StackPointer;

    IC_STACKTRACE::Trace Stack;
};

struct ModuleCapturedEvent
{
    DWORD ThreadId;
    ModuleHookOperation Operation;
    const char *FunctionName;
    const char *SourceModule;
    void *Caller;
    HMODULE ModuleHandle;
    std::vector<std::uint8_t> NameSample;
    std::uint64_t Args[4];

    IC_STACKTRACE::Trace Stack;
};

class WinsockHookController
{
  public:
    WinsockHookController() = default;
    ~WinsockHookController() = default;

    bool Initialize() noexcept;
    void Shutdown() noexcept;

    std::vector<WinsockCapturedEvent> ConsumeEvents();

  private:
    static void KeWinsockHookCallback(const WinsockHookContext &context) noexcept;
    static void EnqueueEvent(const WinsockHookContext &context);

    static bool s_Initialized;
    static std::mutex s_QueueMutex;
    static std::vector<WinsockCapturedEvent> s_Queue;
};

class NtHookController
{
  public:
    NtHookController() = default;
    ~NtHookController() = default;

    bool Initialize() noexcept;
    void Shutdown() noexcept;

    std::vector<NtCapturedEvent> ConsumeEvents();

  private:
    static void KeNtHookCallback(const NtHookContext &context) noexcept;
    static void EnqueueEvent(const NtHookContext &context);

    static bool s_Initialized;
    static std::mutex s_QueueMutex;
    static std::vector<NtCapturedEvent> s_Queue;
};

class KiHookController
{
  public:
    KiHookController() = default;
    ~KiHookController() = default;

    bool Initialize() noexcept;
    void Shutdown() noexcept;

    std::vector<KiCapturedEvent> ConsumeEvents();

  private:
    static void KeKiHookCallback(const KiHookContext &context) noexcept;
    static void EnqueueEvent(const KiHookContext &context);

    static bool s_Initialized;
    static std::mutex s_QueueMutex;
    static std::vector<KiCapturedEvent> s_Queue;
};

class ModuleHookController
{
  public:
    ModuleHookController() = default;
    ~ModuleHookController() = default;

    bool Initialize() noexcept;
    void Shutdown() noexcept;

    std::vector<ModuleCapturedEvent> ConsumeEvents();

  private:
    static void KeModuleHookCallback(const ModuleHookContext &context) noexcept;
    static void EnqueueEvent(const ModuleHookContext &context);

    static bool s_Initialized;
    static std::mutex s_QueueMutex;
    static std::vector<ModuleCapturedEvent> s_Queue;
};