tmate-msgpack.c 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221
  1. #include "tmate.h"
  2. #include "tmate-protocol.h"
  3. static void on_encoder_buffer_ready(__unused evutil_socket_t fd,
  4. __unused short what, void *arg)
  5. {
  6. struct tmate_encoder *encoder = arg;
  7. encoder->ev_active = false;
  8. if (encoder->ready_callback)
  9. encoder->ready_callback(encoder->userdata, encoder->buffer);
  10. }
  11. static int on_encoder_write(void *userdata, const char *buf, size_t len)
  12. {
  13. struct tmate_encoder *encoder = userdata;
  14. if (evbuffer_add(encoder->buffer, buf, len) < 0)
  15. tmate_fatal("Cannot buffer encoded data");
  16. if (!encoder->ev_active) {
  17. event_active(&encoder->ev_buffer, EV_READ, 0);
  18. encoder->ev_active = true;
  19. }
  20. return 0;
  21. }
  22. /* Really sad hack, but we can get away with it */
  23. #define tmate_encoder_from_pk(pk) ((struct tmate_encoder *)pk)
  24. void msgpack_pack_string(msgpack_packer *pk, const char *str)
  25. {
  26. size_t len = strlen(str);
  27. msgpack_pack_str(pk, len);
  28. msgpack_pack_str_body(pk, str, len);
  29. }
  30. void msgpack_pack_boolean(msgpack_packer *pk, bool value)
  31. {
  32. if (value)
  33. msgpack_pack_true(pk);
  34. else
  35. msgpack_pack_false(pk);
  36. }
  37. void tmate_encoder_init(struct tmate_encoder *encoder,
  38. tmate_encoder_write_cb *callback,
  39. void *userdata)
  40. {
  41. msgpack_packer_init(&encoder->pk, encoder, &on_encoder_write);
  42. encoder->buffer = evbuffer_new();
  43. encoder->ready_callback = callback;
  44. encoder->userdata = userdata;
  45. if (!encoder->buffer)
  46. tmate_fatal("Can't allocate buffer");
  47. event_set(&encoder->ev_buffer, -1,
  48. EV_READ | EV_PERSIST, on_encoder_buffer_ready, encoder);
  49. event_add(&encoder->ev_buffer, NULL);
  50. encoder->ev_active = false;
  51. }
  52. void tmate_encoder_destroy(struct tmate_encoder *encoder)
  53. {
  54. /* encoder->pk doesn't need any cleanup */
  55. evbuffer_free(encoder->buffer);
  56. event_del(&encoder->ev_buffer);
  57. memset(encoder, 0, sizeof(*encoder));
  58. }
  59. void tmate_encoder_set_ready_callback(struct tmate_encoder *encoder,
  60. tmate_encoder_write_cb *callback,
  61. void *userdata)
  62. {
  63. encoder->ready_callback = callback;
  64. encoder->userdata = userdata;
  65. if (encoder->ready_callback)
  66. encoder->ready_callback(encoder->userdata, encoder->buffer);
  67. }
  68. void tmate_decoder_error(void)
  69. {
  70. /* TODO Don't kill the session, disconnect */
  71. tmate_print_stack_trace();
  72. tmate_fatal("Received a bad message");
  73. }
  74. void init_unpacker(struct tmate_unpacker *uk, msgpack_object obj)
  75. {
  76. if (obj.type != MSGPACK_OBJECT_ARRAY)
  77. tmate_decoder_error();
  78. uk->argv = obj.via.array.ptr;
  79. uk->argc = obj.via.array.size;
  80. }
  81. int64_t unpack_int(struct tmate_unpacker *uk)
  82. {
  83. int64_t val;
  84. if (uk->argc == 0)
  85. tmate_decoder_error();
  86. if (uk->argv[0].type != MSGPACK_OBJECT_POSITIVE_INTEGER &&
  87. uk->argv[0].type != MSGPACK_OBJECT_NEGATIVE_INTEGER)
  88. tmate_decoder_error();
  89. val = uk->argv[0].via.i64;
  90. uk->argv++;
  91. uk->argc--;
  92. return val;
  93. }
  94. bool unpack_bool(struct tmate_unpacker *uk)
  95. {
  96. bool val;
  97. if (uk->argc == 0)
  98. tmate_decoder_error();
  99. if (uk->argv[0].type != MSGPACK_OBJECT_BOOLEAN)
  100. tmate_decoder_error();
  101. val = uk->argv[0].via.boolean;
  102. uk->argv++;
  103. uk->argc--;
  104. return val;
  105. }
  106. void unpack_buffer(struct tmate_unpacker *uk, const char **buf, size_t *len)
  107. {
  108. if (uk->argc == 0)
  109. tmate_decoder_error();
  110. if (uk->argv[0].type != MSGPACK_OBJECT_STR &&
  111. uk->argv[0].type != MSGPACK_OBJECT_BIN)
  112. tmate_decoder_error();
  113. *len = uk->argv[0].via.str.size;
  114. *buf = uk->argv[0].via.str.ptr;
  115. uk->argv++;
  116. uk->argc--;
  117. }
  118. char *unpack_string(struct tmate_unpacker *uk)
  119. {
  120. const char *buf;
  121. char *alloc_buf;
  122. size_t len;
  123. unpack_buffer(uk, &buf, &len);
  124. alloc_buf = xmalloc(len + 1);
  125. memcpy(alloc_buf, buf, len);
  126. alloc_buf[len] = '\0';
  127. return alloc_buf;
  128. }
  129. void unpack_array(struct tmate_unpacker *uk, struct tmate_unpacker *nested)
  130. {
  131. if (uk->argc == 0)
  132. tmate_decoder_error();
  133. init_unpacker(nested, uk->argv[0]);
  134. uk->argv++;
  135. uk->argc--;
  136. }
  137. #define UNPACKER_RESERVE_SIZE 1024
  138. void tmate_decoder_init(struct tmate_decoder *decoder, tmate_decoder_reader *reader,
  139. void *userdata)
  140. {
  141. if (!msgpack_unpacker_init(&decoder->unpacker, UNPACKER_RESERVE_SIZE))
  142. tmate_fatal("Cannot initialize the unpacker");
  143. decoder->reader = reader;
  144. decoder->userdata = userdata;
  145. }
  146. void tmate_decoder_destroy(struct tmate_decoder *decoder)
  147. {
  148. msgpack_unpacker_destroy(&decoder->unpacker);
  149. memset(decoder, 0, sizeof(*decoder));
  150. }
  151. void tmate_decoder_get_buffer(struct tmate_decoder *decoder,
  152. char **buf, size_t *len)
  153. {
  154. if (!msgpack_unpacker_reserve_buffer(&decoder->unpacker, UNPACKER_RESERVE_SIZE))
  155. tmate_fatal("cannot expand decoder buffer");
  156. *buf = msgpack_unpacker_buffer(&decoder->unpacker);
  157. *len = msgpack_unpacker_buffer_capacity(&decoder->unpacker);
  158. }
  159. void tmate_decoder_commit(struct tmate_decoder *decoder, size_t len)
  160. {
  161. struct tmate_unpacker _uk, *uk = &_uk;
  162. msgpack_unpacked result;
  163. msgpack_unpacker_buffer_consumed(&decoder->unpacker, len);
  164. msgpack_unpacked_init(&result);
  165. while (msgpack_unpacker_next(&decoder->unpacker, &result)) {
  166. init_unpacker(uk, result.data);
  167. decoder->reader(decoder->userdata, uk);
  168. }
  169. msgpack_unpacked_destroy(&result);
  170. }