From e69006fc483a4b9ab9f4cd3ac9f67aa3ddb6ecfb Mon Sep 17 00:00:00 2001
From: Andri Yngvason <andri@yngvason.is>
Date: Sun, 27 Sep 2020 17:17:22 +0000
Subject: [PATCH] tight: Hold fb and client refs while encoding

---
 include/tight.h |  4 ++--
 src/server.c    | 10 +++++++---
 src/tight.c     |  9 +++++++--
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/include/tight.h b/include/tight.h
index 9d06f47..7c11387 100644
--- a/include/tight.h
+++ b/include/tight.h
@@ -51,7 +51,7 @@ struct tight_encoder {
 
 	struct rfb_pixel_format dfmt;
 	struct rfb_pixel_format sfmt;
-	const struct nvnc_fb* fb;
+	struct nvnc_fb* fb;
 
 	uint32_t n_rects;
 	uint32_t n_jobs;
@@ -71,7 +71,7 @@ int tight_encoder_resize(struct tight_encoder* self, uint32_t width,
 
 int tight_encode_frame(struct tight_encoder* self,
 		const struct rfb_pixel_format* dfmt,
-		const struct nvnc_fb* src,
+		struct nvnc_fb* src,
 		const struct rfb_pixel_format* sfmt,
 		struct pixman_region16* damage,
 		enum tight_quality quality,
diff --git a/src/server.c b/src/server.c
index 26b7a3a..1072461 100644
--- a/src/server.c
+++ b/src/server.c
@@ -541,14 +541,17 @@ static void process_fb_update_requests(struct nvnc_client* client)
 	case RFB_ENCODING_ZRLE:
 		rc = schedule_client_update_fb(client, &damage);
 		break;
-	case RFB_ENCODING_TIGHT:;
-		enum tight_quality quality = client_get_tight_quality(client);
+	case RFB_ENCODING_TIGHT:
+		client_ref(client);
 
-		// TODO: Make sure we don't have use-after-free on client and fb
+		enum tight_quality quality = client_get_tight_quality(client);
 		rc = tight_encode_frame(&client->tight_encoder, &client->pixfmt,
 				fb, &server_fmt, &damage, quality,
 				on_tight_encode_frame_done, client);
 
+		if (rc < 0)
+			client_unref(client);
+
 		pixman_region_fini(&damage);
 		break;
 	default:
@@ -1186,6 +1189,7 @@ static void on_tight_encode_frame_done(struct vec* frame, void* userdata)
 {
 	struct nvnc_client* client = userdata;
 	finish_fb_update(client, frame);
+	client_unref(client);
 }
 
 static void on_client_update_fb_done(void* work)
diff --git a/src/tight.c b/src/tight.c
index c292dfc..692a4c2 100644
--- a/src/tight.c
+++ b/src/tight.c
@@ -402,8 +402,10 @@ static void on_tight_zs_work_done(void* obj)
 	struct tight_zs_worker_ctx* ctx = aml_get_userdata(obj);
 	struct tight_encoder* self = ctx->encoder;
 
-	if (--self->n_jobs == 0)
+	if (--self->n_jobs == 0) {
+		nvnc_fb_unref(self->fb);
 		schedule_tight_finish(self);
+	}
 }
 
 static int tight_schedule_zs_work(struct tight_encoder* self, int index)
@@ -479,7 +481,7 @@ static int schedule_tight_finish(struct tight_encoder* self)
 
 int tight_encode_frame(struct tight_encoder* self,
 		const struct rfb_pixel_format* dfmt,
-		const struct nvnc_fb* src,
+		struct nvnc_fb* src,
 		const struct rfb_pixel_format* sfmt,
 		struct pixman_region16* damage,
 		enum tight_quality quality,
@@ -503,7 +505,10 @@ int tight_encode_frame(struct tight_encoder* self,
 
 	encode_rect_count(&self->dst, self->n_rects);
 
+	nvnc_fb_ref(self->fb);
+
 	if (tight_schedule_encoding_jobs(self) < 0) {
+		nvnc_fb_unref(self->fb);
 		vec_destroy(&self->dst);
 		return -1;
 	}