kmsvnc/server.c

#include <rfb/rfb.h>
#include <fcntl.h>
#include <stdio.h>
#include <errno.h>
#include <sys/mman.h>

#include <xf86drm.h>
#include <libdrm/i915_drm.h>
#include <libdrm/amdgpu_drm.h>
#include <xf86drmMode.h>
#include <linux/uinput.h>
#include <linux/dma-buf.h>
#include <xkbcommon/xkbcommon.h>
#include <math.h>

#define ARRAY_SIZE(x) (sizeof(x) / sizeof(x[0]))
#define FPS 10
#define UINPUT_ABS_MAX INT16_MAX
#define UINPUT_MAX_KEY 256

struct Vec2d
{
    int x;
    int y;
};
struct Vec2d resolution;

#define SLEEPNS (1000000000 / FPS)
static void between_frames()
{
    static struct timespec now = {0, 0}, then = {0, 0}, tmp = {0, 0};

    clock_gettime(CLOCK_MONOTONIC, &now);
    memcpy((char *)&then, (char *)&tmp, sizeof(struct timespec));
    tmp.tv_nsec += SLEEPNS;
    if (tmp.tv_nsec >= 1000000000)
    {
        tmp.tv_sec++;
        tmp.tv_nsec %= 1000000000;
    }
    if (now.tv_sec < tmp.tv_sec || (now.tv_sec == tmp.tv_sec && now.tv_nsec < tmp.tv_nsec))
    {
        then.tv_sec = tmp.tv_sec - now.tv_sec;
        then.tv_nsec = tmp.tv_nsec - now.tv_nsec;
        if (then.tv_nsec < 0)
        {
            then.tv_sec--;
            then.tv_nsec += 1000000000;
        }
        nanosleep(&then, &then);
    }
    memcpy((char *)&now, (char *)&then, sizeof(struct timespec));
}

static void convert_bgrx_to_rgb(const char *in, int width, int height, char *buff)
{
    for (int y = 0; y < height; y++)
    {
        for (int x = 0; x < width; x++)
        {
            buff[(y * width + x) * 4] = in[(y * width + x) * 4 + 2];
            buff[(y * width + x) * 4 + 1] = in[(y * width + x) * 4 + 1];
            buff[(y * width + x) * 4 + 2] = in[(y * width + x) * 4];
        }
    }
}

char *kms_convert_buf = NULL;
size_t kms_convert_buf_len = 0;
char *kms_cpy_tmp_buf = NULL;
size_t kms_cpy_tmp_buf_len = 0;
static inline void convert_kmsbuf(const int XSTRIPE, const int YSTRIPE, const char *in, int width, int height, char *buff)
{
    if (width % XSTRIPE)
    {
        return;
    }
    if (height % YSTRIPE)
    {
        int sno = (width / XSTRIPE) + (height / YSTRIPE) * (width / XSTRIPE);
        int ord = (width % XSTRIPE) + (height % YSTRIPE) * XSTRIPE;
        int max_offset = sno * XSTRIPE * YSTRIPE + ord;
        if (kms_cpy_tmp_buf_len < max_offset * 4 + 4)
        {
            if (kms_cpy_tmp_buf)
                free(kms_convert_buf);
            kms_cpy_tmp_buf = malloc(max_offset * 4 + 4);
            kms_cpy_tmp_buf_len = max_offset * 4 + 4;
        }
        memcpy(kms_cpy_tmp_buf, in, max_offset * 4 + 4);
        in = (const char *)kms_cpy_tmp_buf;
    }
    if (kms_convert_buf_len < width * height * 4)
    {
        if (kms_convert_buf)
            free(kms_convert_buf);
        kms_convert_buf = malloc(width * height * 4);
        kms_convert_buf_len = width * height * 4;
    }
    for (int y = 0; y < height; y++)
    {
        for (int x = 0; x < width; x++)
        {
            int sno = (x / XSTRIPE) + (y / YSTRIPE) * (width / XSTRIPE);
            int ord = (x % XSTRIPE) + (y % YSTRIPE) * XSTRIPE;
            int offset = sno * XSTRIPE * YSTRIPE + ord;
            memcpy(kms_convert_buf + (x + y * width) * 4, in + offset * 4, 4);
        }
    }
    convert_bgrx_to_rgb(kms_convert_buf, width, height, buff);
}

#define XSTRIPE_INTEL 128
#define YSTRIPE_INTEL 8
#define XSTRIPE_NVIDIA 16
#define YSTRIPE_NVIDIA 128

