Do more zrle work

bitmap.h

@@ -0,0 +1,90 @@
+/*
+ * Copyright (c) 2019 Andri Yngvason
+ * 
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ * 
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ * 
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+#pragma once
+
+#include <stdint.h>
+#include <stdlib.h>
+#include <unistd.h>
+#include <stdbool.h>
+
+#define UDIV_UP(a, b) (((a) + (b) - 1) / (b))
+
+struct bitmap {
+	size_t n_elem;
+	uint64_t data[0];
+};
+
+static inline struct bitmap *bitmap_alloc(size_t bitlen)
+{
+	size_t n_elem = UDIV_UP(bitlen, 64);
+
+	struct bitmap *self = calloc(1, sizeof(*self) + n_elem * sizeof(*self->data));
+	if (!self)
+		return NULL;
+
+	self->n_elem = n_elem;
+
+	return self;
+}
+
+static inline void bitmap_clear_all(struct bitmap *self)
+{
+	for (size_t i = 0; i < self->n_elem; ++i)
+		self->data[i] = 0;
+}
+
+static inline int bitmap_is_empty(const struct bitmap *self)
+{
+	for (size_t i = 0; i < self->n_elem; ++i)
+		if (self->data[i])
+			return 0;
+
+	return 1;
+}
+
+static inline int bitmap_is_set(const struct bitmap *self, int index)
+{
+	return !!(self->data[index / 64] & (1ULL << (index % 64)));
+}
+
+static inline void bitmap_clear(struct bitmap* self, int index)
+{
+	self->data[index / 64] &= ~(1ULL << (index % 64));
+}
+
+static inline void bitmap_set_cond(struct bitmap* self, int index, bool cond)
+{
+	self->data[index / 64] |= ((uint64_t)cond) << (index % 64);
+}
+
+static inline void bitmap_set(struct bitmap* self, int index)
+{
+	bitmap_set_cond(self, index, true);
+}
+
+static inline int bitmap_runlength(const struct bitmap *self, int start)
+{
+        int r = 0;
+        while (bitmap_is_set(self, start + r)) ++r;
+        return r;
+}

likely.h

@@ -0,0 +1,5 @@
+#pragma once
+
+#define likely(x)       __builtin_expect(!!(x), 1)
+#define unlikely(x)     __builtin_expect(!!(x), 0)
+

rfb-proto.h

@@ -44,6 +44,13 @@ enum rfb_encodings {
 	RFB_ENCODING_DESKTOPSIZE = -223,
 };
 
+enum rfb_server_to_client_msg_type {
+        RFB_SERVER_TO_CLIENT_FRAMEBUFFER_UPDATE = 0,
+        RFB_SERVER_TO_CLIENT_SET_COLOUR_MAP_ENTRIES = 1,
+        RFB_SERVER_TO_CLIENT_BELL = 2,
+        RFB_SERVER_TO_CLIENT_SERVER_CUT_TEXT = 3,
+};
+
 struct rfb_security_types_msg {
         uint8_t n;
         uint8_t types[1];
@@ -113,6 +120,21 @@ struct rfb_client_cut_text_msg {
         char test[0];
 } RFB_PACKED;
 
+struct rfb_server_fb_rect {
+        uint16_t x;
+        uint16_t y;
+        uint16_t width;
+        uint16_t height;
+        int32_t encoding;
+} RFB_PACKED;
+
+struct rfb_server_fb_update_msg {
+        uint8_t type;
+        uint8_t padding;
+        uint16_t n_rects;
+        struct rfb_server_fb_rect rects[0];
+} RFB_PACKED;
+
 static inline int rfb_send_security_types(void *client)
 {
         struct rfb_security_types_msg payload = {

server.c

@@ -1,4 +1,5 @@
 #include "rfb-proto.h"
+#include "util.h"
 
 #include <stdlib.h>
 #include <unistd.h>
@@ -43,12 +44,6 @@ struct vnc_client {
 	uint8_t msg_buffer[MSG_BUFFER_SIZE];
 };
 
-struct vnc_write_request {
-	uv_write_t request;
-	uv_write_cb on_done;
-	uv_buf_t buffer;
-};
-
 LIST_HEAD(vnc_client_list, vnc_client);
 
 struct vnc_display {
@@ -77,29 +72,6 @@ static const char* fourcc_to_string(uint32_t fourcc)
 	return buffer;
 }
 
-static void on_write_req_done(uv_write_t *req, int status)
-{
-	struct vnc_write_request *self = (struct vnc_write_request*)req;
-	if (self->on_done)
-		self->on_done(req, status);
-	free(self);
-}
-
-static int vnc__write(uv_stream_t *stream, const void *payload, size_t size,
-		      uv_write_cb on_done)
-{
-	struct vnc_write_request *req = calloc(1, sizeof(*req));
-	if (!req)
-		return -1;
-
-	req->buffer.base = (char*)payload;
-	req->buffer.len = size;
-	req->on_done = on_done;
-
-	return uv_write(&req->request, stream, &req->buffer, 1,
-			on_write_req_done);
-}
-
 static void allocate_read_buffer(uv_handle_t *handle, size_t suggested_size,
 				 uv_buf_t *buf)
 {

util.c

@@ -0,0 +1,28 @@
+#include "util.h"
+
+#include <uv.h>
+#include <stdlib.h>
+#include <unistd.h>
+
+static void on_write_req_done(uv_write_t *req, int status)
+{
+	struct vnc_write_request *self = (struct vnc_write_request*)req;
+	if (self->on_done)
+		self->on_done(req, status);
+	free(self);
+}
+
+int vnc__write(uv_stream_t *stream, const void *payload, size_t size,
+	       uv_write_cb on_done)
+{
+	struct vnc_write_request *req = calloc(1, sizeof(*req));
+	if (!req)
+		return -1;
+
+	req->buffer.base = (char*)payload;
+	req->buffer.len = size;
+	req->on_done = on_done;
+
+	return uv_write(&req->request, stream, &req->buffer, 1,
+			on_write_req_done);
+}

util.h

@@ -0,0 +1,16 @@
+#ifndef _VNC_UTIL_H_
+#define _VNC_UTIL_H_
+
+#include <uv.h>
+#include <unistd.h>
+
+struct vnc_write_request {
+	uv_write_t request;
+	uv_write_cb on_done;
+	uv_buf_t buffer;
+};
+
+int vnc__write(uv_stream_t *stream, const void *payload, size_t size,
+               uv_write_cb on_done);
+
+#endif /* _VNC_UTIL_H_ */

