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; }