From 8fc5c18b287d4c2da307172982107e30f63acb3d Mon Sep 17 00:00:00 2001
From: Andri Yngvason <andri@yngvason.is>
Date: Sat, 4 Apr 2020 14:00:03 +0000
Subject: [PATCH] Don't show partial frames to new clients

---
 include/common.h  |  1 +
 include/fb.h      |  3 +++
 include/neatvnc.h |  7 +++++++
 src/fb.c          | 12 ++++++++++++
 src/server.c      | 20 +++++++++++++++++---
 5 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/include/common.h b/include/common.h
index 6f0a9cf..c9a95fb 100644
--- a/include/common.h
+++ b/include/common.h
@@ -70,6 +70,7 @@ struct nvnc_client {
 	struct pixman_region16 damage;
 	int n_pending_requests;
 	bool is_updating;
+	bool needs_whole_frame;
 	nvnc_client_fn cleanup_fn;
 	z_stream z_stream;
 	struct tight_encoder tight_encoder;
diff --git a/include/fb.h b/include/fb.h
index e101236..adb1ff0 100644
--- a/include/fb.h
+++ b/include/fb.h
@@ -4,10 +4,13 @@
 #include <stdint.h>
 #include <stdatomic.h>
 
+#include "neatvnc.h"
+
 struct nvnc_fb {
 	int ref;
 	void* addr;
 	atomic_bool is_locked;
+	enum nvnc_fb_flags flags;
 	size_t size;
 	uint16_t width;
 	uint16_t height;
diff --git a/include/neatvnc.h b/include/neatvnc.h
index e0d2a30..27fb75a 100644
--- a/include/neatvnc.h
+++ b/include/neatvnc.h
@@ -32,6 +32,10 @@ enum nvnc_button_mask {
 	NVNC_SCROLL_DOWN = 1 << 4,
 };
 
+enum nvnc_fb_flags {
+	NVNC_FB_PARTIAL = 1 << 0, // The buffer contains only the damaged region
+};
+
 typedef void (*nvnc_key_fn)(struct nvnc_client*, uint32_t keysym,
                             bool is_pressed);
 typedef void (*nvnc_pointer_fn)(struct nvnc_client*, uint16_t x, uint16_t y,
@@ -76,6 +80,9 @@ void nvnc_fb_unref(struct nvnc_fb* fb);
 bool nvnc_fb_lock(struct nvnc_fb*);
 void nvnc_fb_unlock(struct nvnc_fb*);
 
+enum nvnc_fb_flags nvnc_fb_get_flags(const struct nvnc_fb*);
+void nvnc_fb_set_flags(struct nvnc_fb*, enum nvnc_fb_flags);
+
 void* nvnc_fb_get_addr(const struct nvnc_fb* fb);
 uint16_t nvnc_fb_get_width(const struct nvnc_fb* fb);
 uint16_t nvnc_fb_get_height(const struct nvnc_fb* fb);
diff --git a/src/fb.c b/src/fb.c
index 59e44e7..37c2d1e 100644
--- a/src/fb.c
+++ b/src/fb.c
@@ -53,6 +53,18 @@ void nvnc_fb_unlock(struct nvnc_fb* fb)
 	fb->is_locked = false;
 }
 
+EXPORT
+enum nvnc_fb_flags nvnc_fb_get_flags(const struct nvnc_fb* fb)
+{
+	return fb->flags;
+}
+
+EXPORT
+void nvnc_fb_set_flags(struct nvnc_fb* fb, enum nvnc_fb_flags flags)
+{
+	fb->flags = flags;
+}
+
 EXPORT
 void* nvnc_fb_get_addr(const struct nvnc_fb* fb)
 {
diff --git a/src/server.c b/src/server.c
index 90ed9f4..69fd400 100644
--- a/src/server.c
+++ b/src/server.c
@@ -474,13 +474,23 @@ static void process_fb_update_requests(struct nvnc_client* client)
 	if (client->is_updating || client->n_pending_requests == 0)
 		return;
 
-	if (!nvnc_fb_lock(client->server->frame))
+	struct nvnc_fb* fb = client->server->frame;
+
+	if (!nvnc_fb_lock(fb))
 		return;
 
-	client->is_updating = true;
+	if (client->needs_whole_frame && (fb->flags & NVNC_FB_PARTIAL))
+		goto abort;
 
+	client->is_updating = true;
 	if (schedule_client_update_fb(client) < 0)
-		nvnc_fb_unlock(client->server->frame);
+		goto abort;
+
+	return;
+
+abort:
+	nvnc_fb_unlock(fb);
+	client->is_updating = false;
 }
 
 static int on_client_fb_update_request(struct nvnc_client* client)
@@ -500,6 +510,9 @@ static int on_client_fb_update_request(struct nvnc_client* client)
 	int width = ntohs(msg->width);
 	int height = ntohs(msg->height);
 
+	if (!incremental)
+		client->needs_whole_frame = true;
+
 	client->n_pending_requests++;
 
 	/* Note: The region sent from the client is ignored for incremental
@@ -911,6 +924,7 @@ void on_client_update_fb_done(void* work)
 		stream_send(client->net_stream, payload, on_write_frame_done,
 		            client);
 		DTRACE_PROBE1(neatvnc, send_fb_done, client);
+		client->needs_whole_frame = false;
 	} else {
 		client->is_updating = false;
 		vec_destroy(frame);