Reusing protocol code from tmate-slave

tmate-msgpack.c

@@ -0,0 +1,206 @@
+#include "tmate.h"
+#include "tmate-protocol.h"
+
+static void on_encoder_buffer_ready(__unused evutil_socket_t fd,
+				    __unused short what, void *arg)
+{
+	struct tmate_encoder *encoder = arg;
+
+	encoder->ev_active = false;
+	if (encoder->ready_callback)
+		encoder->ready_callback(encoder->userdata, encoder->buffer);
+}
+
+static int on_encoder_write(void *userdata, const char *buf, size_t len)
+{
+	struct tmate_encoder *encoder = userdata;
+
+	if (evbuffer_add(encoder->buffer, buf, len) < 0)
+		tmate_fatal("Cannot buffer encoded data");
+
+	if (!encoder->ev_active) {
+		event_active(&encoder->ev_buffer, EV_READ, 0);
+		encoder->ev_active = true;
+	}
+
+	return 0;
+}
+
+/* Really sad hack, but we can get away with it */
+#define tmate_encoder_from_pk(pk) ((struct tmate_encoder *)pk)
+
+void msgpack_pack_string(msgpack_packer *pk, const char *str)
+{
+	size_t len = strlen(str);
+
+	msgpack_pack_str(pk, len);
+	msgpack_pack_str_body(pk, str, len);
+}
+
+void msgpack_pack_boolean(msgpack_packer *pk, bool value)
+{
+	if (value)
+		msgpack_pack_true(pk);
+	else
+		msgpack_pack_false(pk);
+}
+
+void tmate_encoder_init(struct tmate_encoder *encoder,
+			tmate_encoder_write_cb *callback,
+			void *userdata)
+{
+	msgpack_packer_init(&encoder->pk, encoder, &on_encoder_write);
+	encoder->buffer = evbuffer_new();
+	encoder->ready_callback = callback;
+	encoder->userdata = userdata;
+
+	if (!encoder->buffer)
+		tmate_fatal("Can't allocate buffer");
+
+	event_set(&encoder->ev_buffer, -1,
+		  EV_READ | EV_PERSIST, on_encoder_buffer_ready, encoder);
+
+	event_add(&encoder->ev_buffer, NULL);
+
+	encoder->ev_active = false;
+}
+
+void tmate_encoder_set_ready_callback(struct tmate_encoder *encoder,
+				      tmate_encoder_write_cb *callback,
+				      void *userdata)
+{
+	encoder->ready_callback = callback;
+	encoder->userdata = userdata;
+	encoder->ready_callback(encoder->userdata, encoder->buffer);
+}
+
+void tmate_decoder_error(void)
+{
+	/* TODO Don't kill the session, disconnect */
+	tmate_print_stack_trace();
+	tmate_fatal("Received a bad message");
+}
+
+void init_unpacker(struct tmate_unpacker *uk, msgpack_object obj)
+{
+	if (obj.type != MSGPACK_OBJECT_ARRAY)
+		tmate_decoder_error();
+
+	uk->argv = obj.via.array.ptr;
+	uk->argc = obj.via.array.size;
+}
+
+int64_t unpack_int(struct tmate_unpacker *uk)
+{
+	int64_t val;
+
+	if (uk->argc == 0)
+		tmate_decoder_error();
+
+	if (uk->argv[0].type != MSGPACK_OBJECT_POSITIVE_INTEGER &&
+	    uk->argv[0].type != MSGPACK_OBJECT_NEGATIVE_INTEGER)
+		tmate_decoder_error();
+
+	val = uk->argv[0].via.i64;
+
+	uk->argv++;
+	uk->argc--;
+
+	return val;
+}
+
+bool unpack_bool(struct tmate_unpacker *uk)
+{
+	bool val;
+
+	if (uk->argc == 0)
+		tmate_decoder_error();
+
+	if (uk->argv[0].type != MSGPACK_OBJECT_BOOLEAN)
+		tmate_decoder_error();
+
+	val = uk->argv[0].via.boolean;
+
+	uk->argv++;
+	uk->argc--;
+
+	return val;
+}
+
+void unpack_buffer(struct tmate_unpacker *uk, const char **buf, size_t *len)
+{
+	if (uk->argc == 0)
+		tmate_decoder_error();
+
+	if (uk->argv[0].type != MSGPACK_OBJECT_STR &&
+	    uk->argv[0].type != MSGPACK_OBJECT_BIN)
+		tmate_decoder_error();
+
+	*len = uk->argv[0].via.str.size;
+	*buf = uk->argv[0].via.str.ptr;
+
+	uk->argv++;
+	uk->argc--;
+}
+
+char *unpack_string(struct tmate_unpacker *uk)
+{
+	const char *buf;
+	char *alloc_buf;
+	size_t len;
+
+	unpack_buffer(uk, &buf, &len);
+
+	alloc_buf = xmalloc(len + 1);
+	memcpy(alloc_buf, buf, len);
+	alloc_buf[len] = '\0';
+
+	return alloc_buf;
+}
+
+void unpack_array(struct tmate_unpacker *uk, struct tmate_unpacker *nested)
+{
+	if (uk->argc == 0)
+		tmate_decoder_error();
+
+	init_unpacker(nested, uk->argv[0]);
+
+	uk->argv++;
+	uk->argc--;
+}
+
+void tmate_decoder_init(struct tmate_decoder *decoder, tmate_decoder_reader *reader,
+			void *userdata)
+{
+	if (!msgpack_unpacker_init(&decoder->unpacker, TMATE_MAX_MESSAGE_SIZE))
+		tmate_fatal("Cannot initialize the unpacker");
+	decoder->reader = reader;
+	decoder->userdata = userdata;
+}
+
+void tmate_decoder_get_buffer(struct tmate_decoder *decoder,
+			      char **buf, size_t *len)
+{
+	ssize_t current_size = msgpack_unpacker_message_size(&decoder->unpacker);
+	if (!msgpack_unpacker_reserve_buffer(&decoder->unpacker,
+					     TMATE_MAX_MESSAGE_SIZE - current_size))
+		tmate_fatal("cannot expand decoder buffer");
+
+	*buf = msgpack_unpacker_buffer(&decoder->unpacker);
+	*len = msgpack_unpacker_buffer_capacity(&decoder->unpacker);
+}
+
+void tmate_decoder_commit(struct tmate_decoder *decoder, size_t len)
+{
+	struct tmate_unpacker _uk, *uk = &_uk;
+	msgpack_unpacked result;
+
+	msgpack_unpacker_buffer_consumed(&decoder->unpacker, len);
+
+	msgpack_unpacked_init(&result);
+	while (msgpack_unpacker_next(&decoder->unpacker, &result)) {
+		init_unpacker(uk, result.data);
+		decoder->reader(decoder->userdata, uk);
+	}
+	msgpack_unpacked_destroy(&result);
+}