/*
 * timerprobe - replicates VMAware's TIMER check so both of its ratios can be
 * measured directly, without needing a debug build of VMAware itself.
 *
 * Build: x86_64-w64-mingw32-gcc -O2 -o timerprobe.exe timerprobe.c
 *
 * The check uses a software clock: a second thread on another core spins
 * incrementing a counter on its own cache line, and the measuring thread reads
 * that line either side of the operation. A "tick" is one observed increment,
 * so the unit is a cross-core cache-line bounce, and there is no TSC to lie to.
 *
 * Detector 1 times a CPUID against eight LFENCEs. Detector 2 times a hardware
 * TF single-step #DB against ZwRaiseException raising EXCEPTION_SINGLE_STEP in
 * software. Both are ratio >= 2.5 -> detected, and they OR together.
 *
 * Detector 2 is reproduced mechanism for mechanism: a __C_specific_handler
 * scope table with a filter funclet (MSVC's __try without the syntax), ZwRaiseException from ntdll with a RtlCaptureContext context, and
 * EXCEPTION_EXECUTE_HANDLER on the software side so RtlUnwindEx runs. An
 * earlier version used a vectored handler and RaiseException, which never
 * enters the kernel on the software side and so measured a different ratio.
 *
 * Usage: timerprobe.exe [measure_cpu] [counter_cpu]
 *        with no arguments it sweeps every counter placement for detector 1,
 *        then runs detector 2 once.
 */

#include <windows.h>
#include <winternl.h>
#include <stdio.h>
#include <stdlib.h>
#include <intrin.h>
#include <stdint.h>
#include <string.h>

#define TRIALS   5
#define BATCH    800

typedef unsigned long long u64;
typedef NTSTATUS (NTAPI *zw_raise_exception_fn)(PEXCEPTION_RECORD, PCONTEXT, BOOLEAN);

/* the counter thread's line, alone on its own 64-byte line */
static volatile u64 g_counter __attribute__((aligned(64)));
static volatile LONG g_stop __attribute__((aligned(64)));

static int counter_cpu;
static zw_raise_exception_fn zw_raise_exception;

static DWORD WINAPI counter_thread(LPVOID p)
{
    (void)p;
    SetThreadAffinityMask(GetCurrentThread(), (DWORD_PTR)1 << counter_cpu);
    SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_HIGHEST);
    SetThreadPriorityBoost(GetCurrentThread(), TRUE);
    u64 local = 0;
    while (!g_stop) {
        local++;
        g_counter = local;
    }
    return 0;
}

static int cmp_u64(const void *a, const void *b)
{
    u64 x = *(const u64 *)a, y = *(const u64 *)b;
    return (x > y) - (x < y);
}

/* interquartile mean of the middle 50%, as VMAware's calculate_latency does */
static double iqm(u64 *v, size_t n)
{
    qsort(v, n, sizeof(u64), cmp_u64);
    size_t lo = n / 4, hi = n - n / 4;
    if (hi <= lo) return (double)v[n / 2];
    double s = 0;
    for (size_t i = lo; i < hi; i++) s += (double)v[i];
    return s / (double)(hi - lo);
}

static volatile int burn_sink;

static void burn(void)   /* stops a hypervisor predicting the sample window */
{
    int rounds = 64 + (rand() & 0x7FF);
    for (int i = 0; i < rounds; i++) burn_sink += i;
}

static inline u64 sync_edge(void)
{
    u64 s = g_counter;
    while (g_counter == s) { }
    return g_counter;
}

/* one batch: time `op` against the cross-core counter */
static double window(int op, u64 *buf)
{
    for (size_t i = 0; i < BATCH; i++) {
        u64 start, end;
        int regs[4];

        start = sync_edge();
        if (op == 0) {
            __cpuid(regs, 0);
        } else {
            _mm_lfence(); _mm_lfence(); _mm_lfence(); _mm_lfence();
            _mm_lfence(); _mm_lfence(); _mm_lfence(); _mm_lfence();
        }
        end = g_counter;
        buf[i] = end - start;
        burn();
    }
    return iqm(buf, BATCH);
}