static void convert_nvidia_kmsbuf(const char *in, int width, int height, char *buff)
{
    convert_kmsbuf(XSTRIPE_NVIDIA, YSTRIPE_NVIDIA, in, width, height, buff);
}
static void convert_intel_kmsbuf(const char *in, int width, int height, char *buff)
{
    convert_kmsbuf(XSTRIPE_INTEL, YSTRIPE_INTEL, in, width, height, buff);
}

struct vnc_drm_attr
{
    void (*sync_start)(int);
    void (*sync_end)(int);
    void (*convert)(const char *, int, int, char *);
};

static inline void drm_sync(int drmfd, uint64_t flags)
{
    int ioctl_err;
    struct dma_buf_sync sync = {
        .flags = flags,
    };
    if (ioctl_err = ioctl(drmfd, DMA_BUF_IOCTL_SYNC, &sync))
        fprintf(stderr, "DRM ioctl error %d on line %d\n", ioctl_err, __LINE__);
}
static void drm_sync_start(int drmfd)
{
    drm_sync(drmfd, DMA_BUF_SYNC_START | DMA_BUF_SYNC_READ);
}
static void drm_sync_end(int drmfd)
{
    drm_sync(drmfd, DMA_BUF_SYNC_END | DMA_BUF_SYNC_READ);
}
static void drm_sync_noop(int drmfd)
{
}

struct vnc_xkb
{
    struct xkb_context *ctx;
    struct xkb_keymap *map;
};
struct vnc_xkb xkb;

int uinput_fd = 0;
char keystate[UINPUT_MAX_KEY];
static void xkb_init()
{
    xkb.ctx = xkb_context_new(XKB_CONTEXT_NO_FLAGS);
    if (xkb.ctx == NULL)
    {
        fprintf(stderr, "Failed to create XKB context\n");
        exit(1);
    }
    struct xkb_rule_names names = {
        .rules = "",
        .model = "",
        .layout = "us",
        .variant = "",
        .options = ""};
    xkb.map = xkb_keymap_new_from_names(xkb.ctx, &names, 0);
    if (xkb.map == NULL)
    {
        fprintf(stderr, "Failed to create XKB map\n");
        exit(1);
    }
    // printf("xkb: keymap = %s\n", xkb_keymap_get_as_string(xkb.map, XKB_KEYMAP_USE_ORIGINAL_FORMAT));
}

static void uinput_init()
{
    struct uinput_setup usetup;
    uinput_fd = open("/dev/uinput", O_WRONLY | O_NONBLOCK);
    if (uinput_fd <= 0)
    {
        fprintf(stderr, "Failed to open uinput\n");
        exit(1);
    }
    ioctl(uinput_fd, UI_SET_EVBIT, EV_KEY);
    ioctl(uinput_fd, UI_SET_EVBIT, EV_SYN);
    for (int i = 0; i < UINPUT_MAX_KEY; i++)
    {
        ioctl(uinput_fd, UI_SET_KEYBIT, i);
    }

    ioctl(uinput_fd, UI_SET_EVBIT, EV_ABS);
    ioctl(uinput_fd, UI_SET_ABSBIT, ABS_X);
    ioctl(uinput_fd, UI_SET_ABSBIT, ABS_Y);

    ioctl(uinput_fd, UI_SET_KEYBIT, BTN_LEFT);
    ioctl(uinput_fd, UI_SET_KEYBIT, BTN_MIDDLE);
    ioctl(uinput_fd, UI_SET_KEYBIT, BTN_RIGHT);
    ioctl(uinput_fd, UI_SET_EVBIT, EV_REL);
    ioctl(uinput_fd, UI_SET_RELBIT, REL_WHEEL);

    struct uinput_abs_setup abs;
    memset(&abs, 0, sizeof(abs));
    abs.absinfo.maximum = UINPUT_ABS_MAX;
    abs.absinfo.minimum = 0;
    abs.code = ABS_X;
    ioctl(uinput_fd, UI_ABS_SETUP, &abs);
    abs.code = ABS_Y;
    ioctl(uinput_fd, UI_ABS_SETUP, &abs);

    memset(&usetup, 0, sizeof(usetup));
    usetup.id.bustype = BUS_USB;
    usetup.id.vendor = 0x0011;
    usetup.id.product = 0x4514;
    strcpy(usetup.name, "kmsvnc-device");

    ioctl(uinput_fd, UI_DEV_SETUP, &usetup);
    ioctl(uinput_fd, UI_DEV_CREATE);

    memset(keystate, 0, UINPUT_MAX_KEY);
}