vec.c

@@ -0,0 +1,82 @@
+#include "vec.h"
+#include "likely.h"
+
+#include <stdlib.h>
+#include <string.h>
+
+int vec_init(struct vec* vec, size_t cap)
+{
+	memset(vec, 0, sizeof(*vec));
+	return vec_reserve(vec, cap);
+}
+
+void vec_destroy(struct vec* vec)
+{
+	free(vec->data);
+}
+
+int vec_reserve(struct vec* vec, size_t size)
+{
+	if (likely(size <= vec->cap))
+		return 0;
+
+	void* data = realloc(vec->data, size);
+	if (unlikely(!data))
+		return -1;
+
+	vec->cap = size;
+	vec->data = data;
+
+	return 0;
+}
+
+static int vec__grow(struct vec* vec, size_t size)
+{
+	if (likely(vec->len + size < vec->cap))
+		return 0;
+	
+	return vec_reserve(vec, 2 * (vec->len + size));
+}
+
+int vec_assign(struct vec* vec, const void* data, size_t size)
+{
+	vec->len = 0;
+
+	if (unlikely(vec_reserve(vec, size) < 0))
+		return -1;
+
+	vec->len = size;
+	memcpy(vec->data, data, size);
+
+	return 0;
+}
+
+int vec_append(struct vec* vec, const void* data, size_t size)
+{
+	if (unlikely(vec__grow(vec, size) < 0))
+		return -1;
+
+	char* p = vec->data;
+	memcpy(&p[vec->len], data, size);
+	vec->len += size;
+
+	return 0;
+}
+
+void* vec_append_zero(struct vec* vec, size_t size)
+{
+	if (unlikely(vec__grow(vec, size) < 0))
+		return NULL;
+
+	char* p = vec->data;
+	void* r = &p[vec->len];
+	memset(r, 0, size);
+	vec->len += size;
+
+	return r;
+}
+
+void vec_bzero(struct vec* vec)
+{
+	memset(vec->data, 0, vec->cap);
+}