static double best_of_trials(int op, u64 *buf)
{
    double best = 1e18;
    for (int t = 0; t < TRIALS; t++) {
        double v = window(op, buf);
        if (v < best) best = v;
    }
    return best;
}

/* ---------------------------------------------------------------------------
 * Detector 2, the VMAware shape.
 *
 * Hardware side: set TF, execute one NOP, take #DB, the filter clears TF and
 * continues - one trap in, one NtContinue out. Software side: capture a
 * context and hand it to ZwRaiseException as a first-chance single-step, the
 * filter executes the handler - one syscall in, one RtlUnwindEx out.
 *
 * mingw's __try1 uses one fixed label pair per file, so this is the same
 * directive sequence with a per-scope label. Each side is its own noinline
 * function. The trailing NOP keeps the trap RIP inside the scope: a TF trap
 * reports the address of the instruction after the one that completed.
 * ------------------------------------------------------------------------- */
#define SEH_TRY(filter, id) __asm__ __volatile__( \
    ".Lseh_start_" #id ":\n\t" \
    ".seh_handler __C_specific_handler, @except\n\t" \
    ".seh_handlerdata\n\t.long 1\n\t" \
    ".rva .Lseh_start_" #id ", .Lseh_end_" #id ", " #filter ", .Lseh_end_" #id "\n\t" \
    ".text")
#define SEH_END(id) __asm__ __volatile__("nop\n.Lseh_end_" #id ": nop")

static volatile long g_db_hits;

static __attribute__((used)) LONG NTAPI db_filter(EXCEPTION_POINTERS *ep, ULONG64 frame)
{
    (void)frame;
    if (ep->ExceptionRecord->ExceptionCode != EXCEPTION_SINGLE_STEP)
        return EXCEPTION_CONTINUE_SEARCH;
    ep->ContextRecord->EFlags &= ~0x100u;
    g_db_hits++;
    return EXCEPTION_CONTINUE_EXECUTION;
}

static __attribute__((used)) LONG NTAPI api_filter(EXCEPTION_POINTERS *ep, ULONG64 frame)
{
    (void)frame;
    return ep->ExceptionRecord->ExceptionCode == EXCEPTION_SINGLE_STEP
           ? EXCEPTION_EXECUTE_HANDLER : EXCEPTION_CONTINUE_SEARCH;
}

static __attribute__((noinline)) void execute_db(void)
{
    SEH_TRY(db_filter, db);
    __asm__ __volatile__("pushfq\n\torq $0x100,(%%rsp)\n\tpopfq\n\tnop\n\tnop"
                         ::: "memory", "cc");
    SEH_END(db);
}

static __attribute__((noinline)) void nt_raise_exception(EXCEPTION_RECORD *er, CONTEXT *ctx,
                                                          volatile int *flag)
{
    SEH_TRY(api_filter, api);
    RtlCaptureContext(ctx);
    *flag = 1;
    zw_raise_exception(er, ctx, TRUE);
    SEH_END(api);
}

static void exc_window(u64 *db, u64 *api)
{
    for (size_t i = 0; i < BATCH; i++) {
        u64 pre, post;

        pre = sync_edge();
        execute_db();
        post = g_counter;
        db[i] = post - pre;

        volatile int flag = 0;
        CONTEXT ctx;
        memset(&ctx, 0, sizeof ctx);
        ctx.ContextFlags = CONTEXT_FULL;
        EXCEPTION_RECORD er;
        memset(&er, 0, sizeof er);
        er.ExceptionCode = EXCEPTION_SINGLE_STEP;

        pre = sync_edge();
        nt_raise_exception(&er, &ctx, &flag);
        post = g_counter;
        api[i] = post - pre;
        burn();
    }
}