struct key_iter_search
{
    xkb_keysym_t keysym;

    xkb_keycode_t keycode;
    xkb_level_index_t level;
};

static void key_iter(struct xkb_keymap *xkb, xkb_keycode_t key, void *data)
{
    struct key_iter_search *search = data;
    if (search->keycode != XKB_KEYCODE_INVALID)
    {
        return; // We are done
    }
    xkb_level_index_t num_levels = xkb_keymap_num_levels_for_key(xkb, key, 0);
    for (xkb_level_index_t i = 0; i < num_levels; i++)
    {
        const xkb_keysym_t *syms;
        int num_syms = xkb_keymap_key_get_syms_by_level(xkb, key, 0, i, &syms);
        for (int k = 0; k < num_syms; k++)
        {
            if (syms[k] == search->keysym)
            {
                search->keycode = key;
                search->level = i;
                break;
                goto end;
            }
        }
    }
end:
    return;
}

static void rfb_key_hook(rfbBool down, rfbKeySym keysym, rfbClientPtr cl)
{
    struct key_iter_search search = {
        .keysym = keysym,
        .keycode = XKB_KEYCODE_INVALID,
        .level = 0,
    };
    xkb_keymap_key_for_each(xkb.map, key_iter, &search);
    if (search.keycode == XKB_KEYCODE_INVALID)
    {
        fprintf(stderr, "Keysym %04x not found in our keymap\n", keysym);
        return;
    }
    // printf("key %s, keysym %04x, keycode %u\n", down ? "down" : "up", keysym, search.keycode);
    if (search.keycode >= UINPUT_MAX_KEY)
    {
        fprintf(stderr, "Keycode %d >= %d\n", search.keycode, UINPUT_MAX_KEY);
        return;
    }
    if (down != keystate[search.keycode])
    {
        struct input_event ies[] = {
            {
                .type = EV_KEY,
                .code = search.keycode - 8, // magic
                .value = down,
                .time.tv_sec = 0,
                .time.tv_usec = 0,
            },
            {
                .type = EV_SYN,
                .code = SYN_REPORT,
                .value = 0,
            },
        };
        for (int i = 0; i < ARRAY_SIZE(ies); i++)
        {
            write(uinput_fd, &ies[i], sizeof(ies[0]));
        }

        keystate[search.keycode] = down;
    }
}

static void rfb_ptr_hook(int mask, int screen_x, int screen_y, rfbClientPtr cl)
{
    // printf("pointer to %d, %d\n", screen_x, screen_y);
    float global_x = (float)screen_x;
    float global_y = (float)screen_y;
    int touch_x = round(global_x / resolution.x * UINPUT_ABS_MAX);
    int touch_y = round(global_y / resolution.y * UINPUT_ABS_MAX);
    struct input_event ies1[] = {
        {
            .type = EV_ABS,
            .code = ABS_X,
            .value = touch_x,
        },
        {
            .type = EV_ABS,
            .code = ABS_Y,
            .value = touch_y,
        },
        {.type = EV_KEY,
         .code = BTN_LEFT,
         .value = !!(mask & 0b1)},
        {.type = EV_KEY,
         .code = BTN_MIDDLE,
         .value = !!(mask & 0b10)},
        {.type = EV_KEY,
         .code = BTN_RIGHT,
         .value = !!(mask & 0b100)},
        {
            .type = EV_SYN,
            .code = SYN_REPORT,
            .value = 0,
        },
    };
    for (int i = 0; i < ARRAY_SIZE(ies1); i++)
    {
        write(uinput_fd, &ies1[i], sizeof(ies1[0]));
    }
    if (mask & 0b11000)
    {
        struct input_event ies2[] = {
            {
                .type = EV_REL,
                .code = REL_WHEEL,
                .value = mask & 0b1000 ? 1 : -1,
            },
            {
                .type = EV_SYN,
                .code = SYN_REPORT,
                .value = 0,
            },
        };
        for (int i = 0; i < ARRAY_SIZE(ies2); i++)
        {
            write(uinput_fd, &ies2[i], sizeof(ies2[0]));
        }
    }
}

