Implement Apple's Diffie-Hellman based security type 30

include/common.h

@@ -43,6 +43,9 @@ enum nvnc_client_state {
 	VNC_CLIENT_STATE_WAITING_FOR_VENCRYPT_VERSION,
 	VNC_CLIENT_STATE_WAITING_FOR_VENCRYPT_SUBTYPE,
 	VNC_CLIENT_STATE_WAITING_FOR_VENCRYPT_PLAIN_AUTH,
+#endif
+#ifdef HAVE_CRYPTO
+	VNC_CLIENT_STATE_WAITING_FOR_APPLE_DH_RESPONSE,
 #endif
 	VNC_CLIENT_STATE_WAITING_FOR_INIT,
 	VNC_CLIENT_STATE_READY,
@@ -53,6 +56,7 @@ struct stream;
 struct aml_handler;
 struct aml_idle;
 struct nvnc_display;
+struct crypto_key;
 
 struct nvnc_common {
 	void* userdata;
@@ -93,6 +97,10 @@ struct nvnc_client {
 	struct encoder* encoder;
 	uint32_t cursor_seq;
 	int quality;
+
+#ifdef HAVE_CRYPTO
+	struct crypto_key* apple_dh_secret;
+#endif
 };
 
 LIST_HEAD(nvnc_client_list, nvnc_client);

include/crypto.h

@@ -0,0 +1,61 @@
+#pragma once
+
+#include <stdint.h>
+#include <stdbool.h>
+#include <unistd.h>
+
+struct crypto_key;
+struct crypto_cipher;
+struct crypto_hash;
+
+enum crypto_cipher_type {
+	CRYPTO_CIPHER_INVALID = 0,
+	CRYPTO_CIPHER_AES128_ECB,
+};
+
+enum crypto_hash_type {
+	CRYPTO_HASH_INVALID = 0,
+	CRYPTO_HASH_MD5,
+};
+
+void crypto_dump_base64(const char* msg, const uint8_t* bytes, size_t len);
+
+void crypto_cleanup(void);
+
+// Key generation
+struct crypto_key* crypto_key_new(int g, const uint8_t *p, uint32_t p_len,
+		const uint8_t* q, uint32_t q_len);
+void crypto_key_del(struct crypto_key* key);
+
+int crypto_key_g(const struct crypto_key* key);
+uint32_t crypto_key_p(const struct crypto_key* key, uint8_t* dst,
+		uint32_t dst_size);
+uint32_t crypto_key_q(const struct crypto_key* key, uint8_t* dst,
+		uint32_t dst_size);
+
+struct crypto_key* crypto_keygen(void);
+
+// Diffie-Hellman
+struct crypto_key* crypto_derive_public_key(const struct crypto_key* priv);
+struct crypto_key* crypto_derive_shared_secret(
+		const struct crypto_key* own_secret,
+		const struct crypto_key* remote_public_key);
+
+// Ciphers
+struct crypto_cipher* crypto_cipher_new(const uint8_t* key,
+		enum crypto_cipher_type type);
+void crypto_cipher_del(struct crypto_cipher* self);
+
+bool crypto_cipher_encrypt(struct crypto_cipher* self, uint8_t* dst,
+		const uint8_t* src, size_t len);
+bool crypto_cipher_decrypt(struct crypto_cipher* self, uint8_t* dst,
+		const uint8_t* src, size_t len);
+
+// Hashing
+struct crypto_hash* crypto_hash_new(enum crypto_hash_type type);
+void crypto_hash_del(struct crypto_hash* self);
+
+void crypto_hash_append(struct crypto_hash* self, const uint8_t* src,
+		size_t len);
+void crypto_hash_digest(struct crypto_hash* self, uint8_t* dst,
+		size_t len);

include/rfb-proto.h

@@ -31,6 +31,7 @@ enum rfb_security_type {
 	RFB_SECURITY_TYPE_VNC_AUTH = 2,
 	RFB_SECURITY_TYPE_TIGHT = 16,
 	RFB_SECURITY_TYPE_VENCRYPT = 19,
+	RFB_SECURITY_TYPE_APPLE_DH = 30,
 };
 
 enum rfb_security_handshake_result {
@@ -108,7 +109,7 @@ enum rfb_resize_status {
 
 struct rfb_security_types_msg {
 	uint8_t n;
-	uint8_t types[1];
+	uint8_t types[0];
 } RFB_PACKED;
 
 struct rfb_error_reason {
@@ -237,3 +238,14 @@ struct rfb_ntp_msg {
 	uint8_t padding[3];
 	uint32_t t0, t1, t2, t3;
 } RFB_PACKED;
+
+struct rfb_apple_dh_server_msg {
+	uint16_t generator;
+	uint16_t key_size;
+	uint8_t modulus_and_key[0];
+} RFB_PACKED;
+
+struct rfb_apple_dh_client_msg {
+	uint8_t encrypted_credentials[128];
+	uint8_t public_key[0];
+} RFB_PACKED;

include/stream.h

@@ -19,6 +19,7 @@
 #include "config.h"
 #include "sys/queue.h"
 #include "rcbuf.h"
+#include "vec.h"
 
 #include <stdint.h>
 #include <stdbool.h>
@@ -45,6 +46,7 @@ enum stream_event {
 };
 
 struct stream;
+struct crypto_cipher;
 
 typedef void (*stream_event_fn)(struct stream*, enum stream_event);
 typedef void (*stream_req_fn)(void*, enum stream_req_status);
@@ -68,6 +70,7 @@ struct stream_impl {
 			stream_req_fn on_done, void* userdata);
 	int (*send_first)(struct stream*, struct rcbuf* payload);
 	void (*exec_and_send)(struct stream*, stream_exec_fn, void* userdata);
+	int (*install_cipher)(struct stream*, struct crypto_cipher*);
 };
 
 struct stream {
@@ -86,6 +89,9 @@ struct stream {
 	uint32_t bytes_received;
 
 	bool cork;
+
+	struct crypto_cipher* cipher;
+	struct vec tmp_buf;
 };
 
 #ifdef ENABLE_WEBSOCKET
@@ -108,3 +114,5 @@ void stream_exec_and_send(struct stream* self, stream_exec_fn, void* userdata);
 #ifdef ENABLE_TLS
 int stream_upgrade_to_tls(struct stream* self, void* context);
 #endif
+
+int stream_install_cipher(struct stream* self, struct crypto_cipher* cipher);

meson.build

@@ -51,6 +51,7 @@ pixman = dependency('pixman-1')
 libturbojpeg = dependency('libturbojpeg', required: get_option('jpeg'))
 gnutls = dependency('gnutls', required: get_option('tls'))
 nettle = dependency('nettle', required: get_option('nettle'))
+gmp = dependency('gmp', required: get_option('nettle'))
 zlib = dependency('zlib')
 gbm = dependency('gbm', required: get_option('gbm'))
 libdrm = dependency('libdrm', required: get_option('h264'))
@@ -118,9 +119,11 @@ if gnutls.found()
 	config.set('ENABLE_TLS', true)
 endif
 
-if nettle.found()
+if nettle.found() and gmp.found()
-	dependencies += nettle
+	dependencies += [ nettle, gmp ]
 	enable_websocket = true
+	config.set('HAVE_CRYPTO', true)
+	sources += 'src/crypto-nettle.c'
 endif
 
 if host_system == 'linux' and get_option('systemtap') and cc.has_header('sys/sdt.h')

src/crypto-nettle.c

@@ -0,0 +1,323 @@
+#include "crypto.h"
+#include "neatvnc.h"
+
+#include <gmp.h>
+#include <nettle/base64.h>
+#include <nettle/base16.h>
+#include <nettle/aes.h>
+#include <nettle/md5.h>
+
+#include <stdint.h>
+#include <stdlib.h>
+#include <stdbool.h>
+
+// TODO: This is linux specific
+#include <sys/random.h>
+
+struct crypto_key {
+	int g;
+	mpz_t p;
+	mpz_t q;
+};
+
+struct crypto_cipher {
+	union {
+		struct aes128_ctx aes128_ecb;
+	} enc_ctx;
+
+	union {
+		struct aes128_ctx aes128_ecb;
+	} dec_ctx;
+
+	bool (*encrypt)(struct crypto_cipher*, uint8_t* dst, const uint8_t* src,
+			size_t len);
+	bool (*decrypt)(struct crypto_cipher*, uint8_t* dst, const uint8_t* src,
+			size_t len);
+};
+
+struct crypto_hash {
+	union {
+		struct md5_ctx md5;
+	} ctx;
+
+	void (*update)(void* ctx, size_t len, const uint8_t* src);
+	void (*digest)(void* ctx, size_t len, uint8_t* dst);
+};
+
+void crypto_dump_base64(const char* msg, const uint8_t* bytes, size_t len)
+{
+	struct base64_encode_ctx ctx = {};
+	size_t buflen = BASE64_ENCODE_LENGTH(len);
+	char* buffer = malloc(buflen + BASE64_ENCODE_FINAL_LENGTH);
+	assert(buffer);
+	nettle_base64_encode_init(&ctx);
+	nettle_base64_encode_update(&ctx, buffer, len, bytes);
+	nettle_base64_encode_final(&ctx, buffer + buflen);
+
+	nvnc_log(NVNC_LOG_DEBUG, "%s: %s", msg, buffer);
+	free(buffer);
+}
+
+struct crypto_key *crypto_key_new(int g, const uint8_t* p, uint32_t p_len,
+		const uint8_t* q, uint32_t q_len)
+{
+	struct crypto_key* self = calloc(1, sizeof(*self));
+	if (!self)
+		return NULL;
+
+	self->g = g;
+
+	int order = 1;
+	int unit_size = 1;
+	int endian = 1;
+	int skip_bits = 0;
+
+	mpz_init(self->p);
+	mpz_import(self->p, p_len, order, unit_size, endian, skip_bits, p);
+
+	mpz_init(self->q);
+	mpz_import(self->q, q_len, order, unit_size, endian, skip_bits, q);
+
+	return self;
+}
+
+void crypto_key_del(struct crypto_key* key)
+{
+	mpz_clear(key->q);
+	mpz_clear(key->p);
+	free(key);
+}
+
+int crypto_key_g(const struct crypto_key* key)
+{
+	return key->g;
+}
+
+uint32_t crypto_key_p(const struct crypto_key* key, uint8_t* dst,
+		uint32_t dst_size)
+{
+	int order = 1; // msb first
+	int unit_size = 1; // byte
+	int endian = 1; // msb first
+	int skip_bits = 0;
+
+	size_t len = 0;
+	mpz_export(dst, &len, order, unit_size, endian, skip_bits, key->p);
+
+	return len;
+}
+
+uint32_t crypto_key_q(const struct crypto_key* key, uint8_t* dst,
+		uint32_t dst_size)
+{
+	int order = 1; // msb first
+	int unit_size = 1; // byte
+	int endian = 1; // msb first
+	int skip_bits = 0;
+
+	size_t len = 0;
+	mpz_export(dst, &len, order, unit_size, endian, skip_bits, key->q);
+
+	return len;
+}
+
+static void initialise_p(mpz_t p)
+{
+	// RFC 3526, section 3
+	static const char s[] =
+		"FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1"
+		"29024E088A67CC74020BBEA63B139B22514A08798E3404DD"
+		"EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245"
+		"E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED"
+		"EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3D"
+		"C2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F"
+		"83655D23DCA3AD961C62F356208552BB9ED529077096966D"
+		"670C354E4ABC9804F1746C08CA18217C32905E462E36CE3B"
+		"E39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9"
+		"DE2BCBF6955817183995497CEA956AE515D2261898FA0510"
+		"15728E5A8AACAA68FFFFFFFFFFFFFFFF";
+
+	char buf[256];
+	size_t len = 0;
+	struct base16_decode_ctx ctx;
+	nettle_base16_decode_init(&ctx);
+	nettle_base16_decode_update(&ctx, &len, (uint8_t*)buf, sizeof(s) - 1, s);
+	nettle_base16_decode_final(&ctx);
+	assert(len == sizeof(buf));
+
+	int order = 1;
+	int unit_size = 1;
+	int endian = 1;
+	int skip_bits = 0;
+
+	mpz_import(p, sizeof(buf), order, unit_size, endian, skip_bits, buf);
+}
+
+static void generate_random(mpz_t n)
+{
+	uint8_t buf[256];
+	getrandom(buf, sizeof(buf), 0);
+
+	int order = 1;
+	int unit_size = 1;
+	int endian = 1;
+	int skip_bits = 0;
+
+	mpz_import(n, sizeof(buf), order, unit_size, endian, skip_bits, buf);
+}
+
+struct crypto_key* crypto_keygen(void)
+{
+	struct crypto_key* self = calloc(1, sizeof(*self));
+	if (!self)
+		return NULL;
+
+	self->g = 2;
+
+	mpz_init(self->p);
+	initialise_p(self->p);
+
+	mpz_init(self->q);
+	generate_random(self->q);
+
+	return self;
+}
+
+struct crypto_key* crypto_derive_public_key(const struct crypto_key* priv)
+{
+	struct crypto_key* pub = calloc(1, sizeof(*pub));
+	if (!pub)
+		return NULL;
+
+	pub->g = priv->g;
+	mpz_set(pub->p, priv->p);
+	mpz_init(pub->q);
+
+	mpz_t g;
+	mpz_init(g);
+	mpz_set_ui(g, priv->g);
+
+	mpz_powm_sec(pub->q, g, priv->q, priv->p);
+	mpz_clear(g);
+
+	return pub;
+}
+
+struct crypto_key* crypto_derive_shared_secret(
+		const struct crypto_key* own_secret,
+		const struct crypto_key* remote_public_key)
+{
+	if (own_secret->g != remote_public_key->g) {
+		return NULL;
+	}
+
+	if (mpz_cmp(own_secret->p, remote_public_key->p) != 0) {
+		return NULL;
+	}
+
+	struct crypto_key* shared = calloc(1, sizeof(*shared));
+	if (!shared)
+		return NULL;
+
+	shared->g = own_secret->g;
+	mpz_set(shared->p, own_secret->p);
+
+	mpz_t g;
+	mpz_init(g);
+	mpz_set_ui(g, own_secret->g);
+
+	mpz_powm_sec(shared->q, remote_public_key->q, own_secret->q,
+			own_secret->p);
+	mpz_clear(g);
+
+	return shared;
+}
+
+static bool crypto_cipher_aes128_ecb_encrypt(struct crypto_cipher* self,
+		uint8_t* dst, const uint8_t* src, size_t len)
+{
+	aes128_encrypt(&self->enc_ctx.aes128_ecb, len, dst, src);
+	return true;
+}
+
+static bool crypto_cipher_aes128_ecb_decrypt(struct crypto_cipher* self,
+		uint8_t* dst, const uint8_t* src, size_t len)
+{
+	aes128_decrypt(&self->dec_ctx.aes128_ecb, len, dst, src);
+	return true;
+}
+
+static struct crypto_cipher* crypto_cipher_new_aes128_ecb(const uint8_t* key)
+{
+	struct crypto_cipher* self = calloc(1, sizeof(*self));
+	if (!self)
+		return NULL;
+
+	aes128_set_encrypt_key(&self->enc_ctx.aes128_ecb, key);
+	aes128_invert_key(&self->dec_ctx.aes128_ecb, &self->enc_ctx.aes128_ecb);
+
+	self->encrypt = crypto_cipher_aes128_ecb_encrypt;
+	self->decrypt = crypto_cipher_aes128_ecb_decrypt;
+
+	return self;
+}
+
+struct crypto_cipher* crypto_cipher_new(const uint8_t* key,
+		enum crypto_cipher_type type)
+{
+	switch (type) {
+	case CRYPTO_CIPHER_AES128_ECB:
+		return crypto_cipher_new_aes128_ecb(key);
+	case CRYPTO_CIPHER_INVALID:
+		break;
+	}
+
+	nvnc_log(NVNC_LOG_PANIC, "Invalid type: %d", type);
+	return NULL;
+}
+
+void crypto_cipher_del(struct crypto_cipher* self)
+{
+	free(self);
+}
+
+bool crypto_cipher_encrypt(struct crypto_cipher* self, uint8_t* dst,
+		const uint8_t* src, size_t len)
+{
+	return self->encrypt(self, dst, src, len);
+}
+
+bool crypto_cipher_decrypt(struct crypto_cipher* self, uint8_t* dst,
+		const uint8_t* src, size_t len)
+{
+	return self->decrypt(self, dst, src, len);
+}
+
+struct crypto_hash* crypto_hash_new(enum crypto_hash_type type)
+{
+	struct crypto_hash* self = calloc(1, sizeof(*self));
+	if (!self)
+		return NULL;
+
+	md5_init(&self->ctx.md5);
+	self->update = (void*)nettle_md5_update;
+	self->digest = (void*)nettle_md5_digest;
+
+	return self;
+}
+
+void crypto_hash_del(struct crypto_hash* self)
+{
+	free(self);
+}
+
+void crypto_hash_append(struct crypto_hash* self, const uint8_t* src,
+		size_t len)
+{
+	self->update(&self->ctx, len, src);
+}
+
+void crypto_hash_digest(struct crypto_hash* self, uint8_t* dst, size_t len)
+{
+	self->digest(&self->ctx, len, dst);
+}

src/server.c

@@ -54,6 +54,10 @@
 #include <gnutls/gnutls.h>
 #endif
 
+#ifdef HAVE_CRYPTO
+#include "crypto.h"
+#endif
+
 #ifndef DRM_FORMAT_INVALID
 #define DRM_FORMAT_INVALID 0
 #endif
@@ -63,6 +67,7 @@
 #endif
 
 #define DEFAULT_NAME "Neat VNC"
+#define SECURITY_TYPES_MAX 3
 
 #define EXPORT __attribute__((visibility("default")))
 
@@ -196,6 +201,8 @@ static int handle_unsupported_version(struct nvnc_client* client)
 
 static int on_version_message(struct nvnc_client* client)
 {
+	struct nvnc* server = client->server;
+
 	if (client->buffer_len - client->buffer_index < 12)
 		return 0;
 
@@ -206,17 +213,31 @@ static int on_version_message(struct nvnc_client* client)
 	if (strcmp(RFB_VERSION_MESSAGE, version_string) != 0)
 		return handle_unsupported_version(client);
 
-	struct rfb_security_types_msg security = { 0 };
+	uint8_t buf[sizeof(struct rfb_security_types_msg) +
-	security.n = 1;
+		SECURITY_TYPES_MAX] = {};
-	security.types[0] = RFB_SECURITY_TYPE_NONE;
+	struct rfb_security_types_msg* security =
+		(struct rfb_security_types_msg*)buf;
 
+	security->n = 0;
+	if (client->server->auth_fn) {
 #ifdef ENABLE_TLS
-	if (client->server->auth_fn)
+		if (server->tls_creds) {
-		security.types[0] = RFB_SECURITY_TYPE_VENCRYPT;
+			security->types[security->n++] = RFB_SECURITY_TYPE_VENCRYPT;
+		}
 #endif
 
-	stream_write(client->net_stream, &security, sizeof(security), NULL,
+#ifdef HAVE_CRYPTO
-	             NULL);
+		security->types[security->n++] = RFB_SECURITY_TYPE_APPLE_DH;
+#endif
+	}
+
+	if (security->n == 0) {
+		security->n = 1;
+		security->types[0] = RFB_SECURITY_TYPE_NONE;
+	}
+
+	stream_write(client->net_stream, security, sizeof(*security) +
+			security->n, NULL, NULL);
 
 	client->state = VNC_CLIENT_STATE_WAITING_FOR_SECURITY;
 	return 12;
@@ -368,6 +389,98 @@ static int on_vencrypt_plain_auth_message(struct nvnc_client* client)
 }
 #endif
 
+#ifdef HAVE_CRYPTO
+static int apple_dh_send_public_key(struct nvnc_client* client)
+{
+	client->apple_dh_secret = crypto_keygen();
+	assert(client->apple_dh_secret);
+
+	struct crypto_key* pub =
+		crypto_derive_public_key(client->apple_dh_secret);
+	assert(pub);
+
+	uint8_t mod[256] = {};
+	int mod_len = crypto_key_p(pub, mod, sizeof(mod));
+	assert(mod_len == sizeof(mod));
+
+	uint8_t q[256] = {};
+	int q_len = crypto_key_q(pub, q, sizeof(q));
+	assert(q_len == sizeof(q));
+
+	struct rfb_apple_dh_server_msg msg = {
+		.generator = htons(crypto_key_g(client->apple_dh_secret)),
+		.key_size = htons(q_len),
+	};
+
+	stream_write(client->net_stream, &msg, sizeof(msg), NULL, NULL);
+	stream_write(client->net_stream, mod, mod_len, NULL, NULL);
+	stream_write(client->net_stream, q, q_len, NULL, NULL);
+
+	crypto_key_del(pub);
+	return 0;
+}
+
+static int on_apple_dh_response(struct nvnc_client* client)
+{
+	struct nvnc* server = client->server;
+
+	struct rfb_apple_dh_client_msg* msg =
+	        (void*)(client->msg_buffer + client->buffer_index);
+
+	uint8_t p[256];
+	int key_len = crypto_key_p(client->apple_dh_secret, p, sizeof(p));
+	assert(key_len == sizeof(p));
+
+	if (client->buffer_len - client->buffer_index < sizeof(*msg) + key_len)
+		return 0;
+
+	int g = crypto_key_g(client->apple_dh_secret);
+
+	struct crypto_key* remote_key = crypto_key_new(g, p, key_len,
+			msg->public_key, key_len);
+	assert(remote_key);
+
+	struct crypto_key* shared_secret =
+		crypto_derive_shared_secret(client->apple_dh_secret, remote_key);
+	assert(shared_secret);
+
+	uint8_t shared_buf[256];
+	crypto_key_q(shared_secret, shared_buf, sizeof(shared_buf));
+	crypto_key_del(shared_secret);
+
+	struct crypto_hash* hash_ctx = crypto_hash_new(CRYPTO_HASH_MD5);
+	crypto_hash_append(hash_ctx, shared_buf, sizeof(shared_buf));
+
+	uint8_t hash[16] = {};
+	crypto_hash_digest(hash_ctx, hash, sizeof(hash));
+	crypto_hash_del(hash_ctx);
+
+	struct crypto_cipher* cipher;
+	cipher = crypto_cipher_new(hash, CRYPTO_CIPHER_AES128_ECB);
+	assert(cipher);
+
+	char username[128] = {};
+	char* password = username + 64;
+
+	crypto_cipher_decrypt(cipher, (uint8_t*)username,
+			msg->encrypted_credentials, sizeof(username));
+	username[63] = '\0';
+	username[127] = '\0';
+
+	if (server->auth_fn(username, password, server->auth_ud)) {
+		nvnc_log(NVNC_LOG_INFO, "User \"%s\" authenticated", username);
+		security_handshake_ok(client);
+		client->state = VNC_CLIENT_STATE_WAITING_FOR_INIT;
+	} else {
+		nvnc_log(NVNC_LOG_INFO, "User \"%s\" rejected", username);
+		security_handshake_failed(client, "Invalid username or password");
+		crypto_cipher_del(cipher);
+	}
+
+	return sizeof(*msg) + key_len;
+}
+#endif // HAVE_CRYPTO
+
 static int on_security_message(struct nvnc_client* client)
 {
 	if (client->buffer_len - client->buffer_index < 1)
@@ -385,6 +498,12 @@ static int on_security_message(struct nvnc_client* client)
 		vencrypt_send_version(client);
 		client->state = VNC_CLIENT_STATE_WAITING_FOR_VENCRYPT_VERSION;
 		break;
+#endif
+#ifdef HAVE_CRYPTO
+	case RFB_SECURITY_TYPE_APPLE_DH:
+		apple_dh_send_public_key(client);
+		client->state = VNC_CLIENT_STATE_WAITING_FOR_APPLE_DH_RESPONSE;
+		break;
 #endif
 	default:
 		security_handshake_failed(client, "Unsupported security type");
@@ -1165,6 +1284,10 @@ static int try_read_client_message(struct nvnc_client* client)
 		return on_vencrypt_subtype_message(client);
 	case VNC_CLIENT_STATE_WAITING_FOR_VENCRYPT_PLAIN_AUTH:
 		return on_vencrypt_plain_auth_message(client);
+#endif
+#ifdef HAVE_CRYPTO
+	case VNC_CLIENT_STATE_WAITING_FOR_APPLE_DH_RESPONSE:
+		return on_apple_dh_response(client);
 #endif
 	case VNC_CLIENT_STATE_READY:
 		return on_client_message(client);

src/stream-gnutls.c

@@ -82,6 +82,8 @@ static int stream_gnutls__flush(struct stream* base)
 {
 	struct stream_gnutls* self = (struct stream_gnutls*)base;
 	while (!TAILQ_EMPTY(&self->base.send_queue)) {
+		assert(self->base.state != STREAM_STATE_CLOSED);
+
 		struct stream_req* req = TAILQ_FIRST(&self->base.send_queue);
 
 		ssize_t rc = gnutls_record_send(self->session,

src/stream-tcp.c

@@ -31,10 +31,23 @@
 #include "stream.h"
 #include "stream-common.h"
 #include "sys/queue.h"
+#include "crypto.h"
+#include "neatvnc.h"
 
 static_assert(sizeof(struct stream) <= STREAM_ALLOC_SIZE,
 		"struct stream has grown too large, increase STREAM_ALLOC_SIZE");
 
+static struct rcbuf* encrypt_rcbuf(struct stream* self, struct rcbuf* payload)
+{
+	uint8_t* ciphertext = malloc(payload->size);
+	assert(ciphertext);
+	crypto_cipher_encrypt(self->cipher, ciphertext, payload->payload,
+			payload->size);
+	struct rcbuf* result = rcbuf_new(ciphertext, payload->size);
+	rcbuf_unref(payload);
+	return result;
+}
+
 static int stream_tcp_close(struct stream* self)
 {
 	if (self->state == STREAM_STATE_CLOSED)
@@ -57,6 +70,8 @@ static int stream_tcp_close(struct stream* self)
 
 static void stream_tcp_destroy(struct stream* self)
 {
+	vec_destroy(&self->tmp_buf);
+	crypto_cipher_del(self->cipher);
 	stream_close(self);
 	aml_unref(self->handler);
 	free(self);
@@ -76,7 +91,9 @@ static int stream_tcp__flush(struct stream* self)
 		if (req->exec) {
 			if (req->payload)
 				rcbuf_unref(req->payload);
-			req->payload = req->exec(self, req->userdata);
+			struct rcbuf* payload  = req->exec(self, req->userdata);
+			req->payload = self->cipher ?
+				encrypt_rcbuf(self, payload) : payload;
 		}
 
 		iov[n_msgs].iov_base = req->payload->payload;
@@ -187,11 +204,27 @@ static ssize_t stream_tcp_read(struct stream* self, void* dst, size_t size)
 	if (self->state != STREAM_STATE_NORMAL)
 		return -1;
 
-	ssize_t rc = read(self->fd, dst, size);
+	uint8_t* read_buffer = dst;
+
+	if (self->cipher) {
+		vec_reserve(&self->tmp_buf, size);
+		read_buffer = self->tmp_buf.data;
+	}
+
+	ssize_t rc = read(self->fd, read_buffer, size);
 	if (rc == 0)
 		stream__remote_closed(self);
 	if (rc > 0)
 		self->bytes_received += rc;
+
+	if (rc > 0 && self->cipher && !crypto_cipher_decrypt(self->cipher, dst,
+				read_buffer, rc)) {
+		nvnc_log(NVNC_LOG_ERROR, "Message authentication failed!");
+		stream__remote_closed(self);
+		errno = EPROTO;
+		return -1;
+	}
+
 	return rc;
 }
 
@@ -205,7 +238,7 @@ static int stream_tcp_send(struct stream* self, struct rcbuf* payload,
 	if (!req)
 		return -1;
 
-	req->payload = payload;
+	req->payload = self->cipher ? encrypt_rcbuf(self, payload) : payload;
 	req->on_done = on_done;
 	req->userdata = userdata;
 
@@ -247,6 +280,14 @@ static void stream_tcp_exec_and_send(struct stream* self,
 	stream_tcp__flush(self);
 }
 
+static int stream_tcp_install_cipher(struct stream* self,
+		struct crypto_cipher* cipher)
+{
+	assert(!self->cipher);
+	self->cipher = cipher;
+	return 0;
+}
+
 static struct stream_impl impl = {
 	.close = stream_tcp_close,
 	.destroy = stream_tcp_destroy,
@@ -254,6 +295,7 @@ static struct stream_impl impl = {
 	.send = stream_tcp_send,
 	.send_first = stream_tcp_send_first,
 	.exec_and_send = stream_tcp_exec_and_send,
+	.install_cipher = stream_tcp_install_cipher,
 };
 
 struct stream* stream_new(int fd, stream_event_fn on_event, void* userdata)

src/stream.c

@@ -65,3 +65,11 @@ void stream_exec_and_send(struct stream* self, stream_exec_fn exec_fn,
 	else
 		stream_send(self, exec_fn(self, userdata), NULL, NULL);
 }
+
+int stream_install_cipher(struct stream* self, struct crypto_cipher* cipher)
+{
+	if (!self->impl->install_cipher) {
+		return -1;
+	}
+	return self->impl->install_cipher(self, cipher);
+}