2 files changed, 22 insertions, 16 deletions

diff --git a/ui/vnc-enc-tight.c b/ui/vnc-enc-tight.c
index e5cba0e5a7..b8581dd2e9 100644
--- a/ui/vnc-enc-tight.c
+++ b/ui/vnc-enc-tight.c
@@ -349,7 +349,7 @@ tight_detect_smooth_image(VncState *vs, int w, int h)
     tight_fill_palette##bpp(VncState *vs, int x, int y,                 \
                             int max, size_t count,                      \
                             uint32_t *bg, uint32_t *fg,                 \
-                            VncPalette **palette) {                     \
+                            VncPalette *palette) {                      \
         uint##bpp##_t *data;                                            \
         uint##bpp##_t c0, c1, ci;                                       \
         int i, n0, n1;                                                  \
@@ -396,23 +396,23 @@ tight_detect_smooth_image(VncState *vs, int w, int h)
             return 0;                                                   \
         }                                                               \
                                                                         \
-        *palette = palette_new(max, bpp);                               \
-        palette_put(*palette, c0);                                      \
-        palette_put(*palette, c1);                                      \
-        palette_put(*palette, ci);                                      \
+        palette_init(palette, max, bpp);                                \
+        palette_put(palette, c0);                                       \
+        palette_put(palette, c1);                                       \
+        palette_put(palette, ci);                                       \
                                                                         \
         for (i++; i < count; i++) {                                     \
             if (data[i] == ci) {                                        \
                 continue;                                               \
             } else {                                                    \
                 ci = data[i];                                           \
-                if (!palette_put(*palette, (uint32_t)ci)) {             \
+                if (!palette_put(palette, (uint32_t)ci)) {              \
                     return 0;                                           \
                 }                                                       \
             }                                                           \
         }                                                               \
                                                                         \
-        return palette_size(*palette);                                  \
+        return palette_size(palette);                                   \
     }
 
 DEFINE_FILL_PALETTE_FUNCTION(8)
@@ -421,7 +421,7 @@ DEFINE_FILL_PALETTE_FUNCTION(32)
 
 static int tight_fill_palette(VncState *vs, int x, int y,
                               size_t count, uint32_t *bg, uint32_t *fg,
-                              VncPalette **palette)
+                              VncPalette *palette)
 {
     int max;
 
@@ -1457,9 +1457,11 @@ static int send_sub_rect_jpeg(VncState *vs, int x, int y, int w, int h,
 }
 #endif
 
+static __thread VncPalette color_count_palette;
+
 static int send_sub_rect(VncState *vs, int x, int y, int w, int h)
 {
-    VncPalette *palette = NULL;
+    VncPalette *palette = &color_count_palette;
     uint32_t bg = 0, fg = 0;
     int colors;
     int ret = 0;
@@ -1488,7 +1490,7 @@ static int send_sub_rect(VncState *vs, int x, int y, int w, int h)
     }
 #endif
 
-    colors = tight_fill_palette(vs, x, y, w * h, &bg, &fg, &palette);
+    colors = tight_fill_palette(vs, x, y, w * h, &bg, &fg, palette);
 
 #ifdef CONFIG_VNC_JPEG
     if (allow_jpeg && vs->tight.quality != (uint8_t)-1) {
@@ -1501,7 +1503,6 @@ static int send_sub_rect(VncState *vs, int x, int y, int w, int h)
     ret = send_sub_rect_nojpeg(vs, x, y, w, h, bg, fg, colors, palette);
 #endif
 
-    palette_destroy(palette);
     return ret;
 }
 
diff --git a/ui/vnc.c b/ui/vnc.c
index 18c0b56c3a..e3f857cc90 100644
--- a/ui/vnc.c
+++ b/ui/vnc.c
@@ -1025,7 +1025,7 @@ static int find_and_clear_dirty_height(VncState *vs,
 static int vnc_update_client(VncState *vs, int has_dirty, bool sync)
 {
     vs->has_dirty += has_dirty;
-    if (vs->need_update && vs->ioc != NULL) {
+    if (vs->need_update && !vs->disconnecting) {
         VncDisplay *vd = vs->vd;
         VncJob *job;
         int y;
@@ -1436,8 +1436,9 @@ static void vnc_jobs_bh(void *opaque)
  * First function called whenever there is more data to be read from
  * the client socket. Will delegate actual work according to whether
  * SASL SSF layers are enabled (thus requiring decryption calls)
+ * Returns 0 on success, -1 if client disconnected
  */
-static void vnc_client_read(VncState *vs)
+static int vnc_client_read(VncState *vs)
 {
     ssize_t ret;
 
@@ -1450,8 +1451,9 @@ static void vnc_client_read(VncState *vs)
     if (!ret) {
         if (vs->disconnecting) {
             vnc_disconnect_finish(vs);
+            return -1;
         }
-        return;
+        return 0;
     }
 
     while (vs->read_handler && vs->input.offset >= vs->read_handler_expect) {
@@ -1461,7 +1463,7 @@ static void vnc_client_read(VncState *vs)
         ret = vs->read_handler(vs, vs->input.buffer, len);
         if (vs->disconnecting) {
             vnc_disconnect_finish(vs);
-            return;
+            return -1;
         }
 
         if (!ret) {
@@ -1470,6 +1472,7 @@ static void vnc_client_read(VncState *vs)
             vs->read_handler_expect = ret;
         }
     }
+    return 0;
 }
 
 gboolean vnc_client_io(QIOChannel *ioc G_GNUC_UNUSED,
@@ -1477,7 +1480,9 @@ gboolean vnc_client_io(QIOChannel *ioc G_GNUC_UNUSED,
 {
     VncState *vs = opaque;
     if (condition & G_IO_IN) {
-        vnc_client_read(vs);
+        if (vnc_client_read(vs) < 0) {
+            return TRUE;
+        }
     }
     if (condition & G_IO_OUT) {
         vnc_client_write(vs);