vec.h

@@ -0,0 +1,41 @@
+#pragma once
+
+#include <stddef.h>
+
+struct vec {
+	void* data;
+	size_t len;
+	size_t cap;
+};
+
+static inline void vec_clear(struct vec* vec)
+{
+	vec->len = 0;
+}
+
+int vec_init(struct vec* vec, size_t cap);
+void vec_destroy(struct vec* vec);
+
+int vec_reserve(struct vec* vec, size_t size);
+
+void vec_bzero(struct vec* vec);
+
+int vec_assign(struct vec* vec, const void* data, size_t size);
+int vec_append(struct vec* vec, const void* data, size_t size);
+void* vec_append_zero(struct vec* vec, size_t size);
+
+
+#define vec_for(elem, vec) \
+	for (elem = (vec)->data; \
+	     ((ptrdiff_t)elem - (ptrdiff_t)(vec)->data) < (ptrdiff_t)(vec)->len; \
+	     ++elem)
+
+#define vec_for_tail(elem, vec) \
+	for (elem = (vec)->data, ++elem; \
+	     ((ptrdiff_t)elem - (ptrdiff_t)(vec)->data) < (ptrdiff_t)(vec)->len; \
+	     ++elem)
+
+#define vec_for_ptr(elem, vec) \
+	__typeof__(elem)* ptr_; \
+	vec_for(ptr_, vec) \
+		if ((elem = *ptr_))

zrle.c

@@ -1,4 +1,7 @@
 #include "rfb-proto.h"
+#include "bitmap.h"
+#include "util.h"
+#include "vec.h"
 
 #include <stdint.h>
 #include <unistd.h>
@@ -6,6 +9,7 @@
 #include <endian.h>
 #include <assert.h>
 #include <zlib.h>
+#include <uv.h>
 
 #define POPCOUNT(x) __builtin_popcount(x)
 
@@ -152,43 +156,146 @@ void zrle_encode_tile(uint8_t *dst, const struct rfb_pixel_format *dst_fmt,
 				  src_fmt, bytes_per_cpixel, width);
 }
 