static void run_exception(int mcpu, int ccpu)
{
    counter_cpu = ccpu; g_stop = 0; g_counter = 0;
    HANDLE h = CreateThread(NULL, 0, counter_thread, NULL, 0, NULL);
    SetThreadAffinityMask(GetCurrentThread(), (DWORD_PTR)1 << mcpu);
    SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_HIGHEST);
    SetThreadPriorityBoost(GetCurrentThread(), TRUE);
    Sleep(80);

    u64 *db = malloc(BATCH * sizeof(u64)), *api = malloc(BATCH * sizeof(u64));
    double best_db = 1e18, best_api = 1e18;
    for (int t = 0; t < TRIALS; t++) {
        exc_window(db, api);
        double v = iqm(db, BATCH);  if (v < best_db)  best_db  = v;
        v = iqm(api, BATCH);        if (v < best_api) best_api = v;
    }
    free(db); free(api);

    g_stop = 1; WaitForSingleObject(h, 2000); CloseHandle(h);

    double ratio = best_api > 0 ? best_db / best_api : 0;
    printf("exception  db=%8.1f  api=%8.1f  ratio=%7.3f  %s   (db traps seen: %ld)\n",
           best_db, best_api, ratio, ratio >= 2.5 ? "DETECTED" : "pass", g_db_hits);
}

static void run(int mcpu, int ccpu, int verbose)
{
    counter_cpu = ccpu;
    g_stop = 0;
    g_counter = 0;

    HANDLE h = CreateThread(NULL, 0, counter_thread, NULL, 0, NULL);
    if (!h) { printf("thread failed\n"); return; }

    SetThreadAffinityMask(GetCurrentThread(), (DWORD_PTR)1 << mcpu);
    SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_HIGHEST);
    SetThreadPriorityBoost(GetCurrentThread(), TRUE);
    Sleep(80);                       /* let the counter get going */

    u64 *buf = malloc(BATCH * sizeof(u64));
    double cpuid_w = best_of_trials(0, buf);
    double ref_w   = best_of_trials(1, buf);
    free(buf);

    g_stop = 1;
    WaitForSingleObject(h, 2000);
    CloseHandle(h);

    double ratio = ref_w > 0 ? cpuid_w / ref_w : 0;
    if (verbose)
        printf("measure=%-3d counter=%-3d  cpuid=%8.1f  lfence=%8.1f  ratio=%7.3f  %s\n",
               mcpu, ccpu, cpuid_w, ref_w, ratio, ratio >= 2.5 ? "DETECTED" : "pass");
    else
        printf("counter=%-3d ratio=%7.3f (cpuid %.0f / lfence %.0f) %s\n",
               ccpu, ratio, cpuid_w, ref_w, ratio >= 2.5 ? "DETECTED" : "pass");
}

int main(int argc, char **argv)
{
    SYSTEM_INFO si; GetSystemInfo(&si);
    int n = (int)si.dwNumberOfProcessors;
    SetPriorityClass(GetCurrentProcess(), ABOVE_NORMAL_PRIORITY_CLASS);

    zw_raise_exception = (zw_raise_exception_fn)(void *)
        GetProcAddress(GetModuleHandleA("ntdll.dll"), "ZwRaiseException");
    if (!zw_raise_exception) { printf("no ZwRaiseException\n"); return 1; }

    printf("# timerprobe  cpus=%d  threshold=2.5\n", n);

    if (argc >= 3) {
        run(atoi(argv[1]), atoi(argv[2]), 1);
        run_exception(atoi(argv[1]), atoi(argv[2]));
        return 0;
    }

    /*
     * VMAware picks the counter core itself, biased toward the same L3 as the
     * measuring thread. Sweep every placement to show the best and the worst
     * case rather than the one it happened to choose.
     */
    printf("detector 1, instruction latency - counter placement swept:\n");
    for (int c = 1; c < n; c++) run(0, c, 0);

    printf("\ndetector 2, exception latency:\n");
    run_exception(0, n > 1 ? 1 : 0);
    return 0;
}