diff --git a/bitmap.h b/bitmap.h new file mode 100644 index 0000000..b15ecb5 --- /dev/null +++ b/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 +#include +#include +#include + +#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; +} diff --git a/likely.h b/likely.h new file mode 100644 index 0000000..74d8686 --- /dev/null +++ b/likely.h @@ -0,0 +1,5 @@ +#pragma once + +#define likely(x) __builtin_expect(!!(x), 1) +#define unlikely(x) __builtin_expect(!!(x), 0) + diff --git a/rfb-proto.h b/rfb-proto.h index d632883..8c1d9b7 100644 --- a/rfb-proto.h +++ b/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 = { diff --git a/server.c b/server.c index c999e6e..28ce5c1 100644 --- a/server.c +++ b/server.c @@ -1,4 +1,5 @@ #include "rfb-proto.h" +#include "util.h" #include #include @@ -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) { diff --git a/util.c b/util.c new file mode 100644 index 0000000..4093dea --- /dev/null +++ b/util.c @@ -0,0 +1,28 @@ +#include "util.h" + +#include +#include +#include + +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); +} diff --git a/util.h b/util.h new file mode 100644 index 0000000..84a61fb --- /dev/null +++ b/util.h @@ -0,0 +1,16 @@ +#ifndef _VNC_UTIL_H_ +#define _VNC_UTIL_H_ + +#include +#include + +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_ */ diff --git a/vec.c b/vec.c new file mode 100644 index 0000000..7c5e00b --- /dev/null +++ b/vec.c @@ -0,0 +1,82 @@ +#include "vec.h" +#include "likely.h" + +#include +#include + +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); +} diff --git a/vec.h b/vec.h new file mode 100644 index 0000000..3f88513 --- /dev/null +++ b/vec.h @@ -0,0 +1,41 @@ +#pragma once + +#include + +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_)) diff --git a/zrle.c b/zrle.c index 419af7d..d887baf 100644 --- a/zrle.c +++ b/zrle.c @@ -1,4 +1,7 @@ #include "rfb-proto.h" +#include "bitmap.h" +#include "util.h" +#include "vec.h" #include #include @@ -6,6 +9,7 @@ #include #include #include +#include #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, - 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) +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) { - /* 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 n_tiles = UDIV_UP(width, 64) * UDIV_UP(height, 64); - - uint8_t *frame = malloc(size); - if (!frame) - return -1; - + 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 }; - uint8_t *tile_buffer = malloc(size * bytes_per_cpixel); - if (!tile_buffer) + struct vec out; + uint8_t *in; + + in = malloc(chunk_size); + if (!in) goto failure; - int boff = 0; + /* 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;) { - if (!bitmap_is_set(tile_mask, i)) + 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) +{ + int rc = -1; + int n_tiles = UDIV_UP(width, 64) * UDIV_UP(height, 64); + + 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; + + struct rfb_server_fb_rect rect = { 0 }; + + for (int i = 0; i < n_tiles;) { + 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; }