lmpack.c (34206B)
1 /* 2 * This module exports three classes, and each instance of those classes has its 3 * own private registry for temporary reference storage(keeping state between 4 * calls). A private registry makes managing memory much easier since all we 5 * have to do is call luaL_unref passing the registry reference when the 6 * instance is collected by the __gc metamethod. 7 * 8 * This private registry is manipulated with `lmpack_ref` / `lmpack_unref` / 9 * `lmpack_geti`, which are analogous to `luaL_ref` / `luaL_unref` / 10 * `lua_rawgeti` but operate on the private registry passed as argument. 11 * 12 * In order to simplify debug registry leaks during normal operation(with the 13 * leak_test.lua script), these `lmpack_*` registry functions will target the 14 * normal lua registry when MPACK_DEBUG_REGISTRY_LEAK is defined during 15 * compilation. 16 */ 17 #define LUA_LIB 18 #include <limits.h> 19 #include <stdlib.h> 20 #include <string.h> 21 #include <stdio.h> 22 23 #include <lauxlib.h> 24 #include <lua.h> 25 #include <luaconf.h> 26 27 #include "nvim/macros_defs.h" 28 29 #include "lmpack.h" 30 31 #include "rpc.h" 32 33 #define UNPACKER_META_NAME "mpack.Unpacker" 34 #define UNPACK_FN_NAME "decode" 35 #define PACKER_META_NAME "mpack.Packer" 36 #define PACK_FN_NAME "encode" 37 #define SESSION_META_NAME "mpack.Session" 38 #define NIL_NAME "mpack.NIL" 39 #define EMPTY_DICT_NAME "mpack.empty_dict" 40 41 /* 42 * TODO(tarruda): When targeting lua 5.3 and being compiled with `long long` 43 * support(not -ansi), we should make use of lua 64 bit integers for 44 * representing msgpack integers, since `double` can't represent the full range. 45 */ 46 47 #ifndef luaL_reg 48 /* Taken from Lua5.1's lauxlib.h */ 49 #define luaL_reg luaL_Reg 50 #endif 51 52 #if LUA_VERSION_NUM > 501 53 #ifndef luaL_register 54 #define luaL_register(L,n,f) luaL_setfuncs(L,f,0) 55 #endif 56 #endif 57 58 typedef struct { 59 lua_State *L; 60 mpack_parser_t *parser; 61 int reg, ext, unpacking, mtdict; 62 char *string_buffer; 63 } Unpacker; 64 65 typedef struct { 66 lua_State *L; 67 mpack_parser_t *parser; 68 int reg, ext, root, packing, mtdict; 69 int is_bin, is_bin_fn; 70 } Packer; 71 72 typedef struct { 73 lua_State *L; 74 int reg; 75 mpack_rpc_session_t *session; 76 struct { 77 int type; 78 mpack_rpc_message_t msg; 79 int method_or_error; 80 int args_or_result; 81 } unpacked; 82 int unpacker; 83 } Session; 84 85 static int lmpack_ref(lua_State *L, int reg) 86 { 87 #ifdef MPACK_DEBUG_REGISTRY_LEAK 88 return luaL_ref(L, LUA_REGISTRYINDEX); 89 #else 90 int rv; 91 lua_rawgeti(L, LUA_REGISTRYINDEX, reg); 92 lua_pushvalue(L, -2); 93 rv = luaL_ref(L, -2); 94 lua_pop(L, 2); 95 return rv; 96 #endif 97 } 98 99 static void lmpack_unref(lua_State *L, int reg, int ref) 100 { 101 #ifdef MPACK_DEBUG_REGISTRY_LEAK 102 luaL_unref(L, LUA_REGISTRYINDEX, ref); 103 #else 104 lua_rawgeti(L, LUA_REGISTRYINDEX, reg); 105 luaL_unref(L, -1, ref); 106 lua_pop(L, 1); 107 #endif 108 } 109 110 static void lmpack_geti(lua_State *L, int reg, int ref) 111 { 112 #ifdef MPACK_DEBUG_REGISTRY_LEAK 113 lua_rawgeti(L, LUA_REGISTRYINDEX, ref); 114 #else 115 lua_rawgeti(L, LUA_REGISTRYINDEX, reg); 116 lua_rawgeti(L, -1, ref); 117 lua_replace(L, -2); 118 #endif 119 } 120 121 /* make a shallow copy of the table on stack and remove it after the copy is 122 * done */ 123 static void lmpack_shallow_copy(lua_State *L) 124 { 125 lua_newtable(L); 126 lua_pushnil(L); 127 while (lua_next(L, -3)) { 128 lua_pushvalue(L, -2); 129 lua_insert(L, -2); 130 lua_settable(L, -4); 131 } 132 lua_remove(L, -2); 133 } 134 135 static mpack_parser_t *lmpack_grow_parser(mpack_parser_t *parser) 136 { 137 mpack_parser_t *old = parser; 138 mpack_uint32_t new_capacity = old->capacity * 2; 139 parser = malloc(MPACK_PARSER_STRUCT_SIZE(new_capacity)); 140 if (!parser) goto end; 141 mpack_parser_init(parser, new_capacity); 142 mpack_parser_copy(parser, old); 143 free(old); 144 end: 145 return parser; 146 } 147 148 static mpack_rpc_session_t *lmpack_grow_session(mpack_rpc_session_t *session) 149 { 150 mpack_rpc_session_t *old = session; 151 mpack_uint32_t new_capacity = old->capacity * 2; 152 session = malloc(MPACK_RPC_SESSION_STRUCT_SIZE(new_capacity)); 153 if (!session) goto end; 154 mpack_rpc_session_init(session, new_capacity); 155 mpack_rpc_session_copy(session, old); 156 free(old); 157 end: 158 return session; 159 } 160 161 static Unpacker *lmpack_check_unpacker(lua_State *L, int index) 162 { 163 return luaL_checkudata(L, index, UNPACKER_META_NAME); 164 } 165 166 static Packer *lmpack_check_packer(lua_State *L, int index) 167 { 168 return luaL_checkudata(L, index, PACKER_META_NAME); 169 } 170 171 static Session *lmpack_check_session(lua_State *L, int index) 172 { 173 return luaL_checkudata(L, index, SESSION_META_NAME); 174 } 175 176 static int lmpack_isnil(lua_State *L, int index) 177 { 178 int rv; 179 if (!lua_isuserdata(L, index)) return 0; 180 lua_getfield(L, LUA_REGISTRYINDEX, NIL_NAME); 181 rv = lua_rawequal(L, -1, -2); 182 lua_pop(L, 1); 183 return rv; 184 } 185 186 static int lmpack_isunpacker(lua_State *L, int index) 187 { 188 int rv; 189 if (!lua_isuserdata(L, index) || !lua_getmetatable(L, index)) return 0; 190 luaL_getmetatable(L, UNPACKER_META_NAME); 191 rv = lua_rawequal(L, -1, -2); 192 lua_pop(L, 2); 193 return rv; 194 } 195 196 static void lmpack_pushnil(lua_State *L) 197 { 198 lua_getfield(L, LUA_REGISTRYINDEX, NIL_NAME); 199 } 200 201 /* adapted from 202 * https://github.com/antirez/lua-cmsgpack/blob/master/lua_cmsgpack.c */ 203 static mpack_uint32_t lmpack_objlen(lua_State *L, int *is_array) 204 { 205 size_t len, max; 206 int isarr; 207 lua_Number n; 208 #ifndef NDEBUG 209 int top = lua_gettop(L); 210 assert(top); 211 #endif 212 213 if ((lua_type(L, -1)) != LUA_TTABLE) { 214 #if LUA_VERSION_NUM >= 502 215 len = lua_rawlen(L, -1); 216 #elif LUA_VERSION_NUM == 501 217 len = lua_objlen(L, -1); 218 #else 219 #error You have either broken or too old Lua installation. This library requires Lua>=5.1 220 #endif 221 goto end; 222 } 223 224 /* count the number of keys and determine if it is an array */ 225 len = 0; 226 max = 0; 227 isarr = 1; 228 lua_pushnil(L); 229 230 while (lua_next(L, -2)) { 231 lua_pop(L, 1); /* pop value */ 232 isarr = isarr 233 && lua_type(L, -1) == LUA_TNUMBER /* lua number */ 234 && (n = lua_tonumber(L, -1)) > 0 /* greater than 0 */ 235 && (size_t)n == n; /* and integer */ 236 max = isarr && (size_t)n > max ? (size_t)n : max; 237 len++; 238 } 239 240 // when len==0, the caller should guess the type! 241 if (len > 0) { 242 *is_array = isarr && max == len; 243 } 244 245 end: 246 if ((size_t)-1 > (mpack_uint32_t)-1 && len > (mpack_uint32_t)-1) 247 /* msgpack spec doesn't allow lengths > 32 bits */ 248 len = (mpack_uint32_t)-1; 249 assert(top == lua_gettop(L)); 250 return (mpack_uint32_t)len; 251 } 252 253 static int lmpack_unpacker_new(lua_State *L) 254 { 255 Unpacker *rv; 256 257 if (lua_gettop(L) > 1) 258 return luaL_error(L, "expecting at most 1 table argument"); 259 260 rv = lua_newuserdata(L, sizeof(*rv)); 261 rv->parser = malloc(sizeof(*rv->parser)); 262 if (!rv->parser) return luaL_error(L, "Failed to allocate memory"); 263 mpack_parser_init(rv->parser, 0); 264 rv->parser->data.p = rv; 265 rv->string_buffer = NULL; 266 rv->L = L; 267 rv->unpacking = 0; 268 luaL_getmetatable(L, UNPACKER_META_NAME); 269 lua_setmetatable(L, -2); 270 271 #ifndef MPACK_DEBUG_REGISTRY_LEAK 272 lua_newtable(L); 273 rv->reg = luaL_ref(L, LUA_REGISTRYINDEX); 274 #endif 275 rv->ext = LUA_NOREF; 276 277 lua_getfield(L, LUA_REGISTRYINDEX, EMPTY_DICT_NAME); 278 rv->mtdict = lmpack_ref(L, rv->reg); 279 280 if (lua_istable(L, 1)) { 281 /* parse options */ 282 lua_getfield(L, 1, "ext"); 283 if (!lua_isnil(L, -1)) { 284 if (!lua_istable(L, -1)) 285 return luaL_error(L, "\"ext\" option must be a table"); 286 lmpack_shallow_copy(L); 287 } 288 rv->ext = lmpack_ref(L, rv->reg); 289 } 290 291 return 1; 292 } 293 294 static int lmpack_unpacker_delete(lua_State *L) 295 { 296 Unpacker *unpacker = lmpack_check_unpacker(L, 1); 297 if (unpacker->ext != LUA_NOREF) 298 lmpack_unref(L, unpacker->reg, unpacker->ext); 299 #ifndef MPACK_DEBUG_REGISTRY_LEAK 300 luaL_unref(L, LUA_REGISTRYINDEX, unpacker->reg); 301 #endif 302 free(unpacker->parser); 303 return 0; 304 } 305 306 static void lmpack_parse_enter(mpack_parser_t *parser, mpack_node_t *node) 307 { 308 Unpacker *unpacker = parser->data.p; 309 lua_State *L = unpacker->L; 310 311 switch (node->tok.type) { 312 case MPACK_TOKEN_NIL: 313 lmpack_pushnil(L); break; 314 case MPACK_TOKEN_BOOLEAN: 315 lua_pushboolean(L, (int)mpack_unpack_boolean(node->tok)); break; 316 case MPACK_TOKEN_UINT: 317 case MPACK_TOKEN_SINT: 318 case MPACK_TOKEN_FLOAT: 319 lua_pushnumber(L, mpack_unpack_number(node->tok)); break; 320 case MPACK_TOKEN_CHUNK: 321 assert(unpacker->string_buffer); 322 memcpy(unpacker->string_buffer + MPACK_PARENT_NODE(node)->pos, 323 node->tok.data.chunk_ptr, node->tok.length); 324 break; 325 case MPACK_TOKEN_BIN: 326 case MPACK_TOKEN_STR: 327 case MPACK_TOKEN_EXT: 328 unpacker->string_buffer = malloc(node->tok.length); 329 if (!unpacker->string_buffer) luaL_error(L, "Failed to allocate memory"); 330 break; 331 case MPACK_TOKEN_ARRAY: 332 case MPACK_TOKEN_MAP: 333 lua_newtable(L); 334 node->data[0].i = lmpack_ref(L, unpacker->reg); 335 break; 336 } 337 } 338 339 static void lmpack_parse_exit(mpack_parser_t *parser, mpack_node_t *node) 340 { 341 Unpacker *unpacker = parser->data.p; 342 lua_State *L = unpacker->L; 343 mpack_node_t *parent = MPACK_PARENT_NODE(node); 344 345 switch (node->tok.type) { 346 case MPACK_TOKEN_BIN: 347 case MPACK_TOKEN_STR: 348 case MPACK_TOKEN_EXT: 349 lua_pushlstring(L, unpacker->string_buffer, node->tok.length); 350 free(unpacker->string_buffer); 351 unpacker->string_buffer = NULL; 352 if (node->tok.type == MPACK_TOKEN_EXT && unpacker->ext != LUA_NOREF) { 353 /* check if there's a handler for this type */ 354 lmpack_geti(L, unpacker->reg, unpacker->ext); 355 lua_rawgeti(L, -1, node->tok.data.ext_type); 356 if (lua_isfunction(L, -1)) { 357 /* stack: 358 * 359 * -1: ext unpacker function 360 * -2: ext unpackers table 361 * -3: ext string 362 * 363 * We want to call the ext unpacker function with the type and string 364 * as arguments, so push those now 365 */ 366 lua_pushinteger(L, node->tok.data.ext_type); 367 lua_pushvalue(L, -4); 368 lua_call(L, 2, 1); 369 /* stack: 370 * 371 * -1: returned object 372 * -2: ext unpackers table 373 * -3: ext string 374 */ 375 lua_replace(L, -3); 376 } else { 377 /* the last lua_rawgeti should have pushed nil on the stack, 378 * remove it */ 379 lua_pop(L, 1); 380 } 381 /* pop the ext unpackers table */ 382 lua_pop(L, 1); 383 } 384 break; 385 case MPACK_TOKEN_ARRAY: 386 case MPACK_TOKEN_MAP: 387 lmpack_geti(L, unpacker->reg, (int)node->data[0].i); 388 lmpack_unref(L, unpacker->reg, (int)node->data[0].i); 389 if (node->key_visited == 0 && node->tok.type == MPACK_TOKEN_MAP) { 390 lmpack_geti(L, unpacker->reg, unpacker->mtdict); // [table, mtdict] 391 lua_setmetatable(L, -2); // [table] 392 } 393 394 break; 395 default: 396 break; 397 } 398 399 if (parent && parent->tok.type < MPACK_TOKEN_BIN) { 400 /* At this point the parsed object is on the stack. Add it to the parent 401 * container. First put the container on the stack. */ 402 lmpack_geti(L, unpacker->reg, (int)parent->data[0].i); 403 404 if (parent->tok.type == MPACK_TOKEN_ARRAY) { 405 /* Array, save the value on key equal to `parent->pos` */ 406 lua_pushnumber(L, (lua_Number)parent->pos); 407 lua_pushvalue(L, -3); 408 lua_settable(L, -3); 409 } else { 410 assert(parent->tok.type == MPACK_TOKEN_MAP); 411 if (parent->key_visited) { 412 /* save the key on the registry */ 413 lua_pushvalue(L, -2); 414 parent->data[1].i = lmpack_ref(L, unpacker->reg); 415 } else { 416 /* set the key/value pair */ 417 lmpack_geti(L, unpacker->reg, (int)parent->data[1].i); 418 lmpack_unref(L, unpacker->reg, (int)parent->data[1].i); 419 lua_pushvalue(L, -3); 420 lua_settable(L, -3); 421 } 422 } 423 lua_pop(L, 2); /* pop the container/object */ 424 } 425 } 426 427 static int lmpack_unpacker_unpack_str(lua_State *L, Unpacker *unpacker, 428 const char **str, size_t *len) 429 { 430 int rv; 431 432 if (unpacker->unpacking) { 433 return luaL_error(L, "Unpacker instance already working. Use another " 434 "Unpacker or mpack." UNPACK_FN_NAME "() if you " 435 "need to " UNPACK_FN_NAME " from the ext handler"); 436 } 437 438 do { 439 unpacker->unpacking = 1; 440 rv = mpack_parse(unpacker->parser, str, len, lmpack_parse_enter, 441 lmpack_parse_exit); 442 unpacker->unpacking = 0; 443 444 if (rv == MPACK_NOMEM) { 445 unpacker->parser = lmpack_grow_parser(unpacker->parser); 446 if (!unpacker->parser) { 447 return luaL_error(L, "failed to grow Unpacker capacity"); 448 } 449 } 450 } while (rv == MPACK_NOMEM); 451 452 if (rv == MPACK_ERROR) 453 return luaL_error(L, "invalid msgpack string"); 454 455 return rv; 456 } 457 458 static int lmpack_unpacker_unpack(lua_State *L) 459 { 460 int result, argc; 461 lua_Number startpos; 462 size_t len, offset; 463 const char *str, *str_init; 464 Unpacker *unpacker; 465 466 if ((argc = lua_gettop(L)) > 3 || argc < 2) 467 return luaL_error(L, "expecting between 2 and 3 arguments"); 468 469 unpacker = lmpack_check_unpacker(L, 1); 470 unpacker->L = L; 471 472 str_init = str = luaL_checklstring(L, 2, &len); 473 startpos = lua_gettop(L) == 3 ? luaL_checknumber(L, 3) : 1; 474 475 luaL_argcheck(L, startpos > 0, 3, 476 "start position must be greater than zero"); 477 luaL_argcheck(L, (size_t)startpos == startpos, 3, 478 "start position must be an integer"); 479 luaL_argcheck(L, (size_t)startpos <= len, 3, 480 "start position must be less than or equal to the input string length"); 481 482 offset = (size_t)startpos - 1 ; 483 str += offset; 484 len -= offset; 485 result = lmpack_unpacker_unpack_str(L, unpacker, &str, &len); 486 487 if (result == MPACK_EOF) 488 /* if we hit EOF, return nil as the object */ 489 lua_pushnil(L); 490 491 /* also return the new position in the input string */ 492 lua_pushinteger(L, str - str_init + 1); 493 assert(lua_gettop(L) == argc + 2); 494 return 2; 495 } 496 497 static int lmpack_packer_new(lua_State *L) 498 { 499 Packer *rv; 500 501 if (lua_gettop(L) > 1) 502 return luaL_error(L, "expecting at most 1 table argument"); 503 504 rv = lua_newuserdata(L, sizeof(*rv)); 505 rv->parser = malloc(sizeof(*rv->parser)); 506 if (!rv->parser) return luaL_error(L, "failed to allocate parser memory"); 507 mpack_parser_init(rv->parser, 0); 508 rv->parser->data.p = rv; 509 rv->L = L; 510 rv->packing = 0; 511 rv->is_bin = 0; 512 rv->is_bin_fn = LUA_NOREF; 513 luaL_getmetatable(L, PACKER_META_NAME); 514 lua_setmetatable(L, -2); 515 516 #ifndef MPACK_DEBUG_REGISTRY_LEAK 517 lua_newtable(L); 518 rv->reg = luaL_ref(L, LUA_REGISTRYINDEX); 519 #endif 520 rv->ext = LUA_NOREF; 521 522 lua_getfield(L, LUA_REGISTRYINDEX, EMPTY_DICT_NAME); 523 rv->mtdict = lmpack_ref(L, rv->reg); 524 525 if (lua_istable(L, 1)) { 526 /* parse options */ 527 lua_getfield(L, 1, "ext"); 528 if (!lua_isnil(L, -1)) { 529 if (!lua_istable(L, -1)) 530 return luaL_error(L, "\"ext\" option must be a table"); 531 lmpack_shallow_copy(L); 532 } 533 rv->ext = lmpack_ref(L, rv->reg); 534 lua_getfield(L, 1, "is_bin"); 535 if (!lua_isnil(L, -1)) { 536 if (!lua_isboolean(L, -1) && !lua_isfunction(L, -1)) 537 return luaL_error(L, 538 "\"is_bin\" option must be a boolean or function"); 539 rv->is_bin = lua_toboolean(L, -1); 540 if (lua_isfunction(L, -1)) rv->is_bin_fn = lmpack_ref(L, rv->reg); 541 else lua_pop(L, 1); 542 } else { 543 lua_pop(L, 1); 544 } 545 546 } 547 548 return 1; 549 } 550 551 static int lmpack_packer_delete(lua_State *L) 552 { 553 Packer *packer = lmpack_check_packer(L, 1); 554 if (packer->ext != LUA_NOREF) 555 lmpack_unref(L, packer->reg, packer->ext); 556 #ifndef MPACK_DEBUG_REGISTRY_LEAK 557 luaL_unref(L, LUA_REGISTRYINDEX, packer->reg); 558 #endif 559 free(packer->parser); 560 return 0; 561 } 562 563 static void lmpack_unparse_enter(mpack_parser_t *parser, mpack_node_t *node) 564 { 565 int type; 566 Packer *packer = parser->data.p; 567 lua_State *L = packer->L; 568 mpack_node_t *parent = MPACK_PARENT_NODE(node); 569 570 if (parent) { 571 /* get the parent */ 572 lmpack_geti(L, packer->reg, (int)parent->data[0].i); 573 574 if (parent->tok.type > MPACK_TOKEN_MAP) { 575 /* strings are a special case, they are packed as single child chunk 576 * node */ 577 const char *str = lua_tolstring(L, -1, NULL); 578 node->tok = mpack_pack_chunk(str, parent->tok.length); 579 lua_pop(L, 1); 580 return; 581 } 582 583 if (parent->tok.type == MPACK_TOKEN_ARRAY) { 584 /* push the next index */ 585 lua_pushnumber(L, (lua_Number)(parent->pos + 1)); 586 /* push the element */ 587 lua_gettable(L, -2); 588 } else if (parent->tok.type == MPACK_TOKEN_MAP) { 589 int result; 590 /* push the previous key */ 591 lmpack_geti(L, packer->reg, (int)parent->data[1].i); 592 /* push the pair */ 593 result = lua_next(L, -2); 594 assert(result); /* should not be here if the map was fully processed */ 595 (void)result; /* ignore unused warning */ 596 if (parent->key_visited) { 597 /* release the current key */ 598 lmpack_unref(L, packer->reg, (int)parent->data[1].i); 599 /* push key to the top */ 600 lua_pushvalue(L, -2); 601 /* set the key for the next iteration, leaving value on top */ 602 parent->data[1].i = lmpack_ref(L, packer->reg); 603 /* replace key by the value */ 604 lua_replace(L, -2); 605 } else { 606 /* pop value */ 607 lua_pop(L, 1); 608 } 609 } 610 /* remove parent, leaving only the object which will be serialized */ 611 lua_remove(L, -2); 612 } else { 613 /* root object */ 614 lmpack_geti(L, packer->reg, packer->root); 615 } 616 617 type = lua_type(L, -1); 618 619 switch (type) { 620 case LUA_TBOOLEAN: 621 node->tok = mpack_pack_boolean((unsigned)lua_toboolean(L, -1)); 622 break; 623 case LUA_TNUMBER: 624 node->tok = mpack_pack_number(lua_tonumber(L, -1)); 625 break; 626 case LUA_TSTRING: { 627 int is_bin = packer->is_bin; 628 if (is_bin && packer->is_bin_fn != LUA_NOREF) { 629 lmpack_geti(L, packer->reg, packer->is_bin_fn); 630 lua_pushvalue(L, -2); 631 lua_call(L, 1, 1); 632 is_bin = lua_toboolean(L, -1); 633 lua_pop(L, 1); 634 } 635 if (is_bin) node->tok = mpack_pack_bin(lmpack_objlen(L, NULL)); 636 else node->tok = mpack_pack_str(lmpack_objlen(L, NULL)); 637 break; 638 } 639 case LUA_TTABLE: { 640 mpack_uint32_t len; 641 mpack_node_t *n; 642 643 int has_meta = lua_getmetatable(L, -1); 644 int has_mtdict = false; 645 if (has_meta && packer->mtdict != LUA_NOREF) { 646 lmpack_geti(L, packer->reg, packer->mtdict); // [table, metatable, mtdict] 647 has_mtdict = lua_rawequal(L, -1, -2); 648 lua_pop(L, 1); // [table, metatable]; 649 } 650 if (packer->ext != LUA_NOREF && has_meta && !has_mtdict) { 651 /* check if there's a handler for this metatable */ 652 lmpack_geti(L, packer->reg, packer->ext); 653 lua_pushvalue(L, -2); 654 lua_gettable(L, -2); 655 if (lua_isfunction(L, -1)) { 656 lua_Number ext = -1; 657 /* stack: 658 * 659 * -1: ext packer function 660 * -2: ext packers table 661 * -3: metatable 662 * -4: original object 663 * 664 * We want to call the ext packer function with the original object as 665 * argument, so push it on the top 666 */ 667 lua_pushvalue(L, -4); 668 /* handler should return type code and string */ 669 lua_call(L, 1, 2); 670 if (!lua_isnumber(L, -2) || (ext = lua_tonumber(L, -2)) < 0 671 || ext > 127 || (int)ext != ext) 672 luaL_error(L, 673 "the first result from ext packer must be an integer " 674 "between 0 and 127"); 675 if (!lua_isstring(L, -1)) 676 luaL_error(L, 677 "the second result from ext packer must be a string"); 678 node->tok = mpack_pack_ext((int)ext, lmpack_objlen(L, NULL)); 679 /* stack: 680 * 681 * -1: ext string 682 * -2: ext type 683 * -3: ext packers table 684 * -4: metatable 685 * -5: original table 686 * 687 * We want to leave only the returned ext string, so 688 * replace -5 with the string and pop 3 689 */ 690 lua_replace(L, -5); 691 lua_pop(L, 3); 692 break; /* done */ 693 } else { 694 /* stack: 695 * 696 * -1: ext packers table 697 * -2: metatable 698 * -3: original table 699 * 700 * We want to leave only the original table and metatable since they 701 * will be handled below, so pop 1 702 */ 703 lua_pop(L, 1); 704 } 705 } 706 707 if (has_meta) { 708 lua_pop(L, 1); // [table] 709 } 710 711 /* check for cycles */ 712 n = node; 713 while ((n = MPACK_PARENT_NODE(n))) { 714 lmpack_geti(L, packer->reg, (int)n->data[0].i); 715 if (lua_rawequal(L, -1, -2)) { 716 /* break out of cycles with NIL */ 717 node->tok = mpack_pack_nil(); 718 lua_pop(L, 2); 719 lmpack_pushnil(L); 720 goto end; 721 } 722 lua_pop(L, 1); 723 } 724 725 int is_array = !has_mtdict; 726 len = lmpack_objlen(L, &is_array); 727 if (is_array) { 728 node->tok = mpack_pack_array(len); 729 } else { 730 node->tok = mpack_pack_map(len); 731 /* save nil as the previous key to start iteration */ 732 node->data[1].i = LUA_REFNIL; 733 } 734 break; 735 } 736 case LUA_TUSERDATA: 737 if (lmpack_isnil(L, -1)) { 738 node->tok = mpack_pack_nil(); 739 break; 740 } 741 FALLTHROUGH; 742 default: 743 { 744 /* #define FMT */ 745 char errmsg[50]; 746 snprintf(errmsg, 50, "can't serialize object of type %d", type); 747 luaL_error(L, errmsg); 748 } 749 } 750 751 end: 752 node->data[0].i = lmpack_ref(L, packer->reg); 753 } 754 755 static void lmpack_unparse_exit(mpack_parser_t *parser, mpack_node_t *node) 756 { 757 Packer *packer = parser->data.p; 758 lua_State *L = packer->L; 759 if (node->tok.type != MPACK_TOKEN_CHUNK) { 760 /* release the object */ 761 lmpack_unref(L, packer->reg, (int)node->data[0].i); 762 if (node->tok.type == MPACK_TOKEN_MAP) 763 lmpack_unref(L, packer->reg, (int)node->data[1].i); 764 } 765 } 766 767 static int lmpack_packer_pack(lua_State *L) 768 { 769 char *b; 770 size_t bl; 771 int result, argc; 772 Packer *packer; 773 luaL_Buffer buffer; 774 775 if ((argc = lua_gettop(L)) != 2) 776 return luaL_error(L, "expecting exactly 2 arguments"); 777 778 packer = lmpack_check_packer(L, 1); 779 packer->L = L; 780 packer->root = lmpack_ref(L, packer->reg); 781 luaL_buffinit(L, &buffer); 782 b = luaL_prepbuffer(&buffer); 783 bl = LUAL_BUFFERSIZE; 784 785 if (packer->packing) { 786 return luaL_error(L, "Packer instance already working. Use another Packer " 787 "or mpack." PACK_FN_NAME "() if you need to " 788 PACK_FN_NAME " from the ext handler"); 789 } 790 791 do { 792 size_t bl_init = bl; 793 packer->packing = 1; 794 result = mpack_unparse(packer->parser, &b, &bl, lmpack_unparse_enter, 795 lmpack_unparse_exit); 796 packer->packing = 0; 797 798 if (result == MPACK_NOMEM) { 799 packer->parser = lmpack_grow_parser(packer->parser); 800 if (!packer->parser) { 801 return luaL_error(L, "Failed to grow Packer capacity"); 802 } 803 } 804 805 luaL_addsize(&buffer, bl_init - bl); 806 807 if (!bl) { 808 /* buffer empty, resize */ 809 b = luaL_prepbuffer(&buffer); 810 bl = LUAL_BUFFERSIZE; 811 } 812 } while (result == MPACK_EOF || result == MPACK_NOMEM); 813 814 lmpack_unref(L, packer->reg, packer->root); 815 luaL_pushresult(&buffer); 816 assert(lua_gettop(L) == argc); 817 return 1; 818 } 819 820 static int lmpack_session_new(lua_State *L) 821 { 822 Session *rv = lua_newuserdata(L, sizeof(*rv)); 823 rv->session = malloc(sizeof(*rv->session)); 824 if (!rv->session) return luaL_error(L, "Failed to allocate memory"); 825 mpack_rpc_session_init(rv->session, 0); 826 rv->L = L; 827 luaL_getmetatable(L, SESSION_META_NAME); 828 lua_setmetatable(L, -2); 829 #ifndef MPACK_DEBUG_REGISTRY_LEAK 830 lua_newtable(L); 831 rv->reg = luaL_ref(L, LUA_REGISTRYINDEX); 832 #endif 833 rv->unpacker = LUA_REFNIL; 834 rv->unpacked.args_or_result = LUA_NOREF; 835 rv->unpacked.method_or_error = LUA_NOREF; 836 rv->unpacked.type = MPACK_EOF; 837 838 if (lua_istable(L, 1)) { 839 /* parse options */ 840 lua_getfield(L, 1, "unpack"); 841 if (!lmpack_isunpacker(L, -1)) { 842 return luaL_error(L, 843 "\"unpack\" option must be a " UNPACKER_META_NAME " instance"); 844 } 845 rv->unpacker = lmpack_ref(L, rv->reg); 846 } 847 848 return 1; 849 } 850 851 static int lmpack_session_delete(lua_State *L) 852 { 853 Session *session = lmpack_check_session(L, 1); 854 lmpack_unref(L, session->reg, session->unpacker); 855 #ifndef MPACK_DEBUG_REGISTRY_LEAK 856 luaL_unref(L, LUA_REGISTRYINDEX, session->reg); 857 #endif 858 free(session->session); 859 return 0; 860 } 861 862 static int lmpack_session_receive(lua_State *L) 863 { 864 int argc, done, rcount = 3; 865 lua_Number startpos; 866 size_t len; 867 const char *str, *str_init; 868 Session *session; 869 Unpacker *unpacker = NULL; 870 871 if ((argc = lua_gettop(L)) > 3 || argc < 2) 872 return luaL_error(L, "expecting between 2 and 3 arguments"); 873 874 session = lmpack_check_session(L, 1); 875 str_init = str = luaL_checklstring(L, 2, &len); 876 startpos = lua_gettop(L) == 3 ? luaL_checknumber(L, 3) : 1; 877 878 luaL_argcheck(L, startpos > 0, 3, 879 "start position must be greater than zero"); 880 luaL_argcheck(L, (size_t)startpos == startpos, 3, 881 "start position must be an integer"); 882 luaL_argcheck(L, (size_t)startpos <= len, 3, 883 "start position must be less than or equal to the input string length"); 884 885 size_t offset = (size_t)startpos - 1 ; 886 str += offset; 887 len -= offset; 888 889 if (session->unpacker != LUA_REFNIL) { 890 lmpack_geti(L, session->reg, session->unpacker); 891 unpacker = lmpack_check_unpacker(L, -1); 892 unpacker->L = L; 893 rcount += 2; 894 lua_pop(L, 1); 895 } 896 897 for (;;) { 898 int result; 899 900 if (session->unpacked.type == MPACK_EOF) { 901 session->unpacked.type = 902 mpack_rpc_receive(session->session, &str, &len, &session->unpacked.msg); 903 904 if (!unpacker || session->unpacked.type == MPACK_EOF) 905 break; 906 } 907 908 result = lmpack_unpacker_unpack_str(L, unpacker, &str, &len); 909 910 if (result == MPACK_EOF) break; 911 912 if (session->unpacked.method_or_error == LUA_NOREF) { 913 session->unpacked.method_or_error = lmpack_ref(L, session->reg); 914 } else { 915 session->unpacked.args_or_result = lmpack_ref(L, session->reg); 916 break; 917 } 918 } 919 920 done = session->unpacked.type != MPACK_EOF 921 && (session->unpacked.args_or_result != LUA_NOREF || !unpacker); 922 923 if (!done) { 924 lua_pushnil(L); 925 lua_pushnil(L); 926 if (unpacker) { 927 lua_pushnil(L); 928 lua_pushnil(L); 929 } 930 goto end; 931 } 932 933 switch (session->unpacked.type) { 934 case MPACK_RPC_REQUEST: 935 lua_pushstring(L, "request"); 936 lua_pushnumber(L, session->unpacked.msg.id); 937 break; 938 case MPACK_RPC_RESPONSE: 939 lua_pushstring(L, "response"); 940 lmpack_geti(L, session->reg, (int)session->unpacked.msg.data.i); 941 break; 942 case MPACK_RPC_NOTIFICATION: 943 lua_pushstring(L, "notification"); 944 lua_pushnil(L); 945 break; 946 default: 947 /* In most cases the only sane thing to do when receiving invalid 948 * msgpack-rpc is to close the connection, so handle all errors with 949 * this generic message. Later may add more detailed information. */ 950 return luaL_error(L, "invalid msgpack-rpc string"); 951 } 952 953 session->unpacked.type = MPACK_EOF; 954 955 if (unpacker) { 956 lmpack_geti(L, session->reg, session->unpacked.method_or_error); 957 lmpack_geti(L, session->reg, session->unpacked.args_or_result); 958 lmpack_unref(L, session->reg, session->unpacked.method_or_error); 959 lmpack_unref(L, session->reg, session->unpacked.args_or_result); 960 session->unpacked.method_or_error = LUA_NOREF; 961 session->unpacked.args_or_result = LUA_NOREF; 962 } 963 964 end: 965 lua_pushinteger(L, str - str_init + 1); 966 return rcount; 967 } 968 969 static int lmpack_session_request(lua_State *L) 970 { 971 int result; 972 char buf[16], *b = buf; 973 size_t bl = sizeof(buf); 974 Session *session; 975 mpack_data_t data; 976 977 if (lua_gettop(L) > 2 || lua_gettop(L) < 1) 978 return luaL_error(L, "expecting 1 or 2 arguments"); 979 980 session = lmpack_check_session(L, 1); 981 data.i = lua_isnoneornil(L, 2) ? LUA_NOREF : lmpack_ref(L, session->reg); 982 do { 983 result = mpack_rpc_request(session->session, &b, &bl, data); 984 if (result == MPACK_NOMEM) { 985 session->session = lmpack_grow_session(session->session); 986 if (!session->session) 987 return luaL_error(L, "Failed to grow Session capacity"); 988 } 989 } while (result == MPACK_NOMEM); 990 991 assert(result == MPACK_OK); 992 lua_pushlstring(L, buf, sizeof(buf) - bl); 993 return 1; 994 } 995 996 static int lmpack_session_reply(lua_State *L) 997 { 998 int result; 999 char buf[16], *b = buf; 1000 size_t bl = sizeof(buf); 1001 Session *session; 1002 lua_Number id; 1003 1004 if (lua_gettop(L) != 2) 1005 return luaL_error(L, "expecting exactly 2 arguments"); 1006 1007 session = lmpack_check_session(L, 1); 1008 id = lua_tonumber(L, 2); 1009 luaL_argcheck(L, ((size_t)id == id && id >= 0 && id <= 0xffffffff), 2, 1010 "invalid request id"); 1011 result = mpack_rpc_reply(session->session, &b, &bl, (mpack_uint32_t)id); 1012 assert(result == MPACK_OK); 1013 (void)result; /* ignore unused warning */ 1014 lua_pushlstring(L, buf, sizeof(buf) - bl); 1015 return 1; 1016 } 1017 1018 static int lmpack_session_notify(lua_State *L) 1019 { 1020 int result; 1021 char buf[16], *b = buf; 1022 size_t bl = sizeof(buf); 1023 Session *session; 1024 1025 if (lua_gettop(L) != 1) 1026 return luaL_error(L, "expecting exactly 1 argument"); 1027 1028 session = lmpack_check_session(L, 1); 1029 result = mpack_rpc_notify(session->session, &b, &bl); 1030 assert(result == MPACK_OK); 1031 (void)result; /* ignore unused warning */ 1032 lua_pushlstring(L, buf, sizeof(buf) - bl); 1033 return 1; 1034 } 1035 1036 static int lmpack_nil_tostring(lua_State* L) 1037 { 1038 lua_pushfstring(L, NIL_NAME, lua_topointer(L, 1)); 1039 return 1; 1040 } 1041 1042 static int lmpack_unpack(lua_State *L) 1043 { 1044 int result; 1045 size_t len; 1046 const char *str; 1047 Unpacker unpacker; 1048 mpack_parser_t parser; 1049 1050 if (lua_gettop(L) != 1) 1051 return luaL_error(L, "expecting exactly 1 argument"); 1052 1053 str = luaL_checklstring(L, 1, &len); 1054 1055 /* initialize unpacker */ 1056 lua_newtable(L); 1057 unpacker.reg = luaL_ref(L, LUA_REGISTRYINDEX); 1058 unpacker.ext = LUA_NOREF; 1059 unpacker.parser = &parser; 1060 mpack_parser_init(unpacker.parser, 0); 1061 unpacker.parser->data.p = &unpacker; 1062 unpacker.string_buffer = NULL; 1063 unpacker.L = L; 1064 1065 lua_getfield(L, LUA_REGISTRYINDEX, EMPTY_DICT_NAME); 1066 unpacker.mtdict = lmpack_ref(L, unpacker.reg); 1067 1068 result = mpack_parse(&parser, &str, &len, lmpack_parse_enter, 1069 lmpack_parse_exit); 1070 1071 luaL_unref(L, LUA_REGISTRYINDEX, unpacker.reg); 1072 1073 if (result == MPACK_NOMEM) 1074 return luaL_error(L, "object was too deep to unpack"); 1075 else if (result == MPACK_EOF) 1076 return luaL_error(L, "incomplete msgpack string"); 1077 else if (result == MPACK_ERROR) 1078 return luaL_error(L, "invalid msgpack string"); 1079 else if (result == MPACK_OK && len) 1080 return luaL_error(L, "trailing data in msgpack string"); 1081 1082 assert(result == MPACK_OK); 1083 return 1; 1084 } 1085 1086 static int lmpack_pack(lua_State *L) 1087 { 1088 char *b; 1089 size_t bl; 1090 int result; 1091 Packer packer; 1092 mpack_parser_t parser; 1093 luaL_Buffer buffer; 1094 1095 if (lua_gettop(L) != 1) 1096 return luaL_error(L, "expecting exactly 1 argument"); 1097 1098 /* initialize packer */ 1099 lua_newtable(L); 1100 packer.reg = luaL_ref(L, LUA_REGISTRYINDEX); 1101 packer.ext = LUA_NOREF; 1102 packer.parser = &parser; 1103 mpack_parser_init(packer.parser, 0); 1104 packer.parser->data.p = &packer; 1105 packer.is_bin = 0; 1106 packer.L = L; 1107 packer.root = lmpack_ref(L, packer.reg); 1108 1109 lua_getfield(L, LUA_REGISTRYINDEX, EMPTY_DICT_NAME); 1110 packer.mtdict = lmpack_ref(L, packer.reg); 1111 1112 1113 luaL_buffinit(L, &buffer); 1114 b = luaL_prepbuffer(&buffer); 1115 bl = LUAL_BUFFERSIZE; 1116 1117 do { 1118 size_t bl_init = bl; 1119 result = mpack_unparse(packer.parser, &b, &bl, lmpack_unparse_enter, 1120 lmpack_unparse_exit); 1121 1122 if (result == MPACK_NOMEM) { 1123 lmpack_unref(L, packer.reg, packer.root); 1124 luaL_unref(L, LUA_REGISTRYINDEX, packer.reg); 1125 return luaL_error(L, "object was too deep to pack"); 1126 } 1127 1128 luaL_addsize(&buffer, bl_init - bl); 1129 1130 if (!bl) { 1131 /* buffer empty, resize */ 1132 b = luaL_prepbuffer(&buffer); 1133 bl = LUAL_BUFFERSIZE; 1134 } 1135 } while (result == MPACK_EOF); 1136 1137 lmpack_unref(L, packer.reg, packer.root); 1138 luaL_unref(L, LUA_REGISTRYINDEX, packer.reg); 1139 luaL_pushresult(&buffer); 1140 return 1; 1141 } 1142 1143 static const luaL_reg unpacker_methods[] = { 1144 {"__call", lmpack_unpacker_unpack}, 1145 {"__gc", lmpack_unpacker_delete}, 1146 {NULL, NULL} 1147 }; 1148 1149 static const luaL_reg packer_methods[] = { 1150 {"__call", lmpack_packer_pack}, 1151 {"__gc", lmpack_packer_delete}, 1152 {NULL, NULL} 1153 }; 1154 1155 static const luaL_reg session_methods[] = { 1156 {"receive", lmpack_session_receive}, 1157 {"request", lmpack_session_request}, 1158 {"reply", lmpack_session_reply}, 1159 {"notify", lmpack_session_notify}, 1160 {"__gc", lmpack_session_delete}, 1161 {NULL, NULL} 1162 }; 1163 1164 static const luaL_reg mpack_functions[] = { 1165 {"Unpacker", lmpack_unpacker_new}, 1166 {"Packer", lmpack_packer_new}, 1167 {"Session", lmpack_session_new}, 1168 {UNPACK_FN_NAME, lmpack_unpack}, 1169 {PACK_FN_NAME, lmpack_pack}, 1170 {NULL, NULL} 1171 }; 1172 1173 int luaopen_mpack(lua_State *L) 1174 { 1175 /* Unpacker */ 1176 luaL_newmetatable(L, UNPACKER_META_NAME); 1177 lua_pushvalue(L, -1); 1178 lua_setfield(L, -2, "__index"); 1179 luaL_register(L, NULL, unpacker_methods); 1180 lua_pop(L, 1); 1181 /* Packer */ 1182 luaL_newmetatable(L, PACKER_META_NAME); 1183 lua_pushvalue(L, -1); 1184 lua_setfield(L, -2, "__index"); 1185 luaL_register(L, NULL, packer_methods); 1186 lua_pop(L, 1); 1187 /* Session */ 1188 luaL_newmetatable(L, SESSION_META_NAME); 1189 lua_pushvalue(L, -1); 1190 lua_setfield(L, -2, "__index"); 1191 luaL_register(L, NULL, session_methods); 1192 lua_pop(L, 1); 1193 /* NIL */ 1194 /* Check if NIL is already stored in the registry */ 1195 lua_getfield(L, LUA_REGISTRYINDEX, NIL_NAME); 1196 /* If it isn't, create it */ 1197 if (lua_isnil(L, -1)) { 1198 /* Use a constant userdata to represent NIL */ 1199 (void)lua_newuserdata(L, sizeof(void *)); 1200 /* Create a metatable for NIL userdata */ 1201 lua_createtable(L, 0, 1); 1202 lua_pushstring(L, "__tostring"); 1203 lua_pushcfunction(L, lmpack_nil_tostring); 1204 lua_settable(L, -3); 1205 /* Assign the metatable to the userdata object */ 1206 lua_setmetatable(L, -2); 1207 /* Save NIL on the registry so we can access it easily from other functions */ 1208 lua_setfield(L, LUA_REGISTRYINDEX, NIL_NAME); 1209 } 1210 1211 lua_pop(L, 1); 1212 1213 /* module */ 1214 lua_newtable(L); 1215 luaL_register(L, NULL, mpack_functions); 1216 /* save NIL on the module */ 1217 lua_getfield(L, LUA_REGISTRYINDEX, NIL_NAME); 1218 lua_setfield(L, -2, "NIL"); 1219 return 1; 1220 }