From b725a08b47e3264366b90eb74ea0e823e1bfd9a4 Mon Sep 17 00:00:00 2001
From: Andri Yngvason <andri@yngvason.is>
Date: Thu, 14 Apr 2022 18:41:34 +0000
Subject: [PATCH] vnc: Add presentation timestamps

---
 include/vnc.h       |  1 +
 src/vnc.c           | 49 +++++++++++++++++++++++++++++++++++++++++++++
 src/vnc_encodings.c |  4 ++++
 3 files changed, 54 insertions(+)

diff --git a/include/vnc.h b/include/vnc.h
index 586dc51..5689d52 100644
--- a/include/vnc.h
+++ b/include/vnc.h
@@ -39,6 +39,7 @@ struct vnc_client {
 	bool current_rect_is_av_frame;
 	struct vnc_av_frame* av_frames[VNC_CLIENT_MAX_AV_FRAMES];
 	int n_av_frames;
+	uint64_t pts;
 
 	int (*alloc_fb)(struct vnc_client*);
 	void (*update_fb)(struct vnc_client*);
diff --git a/src/vnc.c b/src/vnc.c
index c4bfec5..cb534be 100644
--- a/src/vnc.c
+++ b/src/vnc.c
@@ -19,21 +19,35 @@
 #include <assert.h>
 #include <string.h>
 #include <stdbool.h>
+#include <limits.h>
 #include <pixman.h>
 #include <rfb/rfbclient.h>
 #include <libdrm/drm_fourcc.h>
 #include <libavutil/frame.h>
+#include <byteswap.h>
 
 #include "vnc.h"
 #include "open-h264.h"
 
 #define RFB_ENCODING_OPEN_H264 50
+#define RFB_ENCODING_PTS -1000
+
+#define NO_PTS UINT64_MAX
 
 extern const unsigned short code_map_linux_to_qnum[];
 extern const unsigned int code_map_linux_to_qnum_len;
 
 rfbBool vnc_client_set_format_and_encodings(rfbClient* client);
 
+static uint64_t vnc_client_htonll(uint64_t x)
+{
+#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
+	return bswap_64(x);
+#else
+	return x;
+#endif
+}
+
 static rfbBool vnc_client_alloc_fb(rfbClient* client)
 {
 	struct vnc_client* self = rfbClientGetClientData(client, NULL);
@@ -73,6 +87,7 @@ static void vnc_client_finish_update(rfbClient* client)
 
 	self->update_fb(self);
 
+	self->pts = NO_PTS;
 	pixman_region_clear(&self->damage);
 	vnc_client_clear_av_frames(self);
 }
@@ -90,6 +105,9 @@ static void vnc_client_got_cut_text(rfbClient* client, const char* text,
 static rfbBool vnc_client_handle_open_h264_rect(rfbClient* client,
 		rfbFramebufferUpdateRectHeader* rect_header)
 {
+	if ((int)rect_header->encoding != RFB_ENCODING_OPEN_H264)
+		return FALSE;
+
 	struct vnc_client* self = rfbClientGetClientData(client, NULL);
 	assert(self);
 
@@ -133,9 +151,38 @@ static void vnc_client_init_open_h264(void)
 	rfbClientRegisterExtension(&ext);
 }
 
+static rfbBool vnc_client_handle_pts_rect(rfbClient* client,
+		rfbFramebufferUpdateRectHeader* rect_header)
+{
+	if ((int)rect_header->encoding != RFB_ENCODING_PTS)
+		return FALSE;
+
+	struct vnc_client* self = rfbClientGetClientData(client, NULL);
+	assert(self);
+
+	uint64_t pts_msg = 0;
+	if (!ReadFromRFBServer(self->client, (char*)&pts_msg, sizeof(pts_msg)))
+		return FALSE;
+
+	self->pts = vnc_client_htonll(pts_msg);
+
+	return TRUE;
+}
+
+static void vnc_client_init_pts_ext(void)
+{
+	static int encodings[] = { RFB_ENCODING_PTS, 0 };
+	static rfbClientProtocolExtension ext = {
+		.encodings = encodings,
+		.handleEncoding = vnc_client_handle_pts_rect,
+	};
+	rfbClientRegisterExtension(&ext);
+}
+
 struct vnc_client* vnc_client_create(void)
 {
 	vnc_client_init_open_h264();
+	vnc_client_init_pts_ext();
 
 	struct vnc_client* self = calloc(1, sizeof(*self));
 	if (!self)
@@ -161,6 +208,8 @@ struct vnc_client* vnc_client_create(void)
 	client->FinishedFrameBufferUpdate = vnc_client_finish_update;
 	client->GotXCutText = vnc_client_got_cut_text;
 
+	self->pts = NO_PTS;
+
 	return self;
 
 failure:
diff --git a/src/vnc_encodings.c b/src/vnc_encodings.c
index d6e4853..9c1be76 100644
--- a/src/vnc_encodings.c
+++ b/src/vnc_encodings.c
@@ -173,6 +173,10 @@ rfbBool vnc_client_set_format_and_encodings(rfbClient* client)
 	if (se->nEncodings < MAX_ENCODINGS)
 		encs[se->nEncodings++] = rfbClientSwap32IfLE(rfbEncodingQemuExtendedKeyEvent);
 
+	/* pts */
+	if (se->nEncodings < MAX_ENCODINGS)
+		encs[se->nEncodings++] = rfbClientSwap32IfLE(-1000);
+
 	len = sz_rfbSetEncodingsMsg + se->nEncodings * 4;
 
 	se->nEncodings = rfbClientSwap16IfLE(se->nEncodings);