int main(int argc, const char **argv)
{
    if (argc < 2)
    {
        printf("not enough arguments\n");
        return 1;
    }
    xkb_init();
    uinput_init();

    const char *card = argv[1];
    const int drmfd = open(card, O_RDONLY);
    int primefd = -1;
    if (drmfd < 0)
    {
        fprintf(stderr, "card %s open failed: %s\n", card, strerror(errno));
        return 1;
    }
    drmVersionPtr drm_ver = drmGetVersion(drmfd);
    int err;
    int source_plane = 0;
    int source_crtc = 0;
    err = drmSetClientCap(drmfd, DRM_CLIENT_CAP_UNIVERSAL_PLANES, 1);
    if (err < 0)
    {
        perror("Failed to set universal planes capability: primary planes will not be usable");
    }
    drmModePlane *plane = NULL;
    drmModePlaneRes *plane_res = NULL;
    drmModeFB *fb = NULL;
    if (source_plane > 0)
    {
        plane = drmModeGetPlane(drmfd, source_plane);
        if (!plane)
        {
            fprintf(stderr, "Failed to get plane %d: %s\n", source_plane, strerror(errno));
            goto cleanup;
        }
        if (plane->fb_id == 0)
        {
            fprintf(stderr, "Place %d does not have an attached framebuffer\n", source_plane);
        }
    }
    else
    {
        plane_res = drmModeGetPlaneResources(drmfd);
        if (!plane_res)
        {
            perror("Failed to get plane resources");
            goto cleanup;
        }
        int i;
        for (i = 0; i < plane_res->count_planes; i++)
        {
            plane = drmModeGetPlane(drmfd, plane_res->planes[i]);
            if (!plane)
            {
                fprintf(stderr, "Failed to get plane %u: %s\n", plane_res->planes[i], strerror(errno));
                continue;
            }
            printf("Plane %u CRTC %u FB %u\n", plane->plane_id, plane->crtc_id, plane->fb_id);
            if ((source_crtc > 0 && plane->crtc_id != source_crtc) || plane->fb_id == 0)
            {
                // Either not connected to the target source CRTC
                // or not active.
                drmModeFreePlane(plane);
                plane = NULL;
                continue;
            }
            break;
        }
        if (i == plane_res->count_planes)
        {
            if (source_crtc > 0)
            {
                fprintf(stderr, "No usable planes found on CRTC %d\n", source_crtc);
            }
            else
            {
                fprintf(stderr, "No usable planes found\n");
            }
            goto cleanup;
        }
        printf("Using plane %u to locate framebuffers\n", plane->plane_id);
    }
    uint32_t plane_id = plane->plane_id;

    fb = drmModeGetFB(drmfd, plane->fb_id);
    if (!fb)
    {
        fprintf(stderr, "Failed to get framebuffer %u: %s\n", plane->fb_id, strerror(errno));
        goto cleanup;
    }
    printf("Template framebuffer is %u: %ux%u %ubpp %ub depth %u pitch\n", fb->fb_id, fb->width, fb->height, fb->bpp, fb->depth, fb->pitch);

    if (fb->bpp != 32 || fb->depth != 24)
    {
        fprintf(stderr, "Unsupported pixfmt\n");
        goto cleanup;
    }

    if (!fb->handle)
    {
        fprintf(stderr, "No handle set on framebuffer: maybe you need some additional capabilities?\n");
        goto cleanup;
    }

    int ioctl_err = 0;

    printf("drm driver is %s\n", drm_ver->name);
    char *mapped = NULL;
    struct vnc_drm_attr funcs = {
        .sync_start = &drm_sync_start,
        .sync_end = &drm_sync_end,
        .convert = &convert_bgrx_to_rgb,
    };
    err = drmPrimeHandleToFD(drmfd, fb->handle, O_RDWR, &primefd);
    if (err < 0 || primefd < 0)
    {
        perror("Failed to get PRIME fd from framebuffer handle");
        goto cleanup;
    }

    int mmap_fd = primefd;
    size_t mmap_size = fb->width * fb->height * 32 / 8;
    off_t mmap_offset = 0;

    if (strcmp(drm_ver->name, "i915") == 0)
    {
        funcs.convert = &convert_intel_kmsbuf;
    }
    else if (strcmp(drm_ver->name, "amdgpu") == 0)
    {
        struct drm_gem_flink flink;
        flink.handle = fb->handle;
        if (ioctl_err = drmIoctl(drmfd, DRM_IOCTL_GEM_FLINK, &flink))
        {
            fprintf(stderr, "DRM ioctl error %d on line %d\n", ioctl_err, __LINE__);
            goto cleanup;
        }
        struct drm_gem_open open_arg;
        open_arg.name = flink.name;
        printf("global name = %d\n", flink.name);
        if (ioctl_err = drmIoctl(drmfd, DRM_IOCTL_GEM_OPEN, &open_arg))
        {
            fprintf(stderr, "DRM ioctl error %d on line %d\n", ioctl_err, __LINE__);
            goto cleanup;
        }
        union drm_amdgpu_gem_mmap mmap_arg;
        memset(&mmap_arg, 0, sizeof(mmap_arg));
        mmap_arg.in.handle = open_arg.handle;
        if (ioctl_err = drmIoctl(drmfd, DRM_IOCTL_AMDGPU_GEM_MMAP, &mmap_arg))
        {
            fprintf(stderr, "DRM ioctl error %d on line %d\n", ioctl_err, __LINE__);
            goto cleanup;
        }
        mmap_size = open_arg.size;
        mmap_offset = mmap_arg.out.addr_ptr;
        mmap_fd = drmfd;
    }
    else if (strcmp(drm_ver->name, "nvidia-drm") == 0)
    {
        // quirky and slow
        funcs.convert = &convert_nvidia_kmsbuf;

        struct drm_gem_flink flink;
        flink.handle = fb->handle;
        if (ioctl_err = drmIoctl(drmfd, DRM_IOCTL_GEM_FLINK, &flink))
        {
            fprintf(stderr, "DRM ioctl error %d on line %d\n", ioctl_err, __LINE__);
            goto cleanup;
        }
        struct drm_gem_open open_arg;
        open_arg.name = flink.name;
        printf("global name = %d\n", flink.name);
        if (ioctl_err = drmIoctl(drmfd, DRM_IOCTL_GEM_OPEN, &open_arg))
        {
            fprintf(stderr, "DRM ioctl error %d on line %d\n", ioctl_err, __LINE__);
            goto cleanup;
        }
        struct drm_mode_map_dumb mreq;
        memset(&mreq, 0, sizeof(mreq));
        mreq.handle = open_arg.handle;
        if (ioctl_err = drmIoctl(drmfd, DRM_IOCTL_MODE_MAP_DUMB, &mreq))
        {
            fprintf(stderr, "DRM ioctl error %d on line %d\n", ioctl_err, __LINE__);
            goto cleanup;
        }
        mmap_size = open_arg.size;
        mmap_offset = mreq.offset;
        mmap_fd = drmfd;
        funcs.sync_start = &drm_sync_noop;
        funcs.sync_end = &drm_sync_noop;
    }
    else if (strcmp(drm_ver->name, "vmwgfx") == 0 || strcmp(drm_ver->name, "vboxvideo") == 0 || strcmp(drm_ver->name, "virtio_gpu") == 0)
    {
        // virgl does not work
    }
    else
    {
        fprintf(stderr, "Untested drm driver, use at your own risk!\n");
    }

    if (!mapped)
    {
        printf("mapping with size = %d, offset = %d, fd = %d\n", mmap_size, mmap_offset, mmap_fd);
        mapped = mmap(NULL, mmap_size, PROT_READ, MAP_SHARED, mmap_fd, mmap_offset);
        if (mapped == MAP_FAILED)
        {
            perror("mmap");
            goto cleanup;
        }
    }

    size_t buflen = fb->width * fb->height * 32 / 8;
    char *buf = malloc(buflen);
    memset(buf, 0, buflen);

    resolution.x = fb->width;
    resolution.y = fb->height;
    rfbScreenInfoPtr server = rfbGetScreen(0, NULL, fb->width, fb->height, 8, 3, 32 / 8);
    if (!server)
        return 1;
    server->desktopName = "kmsvnc";
    server->frameBuffer = buf;
    server->port = 5900;
    // server->listenInterface = inet_addr("127.0.0.1");
    server->ipv6port = 0;
    server->listen6Interface = NULL;
    server->alwaysShared = (1 == 1);
    server->kbdAddEvent = rfb_key_hook;
    server->ptrAddEvent = rfb_ptr_hook;
    rfbInitServer(server);
    rfbRunEventLoop(server, -1, TRUE);
    while (rfbIsActive(server))
    {
        between_frames();
        if (server->clientHead)
        {
            funcs.sync_start(primefd);
            funcs.convert(mapped, fb->width, fb->height, buf);
            funcs.sync_end(primefd);
            rfbMarkRectAsModified(server, 0, 0, fb->width, fb->height);
        }
    }

cleanup:
    if (drm_ver != NULL)
        drmFreeVersion(drm_ver);
    if (fb != NULL)
        drmModeFreeFB(fb);
    if (uinput_fd > 0)
    {
        ioctl(uinput_fd, UI_DEV_DESTROY);
        close(uinput_fd);
    }
    if (primefd > 0)
        close(primefd);
    close(drmfd);
    return 1;
}