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;
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 (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;
            if (offset < height*width) {
                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;
    }
    else if (strcmp(drm_ver->name, "nvidia-drm") == 0) {
        // quirky and slow
        funcs.convert = &convert_nvidia_kmsbuf;
    }
    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
        //funcs.sync_start = &drm_sync_noop;
        //funcs.sync_end = &drm_sync_noop;
    }
    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;
}