From 3647457f6d769fe3c9518078008ef56be58b70b1 Mon Sep 17 00:00:00 2001
From: Andri Yngvason <andri@yngvason.is>
Date: Sat, 5 Nov 2022 13:20:41 +0000
Subject: [PATCH] server: Properly handle stream write errors

This fixes leaked rcbufs when clients disconnect.
---
 src/server.c | 65 +++++++++++++++++++++++++++++-----------------------
 1 file changed, 36 insertions(+), 29 deletions(-)

diff --git a/src/server.c b/src/server.c
index 1af7892..423079f 100644
--- a/src/server.c
+++ b/src/server.c
@@ -567,10 +567,10 @@ static bool will_send_pts(const struct nvnc_client* client, uint64_t pts)
 	return pts != NVNC_NO_PTS && client_has_encoding(client, RFB_ENCODING_PTS);
 }
 
-static void send_pts_rect(struct nvnc_client* client, uint64_t pts)
+static int send_pts_rect(struct nvnc_client* client, uint64_t pts)
 {
 	if (!will_send_pts(client, pts))
-		return;
+		return 0;
 
 	uint8_t buf[sizeof(struct rfb_server_fb_rect) + 8] = { 0 };
 	struct rfb_server_fb_rect* head = (struct rfb_server_fb_rect*)buf;
@@ -578,7 +578,7 @@ static void send_pts_rect(struct nvnc_client* client, uint64_t pts)
 	uint64_t* msg_pts = (uint64_t*)&buf[sizeof(struct rfb_server_fb_rect)];
 	*msg_pts = nvnc__htonll(pts);
 
-	stream_write(client->net_stream, buf, sizeof(buf), NULL, NULL);
+	return stream_write(client->net_stream, buf, sizeof(buf), NULL, NULL);
 }
 
 static const char* encoding_to_string(enum rfb_encodings encoding)
@@ -1347,9 +1347,8 @@ void nvnc_close(struct nvnc* self)
 	free(self);
 }
 
-static void on_write_frame_done(void* userdata, enum stream_req_status status)
+static void complete_fb_update(struct nvnc_client* client)
 {
-	struct nvnc_client* client = userdata;
 	client->is_updating = false;
 	assert(client->current_fb);
 	nvnc_fb_release(client->current_fb);
@@ -1357,6 +1356,13 @@ static void on_write_frame_done(void* userdata, enum stream_req_status status)
 	client->current_fb = NULL;
 	process_fb_update_requests(client);
 	client_unref(client);
+	DTRACE_PROBE2(neatvnc, update_fb_done, client, pts);
+}
+
+static void on_write_frame_done(void* userdata, enum stream_req_status status)
+{
+	struct nvnc_client* client = userdata;
+	complete_fb_update(client);
 }
 
 static enum rfb_encodings choose_frame_encoding(struct nvnc_client* client,
@@ -1397,31 +1403,32 @@ static void finish_fb_update(struct nvnc_client* client, struct rcbuf* payload,
 {
 	client_ref(client);
 
-	if (client->net_stream->state != STREAM_STATE_CLOSED) {
-		DTRACE_PROBE2(neatvnc, send_fb_start, client, pts);
-		n_rects += will_send_pts(client, pts) ? 1 : 0;
-		struct rfb_server_fb_update_msg update_msg = {
-			.type = RFB_SERVER_TO_CLIENT_FRAMEBUFFER_UPDATE,
-			.n_rects = htons(n_rects),
-		};
-		stream_write(client->net_stream, &update_msg,
-				sizeof(update_msg), NULL, NULL);
-		send_pts_rect(client, pts);
-		rcbuf_ref(payload);
-		stream_send(client->net_stream, payload, on_write_frame_done,
-		            client);
-		DTRACE_PROBE2(neatvnc, send_fb_done, client, pts);
-	} else {
-		client->is_updating = false;
-		assert(client->current_fb);
-		nvnc_fb_release(client->current_fb);
-		nvnc_fb_unref(client->current_fb);
-		client->current_fb = NULL;
-		process_fb_update_requests(client);
-		client_unref(client);
-	}
+	if (client->net_stream->state == STREAM_STATE_CLOSED)
+		goto complete;
 
-	DTRACE_PROBE2(neatvnc, update_fb_done, client, pts);
+	DTRACE_PROBE2(neatvnc, send_fb_start, client, pts);
+	n_rects += will_send_pts(client, pts) ? 1 : 0;
+	struct rfb_server_fb_update_msg update_msg = {
+		.type = RFB_SERVER_TO_CLIENT_FRAMEBUFFER_UPDATE,
+		.n_rects = htons(n_rects),
+	};
+	if (stream_write(client->net_stream, &update_msg,
+			sizeof(update_msg), NULL, NULL) < 0)
+		goto complete;
+
+	if (send_pts_rect(client, pts) < 0)
+		goto complete;
+
+	rcbuf_ref(payload);
+	if (stream_send(client->net_stream, payload,
+				on_write_frame_done, client) < 0)
+		goto complete;
+
+	DTRACE_PROBE2(neatvnc, send_fb_done, client, pts);
+	return;
+
+complete:
+	complete_fb_update(client);
 }
 
 static void on_encode_frame_done(struct encoder* encoder, struct rcbuf* result,