vnc: Add Open H.264 encoding extension

include/open-h264.h

@@ -0,0 +1,16 @@
+#pragma once
+
+#define OPEN_H264_MAX_CONTEXTS 64
+
+#include <stdbool.h>
+#include <rfb/rfbclient.h>
+
+struct AVFrame;
+struct open_h264;
+struct open_h264_context;
+
+struct open_h264* open_h264_create(rfbClient* client);
+void open_h264_destroy(struct open_h264*);
+
+struct AVFrame* open_h264_decode_rect(struct open_h264* self,
+		rfbFramebufferUpdateRectHeader* message);

include/vnc.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020 Andri Yngvason
+ * Copyright (c) 2020 - 2022 Andri Yngvason
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -22,9 +22,24 @@
 #include <wayland-client.h>
 #include <rfb/rfbclient.h>
 
+#define VNC_CLIENT_MAX_AV_FRAMES 64
+
+struct open_h264;
+struct AVFrame;
+
+struct vnc_av_frame {
+	struct AVFrame* frame;
+	int x, y, width, height;
+};
+
 struct vnc_client {
 	rfbClient* client;
 
+	struct open_h264* open_h264;
+	bool current_rect_is_av_frame;
+	struct vnc_av_frame* av_frames[VNC_CLIENT_MAX_AV_FRAMES];
+	int n_av_frames;
+
 	int (*alloc_fb)(struct vnc_client*);
 	void (*update_fb)(struct vnc_client*);
 	void (*cut_text)(struct vnc_client*, const char*, size_t);

meson.build

@@ -37,6 +37,8 @@ drm = dependency('libdrm')
 gbm = dependency('gbm')
 egl = dependency('egl')
 glesv2 = dependency('glesv2')
+lavc = dependency('libavcodec')
+lavu = dependency('libavutil')
 
 libvncserver_opt = cmake.subproject_options()
 libvncserver_opt.add_cmake_defines({
@@ -78,6 +80,7 @@ sources = [
 	'src/renderer.c',
 	'src/renderer-egl.c',
 	'src/buffer.c',
+	'src/open-h264.c',
 ]
 
 dependencies = [
@@ -92,6 +95,8 @@ dependencies = [
 	gbm,
 	egl,
 	glesv2,
+	lavc,
+	lavu,
 	libvncclient,
 	client_protos,
 ]

src/open-h264.c

@@ -0,0 +1,252 @@
+#include "open-h264.h"
+
+#include <stdint.h>
+#include <stdbool.h>
+#include <string.h>
+#include <arpa/inet.h>
+#include <rfb/rfbclient.h>
+#include <libavcodec/avcodec.h>
+
+enum open_h264_flags {
+	OPEN_H264_RESET_CONTEXT = 1 << 0,
+	OPEN_H264_RESET_ALL_CONTEXTS = 1 << 1,
+};
+
+struct open_h264_msg_head {
+	uint32_t length;
+	uint32_t flags;
+} __attribute__((packed));
+
+struct open_h264_context {
+	rfbRectangle rect;
+
+	AVCodecParserContext* parser;
+	AVCodecContext* codec_ctx;
+	AVBufferRef* hwctx_ref;
+};
+
+struct open_h264 {
+	rfbClient* client;
+
+	struct open_h264_context* contexts[OPEN_H264_MAX_CONTEXTS];
+	int n_contexts;
+};
+
+static bool are_rects_equal(const rfbRectangle* a, const rfbRectangle* b)
+{
+	return memcmp(a, b, sizeof(*a)) == 0;
+}
+
+static int find_context_index(const struct open_h264* self,
+		const rfbRectangle* rect)
+{
+	for (int i = 0; i < self->n_contexts; ++i)
+		if (are_rects_equal(&self->contexts[i]->rect, rect))
+			return i;
+	return -1;
+}
+
+static struct open_h264_context* find_context(
+		const struct open_h264* self, const rfbRectangle* rect)
+{
+	int i = find_context_index(self, rect);
+	return i >= 0 ? self->contexts[i] : NULL;
+}
+
+static struct open_h264_context* open_h264_context_create(
+		struct open_h264* self, const rfbRectangle* rect)
+{
+	if (self->n_contexts >= OPEN_H264_MAX_CONTEXTS)
+		return NULL;
+
+	struct open_h264_context* context = calloc(1, sizeof(*context));
+	if (!context)
+		return NULL;
+
+	memcpy(&context->rect, rect, sizeof(context->rect));
+
+	const AVCodec* codec = avcodec_find_decoder(AV_CODEC_ID_H264);
+	if (!codec)
+		goto failure;
+
+	context->parser = av_parser_init(codec->id);
+	if (!context->parser)
+		goto failure;
+
+	context->parser->flags = PARSER_FLAG_COMPLETE_FRAMES;
+
+	context->codec_ctx = avcodec_alloc_context3(codec);
+	if (!context->codec_ctx)
+		goto failure;
+
+	if (av_hwdevice_ctx_create(&context->hwctx_ref, AV_HWDEVICE_TYPE_VAAPI,
+				NULL, NULL, 0) != 0)
+		goto failure;
+
+	context->codec_ctx->hw_device_ctx = av_buffer_ref(context->hwctx_ref);
+
+	if (avcodec_open2(context->codec_ctx, codec, NULL) != 0)
+		goto failure;
+
+	self->contexts[self->n_contexts++] = context;
+	return context;
+
+failure:
+	av_buffer_unref(&context->hwctx_ref);
+	avcodec_free_context(&context->codec_ctx);
+	av_parser_close(context->parser);
+	free(context);
+	return NULL;
+}
+
+static void open_h264_context_destroy(struct open_h264_context* context)
+{
+	av_buffer_unref(&context->hwctx_ref);
+	avcodec_free_context(&context->codec_ctx);
+	av_parser_close(context->parser);
+	free(context);
+}
+
+static struct open_h264_context* get_context(struct open_h264* self,
+		const rfbRectangle* rect)
+{
+	struct open_h264_context* context = find_context(self, rect);
+	return context ? context : open_h264_context_create(self, rect);
+}
+
+
+static void reset_context(struct open_h264* self,
+		const rfbRectangle* rect)
+{
+	int i = find_context_index(self, rect);
+	if (i < 0)
+		return;
+
+	open_h264_context_destroy(self->contexts[i]);
+	--self->n_contexts;
+
+	memmove(&self->contexts[i], &self->contexts[i + 1],
+			self->n_contexts - i);
+}
+
+static void reset_all_contexts(struct open_h264* self)
+{
+	for (int i = 0; i < self->n_contexts; ++i) {
+		open_h264_context_destroy(self->contexts[i]);
+		self->contexts[i] = NULL;
+	}
+}
+
+struct open_h264* open_h264_create(rfbClient* client)
+{
+	struct open_h264* self = calloc(1, sizeof(*self));
+	if (!self)
+		return NULL;
+
+	self->client = client;
+
+	return self;
+}
+
+void open_h264_destroy(struct open_h264* self)
+{
+	if (!self)
+		return;
+
+	reset_all_contexts(self);
+	free(self);
+}
+
+static bool decode_frame(struct open_h264_context* context, AVFrame* frame,
+		AVPacket* packet)
+{
+	int rc;
+
+	rc = avcodec_send_packet(context->codec_ctx, packet);
+	if (rc < 0)
+		return false;
+
+	struct AVFrame* vaapi_frame = av_frame_alloc();
+	if (!vaapi_frame)
+		return false;
+
+	rc = avcodec_receive_frame(context->codec_ctx, vaapi_frame);
+	if (rc < 0)
+		return false;
+
+	frame->format = AV_PIX_FMT_DRM_PRIME;
+
+	rc = av_hwframe_map(frame, vaapi_frame, AV_HWFRAME_MAP_DIRECT);
+	if (rc < 0) {
+		av_frame_free(&vaapi_frame);
+		return false;
+	}
+
+	av_frame_copy_props(frame, vaapi_frame);
+	av_frame_free(&vaapi_frame);
+
+	return true;
+}
+
+struct AVFrame* open_h264_decode_rect(struct open_h264* self,
+		rfbFramebufferUpdateRectHeader* msg)
+{
+	bool have_frame = false;
+
+	struct open_h264_msg_head head = { 0 };
+	if (!ReadFromRFBServer(self->client, (char*)&head, sizeof(head)))
+		return NULL;
+
+	uint32_t length = ntohl(head.length);
+	enum open_h264_flags flags = ntohl(head.flags);
+
+	if (flags & OPEN_H264_RESET_ALL_CONTEXTS) {
+		reset_all_contexts(self);
+	} else if (flags & OPEN_H264_RESET_CONTEXT) {
+		reset_context(self, &msg->r);
+	}
+
+	struct open_h264_context* context = get_context(self, &msg->r);
+	if (!context)
+		return NULL;
+
+	char* data = malloc(length);
+	if (!data)
+		return NULL;
+
+	AVFrame* frame = av_frame_alloc();
+	if (!frame)
+		goto failure;
+
+	AVPacket* packet = av_packet_alloc();
+	if (!packet)
+		goto failure;
+
+	if (!ReadFromRFBServer(self->client, data, length))
+		goto failure;
+
+	const uint8_t* dp = (const uint8_t*)data;
+
+	while (length > 0) {
+		int rc = av_parser_parse2(context->parser, context->codec_ctx,
+				&packet->data, &packet->size, dp, length,
+				AV_NOPTS_VALUE, AV_NOPTS_VALUE, 0);
+		if (rc < 0)
+			goto failure;
+
+		dp += rc;
+		length -= rc;
+
+		// If we get multiple frames per rect, there's no point in
+		// rendering them all, so we just return the last one.
+		if (packet->size != 0)
+			have_frame = decode_frame(context, frame, packet);
+	}
+
+failure:
+	av_packet_free(&packet);
+	if (!have_frame)
+		av_frame_unref(frame);
+	free(data);
+	return have_frame ? frame : NULL;
+}

src/vnc.c

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020 Andri Yngvason
+ * Copyright (c) 2020 - 2022 Andri Yngvason
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -22,8 +22,12 @@
 #include <pixman.h>
 #include <rfb/rfbclient.h>
 #include <libdrm/drm_fourcc.h>
+#include <libavutil/frame.h>
 
 #include "vnc.h"
+#include "open-h264.h"
+
+#define RFB_ENCODING_OPEN_H264 50
 
 extern const unsigned short code_map_linux_to_qnum[];
 extern const unsigned int code_map_linux_to_qnum_len;
@@ -42,10 +46,24 @@ static void vnc_client_update_box(rfbClient* client, int x, int y, int width,
 	struct vnc_client* self = rfbClientGetClientData(client, NULL);
 	assert(self);
 
+	if (self->current_rect_is_av_frame) {
+		self->current_rect_is_av_frame = false;
+		return;
+	}
+
 	pixman_region_union_rect(&self->damage, &self->damage, x, y, width,
 			height);
 }
 
+static void vnc_client_clear_av_frames(struct vnc_client* self)
+{
+	for (int i = 0; i < self->n_av_frames; ++i) {
+		av_frame_unref(self->av_frames[i]->frame);
+		free(self->av_frames[i]);
+	}
+	self->n_av_frames = 0;
+}
+
 static void vnc_client_finish_update(rfbClient* client)
 {
 	struct vnc_client* self = rfbClientGetClientData(client, NULL);
@@ -54,6 +72,7 @@ static void vnc_client_finish_update(rfbClient* client)
 	self->update_fb(self);
 
 	pixman_region_clear(&self->damage);
+	vnc_client_clear_av_frames(self);
 }
 
 static void vnc_client_got_cut_text(rfbClient* client, const char* text,
@@ -66,8 +85,56 @@ static void vnc_client_got_cut_text(rfbClient* client, const char* text,
 		self->cut_text(self, text, len);
 }
 
+static rfbBool vnc_client_handle_open_h264_rect(rfbClient* client,
+		rfbFramebufferUpdateRectHeader* rect_header)
+{
+	struct vnc_client* self = rfbClientGetClientData(client, NULL);
+	assert(self);
+
+	if (!self->open_h264)
+		self->open_h264 = open_h264_create(client);
+
+	if (!self->open_h264)
+		return false;
+
+	AVFrame* frame = open_h264_decode_rect(self->open_h264, rect_header);
+	if (!frame)
+		return false;
+
+	assert(self->n_av_frames < VNC_CLIENT_MAX_AV_FRAMES);
+
+	struct vnc_av_frame* f = calloc(1, sizeof(*f));
+	if (!f) {
+		av_frame_unref(frame);
+		return false;
+	}
+
+	f->frame = frame;
+	f->x = rect_header->r.x;
+	f->y = rect_header->r.y;
+	f->width = rect_header->r.w;
+	f->height = rect_header->r.h;
+
+	self->av_frames[self->n_av_frames++] = f;
+
+	self->current_rect_is_av_frame = true;
+	return true;
+}
+
+static void vnc_client_init_open_h264(void)
+{
+	static int encodings[] = { RFB_ENCODING_OPEN_H264, 0 };
+	static rfbClientProtocolExtension ext = {
+		.encodings = encodings,
+		.handleEncoding = vnc_client_handle_open_h264_rect,
+	};
+	rfbClientRegisterExtension(&ext);
+}
+
 struct vnc_client* vnc_client_create(void)
 {
+	vnc_client_init_open_h264();
+
 	struct vnc_client* self = calloc(1, sizeof(*self));
 	if (!self)
 		return NULL;
@@ -101,6 +168,8 @@ failure:
 
 void vnc_client_destroy(struct vnc_client* self)
 {
+	vnc_client_clear_av_frames(self);
+	open_h264_destroy(self->open_h264);
 	rfbClientCleanup(self->client);
 	free(self);
 }