vatrog-vm-signaling/src/si/vgpu-stream/win32/capture_nvfbc.c

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include "capture_nvfbc.h"
#include "capture-win32.h"   /* capture_thread_arg (win32-private) */
#include "present.h"
#include "cursor.h"          /* cursor_apply_shape / ctx->cursor */
#include "geometry.h"        /* reactive geometry resample on recreate */
#include "stream.h"          /* vgpu_publish_cursor / vgpu_publish_cursor_shape */
#include "nvfbc_tosys_c.h"

typedef struct {
    NvFBCToSys_c*              fbc;
    void*                      buf;
    NvFBC_CreateFunctionExType create;
    HCURSOR                    last_handle;   /* shape-gate by HCURSOR change */
} nvfbc_state;

/* Source the cursor for an NvFBC grab and publish it under the cursor_seq gate. NvFBC reports
 * only HW-cursor visibility (gi.bHWMouseVisible); position is not exposed, so one GetCursorInfo
 * per frame supplies x/y (the minimum possible). Shape is re-extracted only on HCURSOR change.
 * NvFBC composites the cursor itself (draw_cursor_cap==0) → present never reads ctx->cursor for
 * drawing, so no ctx->lock is required around the compose fields here.
 * gi.bProtectedContent / gi.dwSourcePID are available but out of scope (not in the contract). */
static void nvfbc_source_cursor(vgpu_ctx* ctx, nvfbc_state* st,
                                const NvFBCFrameGrabInfo* gi) {
    CURSORINFO ci; ci.cbSize = sizeof ci;
    int vis = gi->bHWMouseVisible ? 1 : 0;
    int x = ctx->cursor.x, y = ctx->cursor.y;
    if (GetCursorInfo(&ci)) {
        x = ci.ptScreenPos.x; y = ci.ptScreenPos.y;
        if (ci.hCursor && ci.hCursor != st->last_handle) {
            cursor_apply_shape(ctx, ci.hCursor);
            st->last_handle = ci.hCursor;
        }
    }
    ctx->cursor.visible = vis; ctx->cursor.x = x; ctx->cursor.y = y;

    vgpu_publish_cursor_shape(&ctx->view,
                              (uint32_t)ctx->cursor.hot_x, (uint32_t)ctx->cursor.hot_y,
                              (uint32_t)ctx->cursor.gw,    (uint32_t)ctx->cursor.gh,
                              (uint32_t)ctx->cursor.cursor_id);
    vgpu_publish_cursor(&ctx->view, (int32_t)x, (int32_t)y, (uint32_t)vis);
}

static NvFBCToSys_c* nvfbc_create(NvFBC_CreateFunctionExType pCreate, void** ppBuf) {
    NvFBCCreateParams cp; memset(&cp, 0, sizeof cp);
    cp.dwVersion       = NVFBC_CREATE_PARAMS_VER;
    cp.dwInterfaceType = NVFBC_TO_SYS_C;
    cp.dwAdapterIdx    = 0;
    if (pCreate(&cp) != NVFBC_SUCCESS || !cp.pNvFBC) return NULL;

    NvFBCToSys_c* fbc = (NvFBCToSys_c*)cp.pNvFBC;
    *ppBuf = NULL;

    NVFBC_TOSYS_SETUP_PARAMS_C sp; memset(&sp, 0, sizeof sp);
    sp.dwVersion = NVFBC_TOSYS_SETUP_PARAMS_VER_C;
    sp.bits      = 1u;                 /* bWithHWCursor = 1 (bit 0) */
    sp.eMode     = NVFBC_TOSYS_ARGB;
    sp.ppBuffer  = ppBuf;
    if (fbc->lpVtbl->NvFBCToSysSetUp(fbc, &sp) != NVFBC_SUCCESS || !*ppBuf) {
        fbc->lpVtbl->NvFBCToSysRelease(fbc);
        return NULL;
    }
    return fbc;
}

