/* cuda_drv_test2.c - minimal CUDA driver-API smoke test (dlopen, no toolkit)
 * with copy-path pause points for live inspection.
 *
 * Same sequence as cuda_drv_test.c, plus:
 *   CDT_CP_PAUSE=<s>  sleep before cuMemcpyHtoD (after cuMemAlloc): the
 *                     context is live, the device pointer is known - probe
 *                     the GPU VA translation while parked.
 *   CDT_HP_PAUSE=<s>  sleep after cuMemcpyHtoD (before DtoH): inspect what
 *                     the HtoD left behind.
 *   CDT_WAIT=<s>      park after the test (existing behaviour).
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <dlfcn.h>
#include <stdint.h>

typedef int CUresult;
typedef int CUdevice;
typedef void *CUcontext;
typedef uint64_t CUdeviceptr;

static void fail_info(CUresult r)
{
    const char *(*errname)(CUresult, const char **) = dlsym(RTLD_DEFAULT, "cuGetErrorName");
    const char *(*errstr)(CUresult, const char **) = dlsym(RTLD_DEFAULT, "cuGetErrorString");
    const char *s = NULL;
    const char *hold = getenv("CDT_HOLD");

    if (errname && errname(r, &s) == 0 && s)
        printf("cuGetErrorName: %s\n", s);
    if (errstr && errstr(r, &s) == 0 && s)
        printf("cuGetErrorString: %s\n", s);
    if (hold && atoi(hold) > 0) {
        printf("CDT_HOLD: sleeping %d s for inspection\n", atoi(hold));
        fflush(stdout);
        sleep((unsigned)atoi(hold));
    }
}

#define CHECK(expr) do { \
    CUresult r = (expr); \
    if (r) { printf("%-24s FAILED rc=%d\n", #expr, r); fail_info(r); return 1; } \
    printf("%-24s OK\n", #expr); \
} while (0)

#define MARK(name) do { printf("-> %s\n", name); } while (0)

static void pause_env(const char *name)
{
    const char *v = getenv(name);

    if (v && atoi(v) > 0) {
        printf("%s: parked %d s\n", name, atoi(v));
        fflush(stdout);
        sleep((unsigned)atoi(v));
    }
}

/*
 * Scan this process's readable mappings for the payload pattern and print
 * every mapping that looks like a device/BAR window. This answers whether the
 * HtoD bytes are visible to the client CPU at all: the DtoH path can only be a
 * CPU read if some mapping of this process contains them.
 */
static void scan_self_for_payload(const char *tag)
{
    const unsigned long long pat = 0x0000000400000001ULL;
    FILE *f = fopen("/proc/self/maps", "r");
    char line[512];
    int hits = 0;

    printf("SCAN %s:\n", tag);
    if (!f) { printf("SCAN %s: no maps\n", tag); return; }
    while (fgets(line, sizeof(line), f)) {
        unsigned long long a = 0, b = 0, off = 0;
        char perms[8] = {0};
        char path[256] = {0};

        if (sscanf(line, "%llx-%llx %7s %llx %*s %*s %255[^\n]",
                   &a, &b, &perms, &off, path) < 4)
            continue;
        if (strstr(path, "vfio") || strstr(path, "nvidia") ||
            strstr(path, "/dev/"))
            printf("MAP %s: %llx-%llx %s off=%llx %s\n", tag, a, b,
                   perms, off, path);
        if (perms[0] != 'r' || b <= a || b - a > (512ULL << 20))
            continue;
        /*
         * Only scan anonymous mappings: reading file-backed or device
         * mappings can SIGBUS (past EOF) or block in a device fault path,
         * either of which kills/hangs the client before cuMemcpyDtoH. The
         * device mapping list above still reveals any BAR window.
         */
        if (path[0] && path[0] != '[')
            continue;
        {
            volatile unsigned long long *p;
            for (p = (volatile unsigned long long *)a;
                 (unsigned long long)p + 8 <= b; p++) {
                if (*p == pat && hits < 16) {
                    hits++;
                    printf("SCAN %s: payload at %p (map %llx-%llx %s %s)\n",
                           tag, (void *)p, a, b, perms, path);
                }
            }
        }
    }
    if (!hits)
        printf("SCAN %s: payload not visible in any mapping\n", tag);
    fclose(f);
}