-ssize_t zrle_encode_frame(uint8_t **result,
+int zrle_encode_adjacent_tiles(uv_stream_t *stream,
+			       const struct rfb_pixel_format *dst_fmt,
+			       uint8_t *src,
+			       const struct rfb_pixel_format *src_fmt,
+			       int n, int x, int y, int width, int height)
+{
+	int r = -1;
+	int zr = Z_STREAM_ERROR;
+	int bytes_per_cpixel = dst_fmt->depth / 8;
+	int chunk_size = bytes_per_cpixel * 64 * 64;
+	z_stream zs = { 0 };
+
+	struct vec out;
+	uint8_t *in;
+
+	in = malloc(chunk_size);
+	if (!in)
+		goto failure;
+
+	/* The output is expected to be around half the size of the input */
+	if (vec_init(&out, chunk_size * n / 2) < 0)
+		goto failure;
+
+	r = deflateInit(&zs, Z_DEFAULT_COMPRESSION);
+	if (r != Z_OK)
+		goto failure;
+
+	/* Reserve space for size */
+	vec_append_zero(&out, 4);
+
+	for (int i = 0; i < n; ++i) {
+		zrle_encode_tile(in, dst_fmt,
+				 ((uint32_t*)src) + x + y * width,
+				 src_fmt, width, width, height);
+
+		zs.next_in = in;
+		zs.avail_in = chunk_size;
+
+		int flush = (i == n - 1) ? Z_FINISH : Z_NO_FLUSH;
+
+		do {
+			zs.next_out = ((Bytef*)out.data) + out.len;
+
+			r = vec_reserve(&out, out.len + chunk_size);
+			if (r < 0)
+				goto failure;
+
+			zs.avail_out = out.cap - out.len;
+
+			zr = deflate(&zs, flush);
+			assert(zr != Z_STREAM_ERROR);
+
+			int have = chunk_size - zs.avail_out;
+			out.len += have;
+		} while (zs.avail_out == 0);
+
+		assert(zs.avail_in == 0);
+	}
+
+	assert(zr == Z_STREAM_END);
+
+	deflateEnd(&zs);
+
+	uint32_t *out_size = out.data;
+	*out_size = htonl(out.len);
+
+	r = vnc__write(stream, out.data, out.len, NULL);
+failure:
+	vec_destroy(&out);
+	free(in);
+	return r;
+#undef CHUNK
+}
+
+int zrle_count_contiguous_regions(struct bitmap *tile_mask, int n_tiles)
+{
+	int r = 0;
+
+	for (int i = 0; i < n_tiles;) {
+		int rl = bitmap_runlength(tile_mask, i);
+		if (rl == 0) {
+			++i;
+		} else {
+			++r;
+			i += rl;
+		}
+	}
+
+	return r;
+}
+
+int zrle_encode_frame(uv_stream_t *stream,
 		      const struct rfb_pixel_format *dst_fmt,
 		      uint8_t *src,
 		      const struct rfb_pixel_format *src_fmt,
 		      int width, int height,
 		      struct bitmap *tile_mask)
 {
-	/* Expect the compressed data to be half the size of the updated regions
-	 */
-	size_t actual_size, size = 64 * bitmap_popcount(tile_mask) / 2;
+	int rc = -1;
 	int n_tiles = UDIV_UP(width, 64) * UDIV_UP(height, 64);
 
-	uint8_t *frame = malloc(size);
-	if (!frame)
+	int n_regions = zrle_count_contiguous_regions(tile_mask, n_tiles);
+	assert(n_regions < UINT16_MAX);
+
+	struct rfb_server_fb_update_msg head = {
+		.type = RFB_SERVER_TO_CLIENT_FRAMEBUFFER_UPDATE,
+		.n_rects = htons(n_regions),
+	};
+
+	rc = vnc__write(stream, &head, sizeof(head), NULL);
+	if (rc < 0)
 		return -1;
 
-	int bytes_per_cpixel = dst_fmt->depth / 8;
-
-	uint8_t *tile_buffer = malloc(size * bytes_per_cpixel);
-	if (!tile_buffer)
-		goto failure;
-
-	int boff = 0;
+	struct rfb_server_fb_rect rect = { 0 };
 
 	for (int i = 0; i < n_tiles;) {
-		if (!bitmap_is_set(tile_mask, i))
+		if (!bitmap_is_set(tile_mask, i)) {
+			i += 1;
 			continue;
+		}
 
 		int x = (i * 64) % UDIV_UP(width, 64);
 		int y = (i * 64) / UDIV_UP(width, 64);
 
+		int adjacent = bitmap_runlength(tile_mask, i);
+		assert(adjacent <= n_tiles - i);
+
+		rect.encoding = htonl(RFB_ENCODING_ZRLE);
+		rect.x = x;
+		rect.y = y;
+		rect.width = 0; /* TODO */
+		rect.height = 0; /* TODO */
+
+		rc = zrle_encode_adjacent_tiles(stream, dst_fmt, src, src_fmt,
+						adjacent, x, y, width, height);
+		if (rc < 0)
+			return rc;
+
+		i += adjacent;
 	}
 
-	*result = frame;
-	return actual_size;
-
-failure:
-	free(frame);
-	return -1;
+	return 0;
 }