static DWORD WINAPI nvfbc_thread(LPVOID param) {
    capture_thread_arg* arg = (capture_thread_arg*)param;
    vgpu_ctx*    ctx = arg->ctx;
    nvfbc_state* st  = (nvfbc_state*)arg->backend_state;
    free(arg);

    NvFBCToSys_c* fbc = st->fbc;
    void*         buf = st->buf;

    for (;;) {
        NvFBCFrameGrabInfo gi; memset(&gi, 0, sizeof gi);
        NVFBC_TOSYS_GRAB_FRAME_PARAMS_C gp; memset(&gp, 0, sizeof gp);
        gp.dwVersion = NVFBC_TOSYS_GRAB_FRAME_PARAMS_VER_C;
        gp.dwFlags   = NVFBC_TOSYS_WAIT_WITH_TIMEOUT_C;
        gp.dwWaitTime = 1000;
        gp.eGMode    = NVFBC_TOSYS_SOURCEMODE_FULL;
        gp.pNvFBCFrameGrabInfo = &gi;

        NVFBCRESULT r = fbc->lpVtbl->NvFBCToSysGrabFrame(fbc, &gp);
        if (r != NVFBC_SUCCESS) {
            if (r == NVFBC_ERROR_INVALIDATED_SESSION || gi.bMustRecreate) {
                fprintf(stderr, "eyes(nvfbc): session invalidated (r=%d), recreating\n", (int)r);
                fbc->lpVtbl->NvFBCToSysRelease(fbc);
                fbc = NULL;
                while (!(fbc = nvfbc_create(st->create, &buf))) Sleep(200);
                st->fbc = fbc; st->buf = buf;
                /* grab session was recreated → display config may have changed: resample */
                geometry_sample_and_publish(ctx, 0, 0);
            } else {
                Sleep(50);
            }
            continue;
        }
        if (gi.dwWidth && gi.dwHeight)
            vgpu_present_submit(ctx, (const uint8_t*)buf,
                                gi.dwWidth, gi.dwHeight, gi.dwBufferWidth * 4u);
        nvfbc_source_cursor(ctx, st, &gi);
    }
    return 0;  /* unreachable; satisfies -Wreturn-type */
}

int nvfbc_start(vgpu_ctx* ctx, int fps) {
    (void)fps;
    HMODULE lib = LoadLibraryA("NvFBC64.dll");
    if (!lib) {
        fprintf(stderr, "eyes(nvfbc): LoadLibrary NvFBC64.dll failed (%lu)\n", GetLastError());
        return 0;
    }
    NvFBC_SetGlobalFlagsType    pSetFlags = (NvFBC_SetGlobalFlagsType)(void*)GetProcAddress(lib, "NvFBC_SetGlobalFlags");
    NvFBC_EnableFunctionType    pEnable   = (NvFBC_EnableFunctionType)(void*)GetProcAddress(lib, "NvFBC_Enable");
    NvFBC_CreateFunctionExType  pCreate   = (NvFBC_CreateFunctionExType)(void*)GetProcAddress(lib, "NvFBC_CreateEx");
    NvFBC_GetStatusExFunctionType pStatus = (NvFBC_GetStatusExFunctionType)(void*)GetProcAddress(lib, "NvFBC_GetStatusEx");
    if (!pEnable || !pCreate || !pStatus) {
        fprintf(stderr, "eyes(nvfbc): missing exports\n");
        return 0;
    }
    if (pSetFlags) pSetFlags(NVFBC_GLOBAL_FLAGS_NO_INITIAL_REFRESH);
    if (pEnable(NVFBC_STATE_ENABLE) != NVFBC_SUCCESS) {
        fprintf(stderr, "eyes(nvfbc): NvFBC_Enable failed\n");
        return 0;
    }
    NvFBCStatusEx stx; memset(&stx, 0, sizeof stx);
    stx.dwVersion = NVFBC_STATUS_VER; stx.dwAdapterIdx = 0;
    if (pStatus(&stx) != NVFBC_SUCCESS || !stx.bIsCapturePossible) {
        fprintf(stderr, "eyes(nvfbc): capture NOT possible on this GPU/license\n");
        return 0;
    }
    void* buf = NULL;
    NvFBCToSys_c* fbc = nvfbc_create(pCreate, &buf);
    if (!fbc) {
        fprintf(stderr, "eyes(nvfbc): CreateEx/ToSysSetUp failed\n");
        return 0;
    }

    nvfbc_state* st = (nvfbc_state*)malloc(sizeof *st);
    if (!st) { fbc->lpVtbl->NvFBCToSysRelease(fbc); return 0; }
    st->fbc = fbc; st->buf = buf; st->create = pCreate; st->last_handle = NULL;

    capture_thread_arg* arg = (capture_thread_arg*)malloc(sizeof *arg);
    if (!arg) { fbc->lpVtbl->NvFBCToSysRelease(fbc); free(st); return 0; }
    arg->ctx = ctx; arg->fps = fps; arg->backend_state = st;

    ctx->backend = VGPU_BK_NVFBC;
    ctx->draw_cursor_cap = 0;   /* NvFBC composites HW cursor itself */

    HANDLE t = CreateThread(NULL, 0, nvfbc_thread, arg, 0, NULL);
    if (!t) { fbc->lpVtbl->NvFBCToSysRelease(fbc); free(st); free(arg); return 0; }
    CloseHandle(t);

    fprintf(stderr, "eyes(nvfbc): session up (ToSys ARGB/BGRA), iface=0x%lx\n",
            (unsigned long)stx.dwNvFBCVersion);
    return 1;
}