int main(void)
{
    setvbuf(stdout, NULL, _IONBF, 0);

    void *h = dlopen("libcuda.so.1", RTLD_NOW);
    if (!h) { printf("dlopen libcuda.so.1 failed: %s\n", dlerror()); return 1; }

    CUresult (*cuInit)(unsigned int) = dlsym(h, "cuInit");
    CUresult (*cuDeviceGetCount)(int *) = dlsym(h, "cuDeviceGetCount");
    CUresult (*cuDeviceGet)(CUdevice *, int) = dlsym(h, "cuDeviceGet");
    CUresult (*cuCtxCreate)(CUcontext *, unsigned int, CUdevice) = dlsym(h, "cuCtxCreate_v2");
    CUresult (*cuMemAlloc)(CUdeviceptr *, size_t) = dlsym(h, "cuMemAlloc_v2");
    CUresult (*cuMemcpyHtoD)(CUdeviceptr, const void *, size_t) = dlsym(h, "cuMemcpyHtoD_v2");
    CUresult (*cuMemcpyDtoH)(void *, CUdeviceptr, size_t) = dlsym(h, "cuMemcpyDtoH_v2");
    CUresult (*cuMemFree)(CUdeviceptr) = dlsym(h, "cuMemFree_v2");
    CUresult (*cuCtxDestroy)(CUcontext) = dlsym(h, "cuCtxDestroy_v2");

    if (!cuInit || !cuDeviceGetCount || !cuDeviceGet || !cuCtxCreate ||
        !cuMemAlloc || !cuMemcpyHtoD || !cuMemcpyDtoH || !cuMemFree || !cuCtxDestroy) {
        printf("dlsym failed\n");
        return 1;
    }

    int count = 0;
    CUdevice dev;
    CUcontext ctx = NULL;
    CUdeviceptr dptr = 0;
    unsigned int buf[64];
    unsigned int back[64];
    size_t i;

    MARK("cuInit");
    CHECK(cuInit(0));
    MARK("cuDeviceGetCount");
    CHECK(cuDeviceGetCount(&count));
    printf("device count = %d\n", count);
    if (count < 1) return 1;

    MARK("cuDeviceGet");
    CHECK(cuDeviceGet(&dev, 0));
    printf("device 0 = %d\n", dev);

    MARK("cuCtxCreate");
    {
        const char *pause = getenv("CDT_PAUSE");
        if (pause && atoi(pause) > 0) {
            printf("CDT_PAUSE: sleeping %d s before cuCtxCreate\n", atoi(pause));
            fflush(stdout);
            sleep((unsigned)atoi(pause));
        }
    }
    CHECK(cuCtxCreate(&ctx, 0, dev));
    printf("context = %p\n", ctx);

    MARK("cuMemAlloc");
    CHECK(cuMemAlloc(&dptr, sizeof(buf)));
    printf("device ptr = 0x%llx (%zu bytes)\n", (unsigned long long)dptr, sizeof(buf));

    for (i = 0; i < 64; i++) buf[i] = (unsigned int)(i * 3 + 1);
    pause_env("CDT_CP_PAUSE");
    MARK("cuMemcpyHtoD");
    CHECK(cuMemcpyHtoD(dptr, buf, sizeof(buf)));
    pause_env("CDT_HP_PAUSE");
    if (getenv("CDT_SCAN"))
        scan_self_for_payload("after-HtoD");
    memset(back, 0, sizeof(back));
    MARK("cuMemcpyDtoH");
    CHECK(cuMemcpyDtoH(back, dptr, sizeof(buf)));
    if (getenv("CDT_SCAN"))
        scan_self_for_payload("after-DtoH");
    {
        size_t matched = 0;
        for (i = 0; i < 64; i++)
            if (buf[i] == back[i]) matched++;
        printf("roundtrip %s matched=%zu/64 back[0..7]=%u,%u,%u,%u,%u,%u,%u,%u expected 1,4,7,10,13,16,19,22\n",
               matched == 64 ? "match" : "MISMATCH", matched,
               back[0], back[1], back[2], back[3],
               back[4], back[5], back[6], back[7]);
    }

    MARK("cuMemFree");
    CHECK(cuMemFree(dptr));
    MARK("cuCtxDestroy");
    CHECK(cuCtxDestroy(ctx));
    printf("ALL OK\n");
    pause_env("CDT_WAIT");
    return 0;
}