diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/mpack/lmpack.c | 1179 |
1 files changed, 1179 insertions, 0 deletions
diff --git a/src/mpack/lmpack.c b/src/mpack/lmpack.c new file mode 100644 index 0000000000..6da28fc263 --- /dev/null +++ b/src/mpack/lmpack.c @@ -0,0 +1,1179 @@ +/* + * This module exports three classes, and each instance of those classes has its + * own private registry for temporary reference storage(keeping state between + * calls). A private registry makes managing memory much easier since all we + * have to do is call luaL_unref passing the registry reference when the + * instance is collected by the __gc metamethod. + * + * This private registry is manipulated with `lmpack_ref` / `lmpack_unref` / + * `lmpack_geti`, which are analogous to `luaL_ref` / `luaL_unref` / + * `lua_rawgeti` but operate on the private registry passed as argument. + * + * In order to simplify debug registry leaks during normal operation(with the + * leak_test.lua script), these `lmpack_*` registry functions will target the + * normal lua registry when MPACK_DEBUG_REGISTRY_LEAK is defined during + * compilation. + */ +#define LUA_LIB +/* for snprintf */ +#define _XOPEN_SOURCE 500 +#include <limits.h> +#include <stdlib.h> +#include <string.h> +#include <stdio.h> + +#include <lauxlib.h> +#include <lua.h> +#include <luaconf.h> + +#ifdef MPACK_USE_SYSTEM +# include <mpack.h> +#else +# define MPACK_API static +# include "mpack-src/src/mpack.c" +#endif + +#define UNPACKER_META_NAME "mpack.Unpacker" +#define PACKER_META_NAME "mpack.Packer" +#define SESSION_META_NAME "mpack.Session" +#define NIL_NAME "mpack.NIL" + +/* + * TODO(tarruda): When targeting lua 5.3 and being compiled with `long long` + * support(not -ansi), we should make use of lua 64 bit integers for + * representing msgpack integers, since `double` can't represent the full range. + */ + +#ifndef luaL_reg +/* Taken from Lua5.1's lauxlib.h */ +#define luaL_reg luaL_Reg +#endif + +#if LUA_VERSION_NUM > 501 +#ifndef luaL_register +#define luaL_register(L,n,f) luaL_setfuncs(L,f,0) +#endif +#endif + +typedef struct { + lua_State *L; + mpack_parser_t *parser; + int reg, ext, unpacking; + char *string_buffer; +} Unpacker; + +typedef struct { + lua_State *L; + mpack_parser_t *parser; + int reg, ext, root, packing; + int is_bin, is_bin_fn; +} Packer; + +typedef struct { + lua_State *L; + int reg; + mpack_rpc_session_t *session; + struct { + int type; + mpack_rpc_message_t msg; + int method_or_error; + int args_or_result; + } unpacked; + int unpacker; +} Session; + +static int lmpack_ref(lua_State *L, int reg) +{ +#ifdef MPACK_DEBUG_REGISTRY_LEAK + return luaL_ref(L, LUA_REGISTRYINDEX); +#else + int rv; + lua_rawgeti(L, LUA_REGISTRYINDEX, reg); + lua_pushvalue(L, -2); + rv = luaL_ref(L, -2); + lua_pop(L, 2); + return rv; +#endif +} + +static void lmpack_unref(lua_State *L, int reg, int ref) +{ +#ifdef MPACK_DEBUG_REGISTRY_LEAK + luaL_unref(L, LUA_REGISTRYINDEX, ref); +#else + lua_rawgeti(L, LUA_REGISTRYINDEX, reg); + luaL_unref(L, -1, ref); + lua_pop(L, 1); +#endif +} + +static void lmpack_geti(lua_State *L, int reg, int ref) +{ +#ifdef MPACK_DEBUG_REGISTRY_LEAK + lua_rawgeti(L, LUA_REGISTRYINDEX, ref); +#else + lua_rawgeti(L, LUA_REGISTRYINDEX, reg); + lua_rawgeti(L, -1, ref); + lua_replace(L, -2); +#endif +} + +/* make a shallow copy of the table on stack and remove it after the copy is + * done */ +static void lmpack_shallow_copy(lua_State *L) +{ + lua_newtable(L); + lua_pushnil(L); + while (lua_next(L, -3)) { + lua_pushvalue(L, -2); + lua_insert(L, -2); + lua_settable(L, -4); + } + lua_remove(L, -2); +} + +static mpack_parser_t *lmpack_grow_parser(mpack_parser_t *parser) +{ + mpack_parser_t *old = parser; + mpack_uint32_t new_capacity = old->capacity * 2; + parser = malloc(MPACK_PARSER_STRUCT_SIZE(new_capacity)); + if (!parser) goto end; + mpack_parser_init(parser, new_capacity); + mpack_parser_copy(parser, old); + free(old); +end: + return parser; +} + +static mpack_rpc_session_t *lmpack_grow_session(mpack_rpc_session_t *session) +{ + mpack_rpc_session_t *old = session; + mpack_uint32_t new_capacity = old->capacity * 2; + session = malloc(MPACK_RPC_SESSION_STRUCT_SIZE(new_capacity)); + if (!session) goto end; + mpack_rpc_session_init(session, new_capacity); + mpack_rpc_session_copy(session, old); + free(old); +end: + return session; +} + +static Unpacker *lmpack_check_unpacker(lua_State *L, int index) +{ + return luaL_checkudata(L, index, UNPACKER_META_NAME); +} + +static Packer *lmpack_check_packer(lua_State *L, int index) +{ + return luaL_checkudata(L, index, PACKER_META_NAME); +} + +static Session *lmpack_check_session(lua_State *L, int index) +{ + return luaL_checkudata(L, index, SESSION_META_NAME); +} + +static int lmpack_isnil(lua_State *L, int index) +{ + int rv; + if (!lua_isuserdata(L, index)) return 0; + lua_getfield(L, LUA_REGISTRYINDEX, NIL_NAME); + rv = lua_rawequal(L, -1, -2); + lua_pop(L, 1); + return rv; +} + +static int lmpack_isunpacker(lua_State *L, int index) +{ + int rv; + if (!lua_isuserdata(L, index) || !lua_getmetatable(L, index)) return 0; + luaL_getmetatable(L, UNPACKER_META_NAME); + rv = lua_rawequal(L, -1, -2); + lua_pop(L, 2); + return rv; +} + +static void lmpack_pushnil(lua_State *L) +{ + lua_getfield(L, LUA_REGISTRYINDEX, NIL_NAME); +} + +/* adapted from + * https://github.com/antirez/lua-cmsgpack/blob/master/lua_cmsgpack.c */ +static mpack_uint32_t lmpack_objlen(lua_State *L, int *is_array) +{ + size_t len, max; + int isarr, type; + lua_Number n; +#ifndef NDEBUG + int top = lua_gettop(L); + assert(top); +#endif + + if ((type = lua_type(L, -1)) != LUA_TTABLE) { +#if LUA_VERSION_NUM >= 502 + len = lua_rawlen(L, -1); +#elif LUA_VERSION_NUM == 501 + len = lua_objlen(L, -1); +#else + #error You have either broken or too old Lua installation. This library requires Lua>=5.1 +#endif + goto end; + } + + /* count the number of keys and determine if it is an array */ + len = 0; + max = 0; + isarr = 1; + lua_pushnil(L); + + while (lua_next(L, -2)) { + lua_pop(L, 1); /* pop value */ + isarr = isarr + && lua_isnumber(L, -1) /* lua number */ + && (n = lua_tonumber(L, -1)) > 0 /* greater than 0 */ + && (size_t)n == n; /* and integer */ + max = isarr && (size_t)n > max ? (size_t)n : max; + len++; + } + + *is_array = isarr && max == len; + +end: + if ((size_t)-1 > (mpack_uint32_t)-1 && len > (mpack_uint32_t)-1) + /* msgpack spec doesn't allow lengths > 32 bits */ + len = (mpack_uint32_t)-1; + assert(top == lua_gettop(L)); + return (mpack_uint32_t)len; +} + +static int lmpack_unpacker_new(lua_State *L) +{ + Unpacker *rv; + + if (lua_gettop(L) > 1) + return luaL_error(L, "expecting at most 1 table argument"); + + rv = lua_newuserdata(L, sizeof(*rv)); + rv->parser = malloc(sizeof(*rv->parser)); + if (!rv->parser) return luaL_error(L, "Failed to allocate memory"); + mpack_parser_init(rv->parser, 0); + rv->parser->data.p = rv; + rv->string_buffer = NULL; + rv->L = L; + rv->unpacking = 0; + luaL_getmetatable(L, UNPACKER_META_NAME); + lua_setmetatable(L, -2); + +#ifndef MPACK_DEBUG_REGISTRY_LEAK + lua_newtable(L); + rv->reg = luaL_ref(L, LUA_REGISTRYINDEX); +#endif + rv->ext = LUA_NOREF; + + if (lua_istable(L, 1)) { + /* parse options */ + lua_getfield(L, 1, "ext"); + if (!lua_isnil(L, -1)) { + if (!lua_istable(L, -1)) + return luaL_error(L, "\"ext\" option must be a table"); + lmpack_shallow_copy(L); + } + rv->ext = lmpack_ref(L, rv->reg); + } + + return 1; +} + +static int lmpack_unpacker_delete(lua_State *L) +{ + Unpacker *unpacker = lmpack_check_unpacker(L, 1); + if (unpacker->ext != LUA_NOREF) + lmpack_unref(L, unpacker->reg, unpacker->ext); +#ifndef MPACK_DEBUG_REGISTRY_LEAK + luaL_unref(L, LUA_REGISTRYINDEX, unpacker->reg); +#endif + free(unpacker->parser); + return 0; +} + +static void lmpack_parse_enter(mpack_parser_t *parser, mpack_node_t *node) +{ + Unpacker *unpacker = parser->data.p; + lua_State *L = unpacker->L; + + switch (node->tok.type) { + case MPACK_TOKEN_NIL: + lmpack_pushnil(L); break; + case MPACK_TOKEN_BOOLEAN: + lua_pushboolean(L, (int)mpack_unpack_boolean(node->tok)); break; + case MPACK_TOKEN_UINT: + case MPACK_TOKEN_SINT: + case MPACK_TOKEN_FLOAT: + lua_pushnumber(L, mpack_unpack_number(node->tok)); break; + case MPACK_TOKEN_CHUNK: + assert(unpacker->string_buffer); + memcpy(unpacker->string_buffer + MPACK_PARENT_NODE(node)->pos, + node->tok.data.chunk_ptr, node->tok.length); + break; + case MPACK_TOKEN_BIN: + case MPACK_TOKEN_STR: + case MPACK_TOKEN_EXT: + unpacker->string_buffer = malloc(node->tok.length); + if (!unpacker->string_buffer) luaL_error(L, "Failed to allocate memory"); + break; + case MPACK_TOKEN_ARRAY: + case MPACK_TOKEN_MAP: + lua_newtable(L); + node->data[0].i = lmpack_ref(L, unpacker->reg); + break; + } +} + +static void lmpack_parse_exit(mpack_parser_t *parser, mpack_node_t *node) +{ + Unpacker *unpacker = parser->data.p; + lua_State *L = unpacker->L; + mpack_node_t *parent = MPACK_PARENT_NODE(node); + + switch (node->tok.type) { + case MPACK_TOKEN_BIN: + case MPACK_TOKEN_STR: + case MPACK_TOKEN_EXT: + lua_pushlstring(L, unpacker->string_buffer, node->tok.length); + free(unpacker->string_buffer); + unpacker->string_buffer = NULL; + if (node->tok.type == MPACK_TOKEN_EXT && unpacker->ext != LUA_NOREF) { + /* check if there's a handler for this type */ + lmpack_geti(L, unpacker->reg, unpacker->ext); + lua_rawgeti(L, -1, node->tok.data.ext_type); + if (lua_isfunction(L, -1)) { + /* stack: + * + * -1: ext unpacker function + * -2: ext unpackers table + * -3: ext string + * + * We want to call the ext unpacker function with the type and string + * as arguments, so push those now + */ + lua_pushinteger(L, node->tok.data.ext_type); + lua_pushvalue(L, -4); + lua_call(L, 2, 1); + /* stack: + * + * -1: returned object + * -2: ext unpackers table + * -3: ext string + */ + lua_replace(L, -3); + } else { + /* the last lua_rawgeti should have pushed nil on the stack, + * remove it */ + lua_pop(L, 1); + } + /* pop the ext unpackers table */ + lua_pop(L, 1); + } + break; + case MPACK_TOKEN_ARRAY: + case MPACK_TOKEN_MAP: + lmpack_geti(L, unpacker->reg, (int)node->data[0].i); + lmpack_unref(L, unpacker->reg, (int)node->data[0].i); + break; + default: + break; + } + + if (parent && parent->tok.type < MPACK_TOKEN_BIN) { + /* At this point the parsed object is on the stack. Add it to the parent + * container. First put the container on the stack. */ + lmpack_geti(L, unpacker->reg, (int)parent->data[0].i); + + if (parent->tok.type == MPACK_TOKEN_ARRAY) { + /* Array, save the value on key equal to `parent->pos` */ + lua_pushnumber(L, (lua_Number)parent->pos); + lua_pushvalue(L, -3); + lua_settable(L, -3); + } else { + assert(parent->tok.type == MPACK_TOKEN_MAP); + if (parent->key_visited) { + /* save the key on the registry */ + lua_pushvalue(L, -2); + parent->data[1].i = lmpack_ref(L, unpacker->reg); + } else { + /* set the key/value pair */ + lmpack_geti(L, unpacker->reg, (int)parent->data[1].i); + lmpack_unref(L, unpacker->reg, (int)parent->data[1].i); + lua_pushvalue(L, -3); + lua_settable(L, -3); + } + } + lua_pop(L, 2); /* pop the container/object */ + } +} + +static int lmpack_unpacker_unpack_str(lua_State *L, Unpacker *unpacker, + const char **str, size_t *len) +{ + int rv; + + if (unpacker->unpacking) { + return luaL_error(L, "Unpacker instance already working. Use another " + "Unpacker or the module's \"unpack\" function if you " + "need to unpack from the ext handler"); + } + + do { + unpacker->unpacking = 1; + rv = mpack_parse(unpacker->parser, str, len, lmpack_parse_enter, + lmpack_parse_exit); + unpacker->unpacking = 0; + + if (rv == MPACK_NOMEM) { + unpacker->parser = lmpack_grow_parser(unpacker->parser); + if (!unpacker->parser) { + unpacker->unpacking = 0; + return luaL_error(L, "failed to grow Unpacker capacity"); + } + } + } while (rv == MPACK_NOMEM); + + if (rv == MPACK_ERROR) + return luaL_error(L, "invalid msgpack string"); + + return rv; +} + +static int lmpack_unpacker_unpack(lua_State *L) +{ + int result, argc; + lua_Number startpos; + size_t len, offset; + const char *str, *str_init; + Unpacker *unpacker; + + if ((argc = lua_gettop(L)) > 3 || argc < 2) + return luaL_error(L, "expecting between 2 and 3 arguments"); + + unpacker = lmpack_check_unpacker(L, 1); + unpacker->L = L; + + str_init = str = luaL_checklstring(L, 2, &len); + startpos = lua_gettop(L) == 3 ? luaL_checknumber(L, 3) : 1; + + luaL_argcheck(L, startpos > 0, 3, + "start position must be greater than zero"); + luaL_argcheck(L, (size_t)startpos == startpos, 3, + "start position must be an integer"); + luaL_argcheck(L, (size_t)startpos <= len, 3, + "start position must be less than or equal to the input string length"); + + offset = (size_t)startpos - 1 ; + str += offset; + len -= offset; + result = lmpack_unpacker_unpack_str(L, unpacker, &str, &len); + + if (result == MPACK_EOF) + /* if we hit EOF, return nil as the object */ + lua_pushnil(L); + + /* also return the new position in the input string */ + lua_pushinteger(L, str - str_init + 1); + assert(lua_gettop(L) == argc + 2); + return 2; +} + +static int lmpack_packer_new(lua_State *L) +{ + Packer *rv; + + if (lua_gettop(L) > 1) + return luaL_error(L, "expecting at most 1 table argument"); + + rv = lua_newuserdata(L, sizeof(*rv)); + rv->parser = malloc(sizeof(*rv->parser)); + if (!rv->parser) return luaL_error(L, "failed to allocate parser memory"); + mpack_parser_init(rv->parser, 0); + rv->parser->data.p = rv; + rv->L = L; + rv->packing = 0; + rv->is_bin = 0; + rv->is_bin_fn = LUA_NOREF; + luaL_getmetatable(L, PACKER_META_NAME); + lua_setmetatable(L, -2); + +#ifndef MPACK_DEBUG_REGISTRY_LEAK + lua_newtable(L); + rv->reg = luaL_ref(L, LUA_REGISTRYINDEX); +#endif + rv->ext = LUA_NOREF; + + if (lua_istable(L, 1)) { + /* parse options */ + lua_getfield(L, 1, "ext"); + if (!lua_isnil(L, -1)) { + if (!lua_istable(L, -1)) + return luaL_error(L, "\"ext\" option must be a table"); + lmpack_shallow_copy(L); + } + rv->ext = lmpack_ref(L, rv->reg); + lua_getfield(L, 1, "is_bin"); + if (!lua_isnil(L, -1)) { + if (!lua_isboolean(L, -1) && !lua_isfunction(L, -1)) + return luaL_error(L, + "\"is_bin\" option must be a boolean or function"); + rv->is_bin = lua_toboolean(L, -1); + if (lua_isfunction(L, -1)) rv->is_bin_fn = lmpack_ref(L, rv->reg); + else lua_pop(L, 1); + } else { + lua_pop(L, 1); + } + + } + + return 1; +} + +static int lmpack_packer_delete(lua_State *L) +{ + Packer *packer = lmpack_check_packer(L, 1); + if (packer->ext != LUA_NOREF) + lmpack_unref(L, packer->reg, packer->ext); +#ifndef MPACK_DEBUG_REGISTRY_LEAK + luaL_unref(L, LUA_REGISTRYINDEX, packer->reg); +#endif + free(packer->parser); + return 0; +} + +static void lmpack_unparse_enter(mpack_parser_t *parser, mpack_node_t *node) +{ + int type; + Packer *packer = parser->data.p; + lua_State *L = packer->L; + mpack_node_t *parent = MPACK_PARENT_NODE(node); + + if (parent) { + /* get the parent */ + lmpack_geti(L, packer->reg, (int)parent->data[0].i); + + if (parent->tok.type > MPACK_TOKEN_MAP) { + /* strings are a special case, they are packed as single child chunk + * node */ + const char *str = lua_tolstring(L, -1, NULL); + node->tok = mpack_pack_chunk(str, parent->tok.length); + lua_pop(L, 1); + return; + } + + if (parent->tok.type == MPACK_TOKEN_ARRAY) { + /* push the next index */ + lua_pushnumber(L, (lua_Number)(parent->pos + 1)); + /* push the element */ + lua_gettable(L, -2); + } else if (parent->tok.type == MPACK_TOKEN_MAP) { + int result; + /* push the previous key */ + lmpack_geti(L, packer->reg, (int)parent->data[1].i); + /* push the pair */ + result = lua_next(L, -2); + assert(result); /* should not be here if the map was fully processed */ + if (parent->key_visited) { + /* release the current key */ + lmpack_unref(L, packer->reg, (int)parent->data[1].i); + /* push key to the top */ + lua_pushvalue(L, -2); + /* set the key for the next iteration, leaving value on top */ + parent->data[1].i = lmpack_ref(L, packer->reg); + /* replace key by the value */ + lua_replace(L, -2); + } else { + /* pop value */ + lua_pop(L, 1); + } + } + /* remove parent, leaving only the object which will be serialized */ + lua_remove(L, -2); + } else { + /* root object */ + lmpack_geti(L, packer->reg, packer->root); + } + + type = lua_type(L, -1); + + switch (type) { + case LUA_TBOOLEAN: + node->tok = mpack_pack_boolean((unsigned)lua_toboolean(L, -1)); + break; + case LUA_TNUMBER: + node->tok = mpack_pack_number(lua_tonumber(L, -1)); + break; + case LUA_TSTRING: { + int is_bin = packer->is_bin; + if (is_bin && packer->is_bin_fn != LUA_NOREF) { + lmpack_geti(L, packer->reg, packer->is_bin_fn); + lua_pushvalue(L, -2); + lua_call(L, 1, 1); + is_bin = lua_toboolean(L, -1); + lua_pop(L, 1); + } + if (is_bin) node->tok = mpack_pack_bin(lmpack_objlen(L, NULL)); + else node->tok = mpack_pack_str(lmpack_objlen(L, NULL)); + break; + } + case LUA_TTABLE: { + int is_array; + mpack_uint32_t len; + mpack_node_t *n; + + if (packer->ext != LUA_NOREF && lua_getmetatable(L, -1)) { + /* check if there's a handler for this metatable */ + lmpack_geti(L, packer->reg, packer->ext); + lua_pushvalue(L, -2); + lua_gettable(L, -2); + if (lua_isfunction(L, -1)) { + lua_Number ext = -1; + /* stack: + * + * -1: ext packer function + * -2: ext packers table + * -3: metatable + * -4: original object + * + * We want to call the ext packer function with the original object as + * argument, so push it on the top + */ + lua_pushvalue(L, -4); + /* handler should return type code and string */ + lua_call(L, 1, 2); + if (!lua_isnumber(L, -2) || (ext = lua_tonumber(L, -2)) < 0 + || ext > 127 || (int)ext != ext) + luaL_error(L, + "the first result from ext packer must be an integer " + "between 0 and 127"); + if (!lua_isstring(L, -1)) + luaL_error(L, + "the second result from ext packer must be a string"); + node->tok = mpack_pack_ext((int)ext, lmpack_objlen(L, NULL)); + /* stack: + * + * -1: ext string + * -2: ext type + * -3: ext packers table + * -4: metatable + * -5: original table + * + * We want to leave only the returned ext string, so + * replace -5 with the string and pop 3 + */ + lua_replace(L, -5); + lua_pop(L, 3); + break; /* done */ + } else { + /* stack: + * + * -1: ext packers table + * -2: metatable + * -3: original table + * + * We want to leave only the original table since it will be handled + * below, so pop 2 + */ + lua_pop(L, 2); + } + } + + /* check for cycles */ + n = node; + while ((n = MPACK_PARENT_NODE(n))) { + lmpack_geti(L, packer->reg, (int)n->data[0].i); + if (lua_rawequal(L, -1, -2)) { + /* break out of cycles with NIL */ + node->tok = mpack_pack_nil(); + lua_pop(L, 2); + lmpack_pushnil(L); + goto end; + } + lua_pop(L, 1); + } + + len = lmpack_objlen(L, &is_array); + if (is_array) { + node->tok = mpack_pack_array(len); + } else { + node->tok = mpack_pack_map(len); + /* save nil as the previous key to start iteration */ + node->data[1].i = LUA_REFNIL; + } + break; + } + case LUA_TUSERDATA: + if (lmpack_isnil(L, -1)) { + node->tok = mpack_pack_nil(); + break; + } + /* Fallthrough */ + default: + { + /* #define FMT */ + char errmsg[50]; + snprintf(errmsg, 50, "can't serialize object of type %d", type); + luaL_error(L, errmsg); + } + } + +end: + node->data[0].i = lmpack_ref(L, packer->reg); +} + +static void lmpack_unparse_exit(mpack_parser_t *parser, mpack_node_t *node) +{ + Packer *packer = parser->data.p; + lua_State *L = packer->L; + if (node->tok.type != MPACK_TOKEN_CHUNK) { + /* release the object */ + lmpack_unref(L, packer->reg, (int)node->data[0].i); + if (node->tok.type == MPACK_TOKEN_MAP) + lmpack_unref(L, packer->reg, (int)node->data[1].i); + } +} + +static int lmpack_packer_pack(lua_State *L) +{ + char *b; + size_t bl; + int result, argc; + Packer *packer; + luaL_Buffer buffer; + + if ((argc = lua_gettop(L)) != 2) + return luaL_error(L, "expecting exactly 2 arguments"); + + packer = lmpack_check_packer(L, 1); + packer->L = L; + packer->root = lmpack_ref(L, packer->reg); + luaL_buffinit(L, &buffer); + b = luaL_prepbuffer(&buffer); + bl = LUAL_BUFFERSIZE; + + if (packer->packing) { + return luaL_error(L, "Packer instance already working. Use another Packer " + "or the module's \"pack\" function if you need to " + "pack from the ext handler"); + } + + do { + size_t bl_init = bl; + packer->packing = 1; + result = mpack_unparse(packer->parser, &b, &bl, lmpack_unparse_enter, + lmpack_unparse_exit); + packer->packing = 0; + + if (result == MPACK_NOMEM) { + packer->parser = lmpack_grow_parser(packer->parser); + if (!packer->parser) { + packer->packing = 0; + return luaL_error(L, "Failed to grow Packer capacity"); + } + } + + luaL_addsize(&buffer, bl_init - bl); + + if (!bl) { + /* buffer empty, resize */ + b = luaL_prepbuffer(&buffer); + bl = LUAL_BUFFERSIZE; + } + } while (result == MPACK_EOF || result == MPACK_NOMEM); + + lmpack_unref(L, packer->reg, packer->root); + luaL_pushresult(&buffer); + assert(lua_gettop(L) == argc); + return 1; +} + +static int lmpack_session_new(lua_State *L) +{ + Session *rv = lua_newuserdata(L, sizeof(*rv)); + rv->session = malloc(sizeof(*rv->session)); + if (!rv->session) return luaL_error(L, "Failed to allocate memory"); + mpack_rpc_session_init(rv->session, 0); + rv->L = L; + luaL_getmetatable(L, SESSION_META_NAME); + lua_setmetatable(L, -2); +#ifndef MPACK_DEBUG_REGISTRY_LEAK + lua_newtable(L); + rv->reg = luaL_ref(L, LUA_REGISTRYINDEX); +#endif + rv->unpacker = LUA_REFNIL; + rv->unpacked.args_or_result = LUA_NOREF; + rv->unpacked.method_or_error = LUA_NOREF; + rv->unpacked.type = MPACK_EOF; + + if (lua_istable(L, 1)) { + /* parse options */ + lua_getfield(L, 1, "unpack"); + if (!lmpack_isunpacker(L, -1)) { + return luaL_error(L, + "\"unpack\" option must be a " UNPACKER_META_NAME " instance"); + } + rv->unpacker = lmpack_ref(L, rv->reg); + } + + return 1; +} + +static int lmpack_session_delete(lua_State *L) +{ + Session *session = lmpack_check_session(L, 1); + lmpack_unref(L, session->reg, session->unpacker); +#ifndef MPACK_DEBUG_REGISTRY_LEAK + luaL_unref(L, LUA_REGISTRYINDEX, session->reg); +#endif + free(session->session); + return 0; +} + +static int lmpack_session_receive(lua_State *L) +{ + int argc, done, rcount = 3; + lua_Number startpos; + size_t len; + const char *str, *str_init; + Session *session; + Unpacker *unpacker = NULL; + + if ((argc = lua_gettop(L)) > 3 || argc < 2) + return luaL_error(L, "expecting between 2 and 3 arguments"); + + session = lmpack_check_session(L, 1); + str_init = str = luaL_checklstring(L, 2, &len); + startpos = lua_gettop(L) == 3 ? luaL_checknumber(L, 3) : 1; + + luaL_argcheck(L, startpos > 0, 3, + "start position must be greater than zero"); + luaL_argcheck(L, (size_t)startpos == startpos, 3, + "start position must be an integer"); + luaL_argcheck(L, (size_t)startpos <= len, 3, + "start position must be less than or equal to the input string length"); + + str += (size_t)startpos - 1; + + if (session->unpacker != LUA_REFNIL) { + lmpack_geti(L, session->reg, session->unpacker); + unpacker = lmpack_check_unpacker(L, -1); + unpacker->L = L; + rcount += 2; + lua_pop(L, 1); + } + + for (;;) { + int result; + + if (session->unpacked.type == MPACK_EOF) { + session->unpacked.type = + mpack_rpc_receive(session->session, &str, &len, &session->unpacked.msg); + + if (!unpacker || session->unpacked.type == MPACK_EOF) + break; + } + + result = lmpack_unpacker_unpack_str(L, unpacker, &str, &len); + + if (result == MPACK_EOF) break; + + if (session->unpacked.method_or_error == LUA_NOREF) { + session->unpacked.method_or_error = lmpack_ref(L, session->reg); + } else { + session->unpacked.args_or_result = lmpack_ref(L, session->reg); + break; + } + } + + done = session->unpacked.type != MPACK_EOF + && (session->unpacked.args_or_result != LUA_NOREF || !unpacker); + + if (!done) { + lua_pushnil(L); + lua_pushnil(L); + if (unpacker) { + lua_pushnil(L); + lua_pushnil(L); + } + goto end; + } + + switch (session->unpacked.type) { + case MPACK_RPC_REQUEST: + lua_pushstring(L, "request"); + lua_pushnumber(L, session->unpacked.msg.id); + break; + case MPACK_RPC_RESPONSE: + lua_pushstring(L, "response"); + lmpack_geti(L, session->reg, (int)session->unpacked.msg.data.i); + break; + case MPACK_RPC_NOTIFICATION: + lua_pushstring(L, "notification"); + lua_pushnil(L); + break; + default: + /* In most cases the only sane thing to do when receiving invalid + * msgpack-rpc is to close the connection, so handle all errors with + * this generic message. Later may add more detailed information. */ + return luaL_error(L, "invalid msgpack-rpc string"); + } + + session->unpacked.type = MPACK_EOF; + + if (unpacker) { + lmpack_geti(L, session->reg, session->unpacked.method_or_error); + lmpack_geti(L, session->reg, session->unpacked.args_or_result); + lmpack_unref(L, session->reg, session->unpacked.method_or_error); + lmpack_unref(L, session->reg, session->unpacked.args_or_result); + session->unpacked.method_or_error = LUA_NOREF; + session->unpacked.args_or_result = LUA_NOREF; + } + +end: + lua_pushinteger(L, str - str_init + 1); + return rcount; +} + +static int lmpack_session_request(lua_State *L) +{ + int result; + char buf[16], *b = buf; + size_t bl = sizeof(buf); + Session *session; + mpack_data_t data; + + if (lua_gettop(L) > 2 || lua_gettop(L) < 1) + return luaL_error(L, "expecting 1 or 2 arguments"); + + session = lmpack_check_session(L, 1); + data.i = lua_isnoneornil(L, 2) ? LUA_NOREF : lmpack_ref(L, session->reg); + do { + result = mpack_rpc_request(session->session, &b, &bl, data); + if (result == MPACK_NOMEM) { + session->session = lmpack_grow_session(session->session); + if (!session->session) + return luaL_error(L, "Failed to grow Session capacity"); + } + } while (result == MPACK_NOMEM); + + assert(result == MPACK_OK); + lua_pushlstring(L, buf, sizeof(buf) - bl); + return 1; +} + +static int lmpack_session_reply(lua_State *L) +{ + int result; + char buf[16], *b = buf; + size_t bl = sizeof(buf); + Session *session; + lua_Number id; + + if (lua_gettop(L) != 2) + return luaL_error(L, "expecting exactly 2 arguments"); + + session = lmpack_check_session(L, 1); + id = lua_tonumber(L, 2); + luaL_argcheck(L, ((size_t)id == id && id >= 0 && id <= 0xffffffff), 2, + "invalid request id"); + result = mpack_rpc_reply(session->session, &b, &bl, (mpack_uint32_t)id); + assert(result == MPACK_OK); + lua_pushlstring(L, buf, sizeof(buf) - bl); + return 1; +} + +static int lmpack_session_notify(lua_State *L) +{ + int result; + char buf[16], *b = buf; + size_t bl = sizeof(buf); + Session *session; + + if (lua_gettop(L) != 1) + return luaL_error(L, "expecting exactly 1 argument"); + + session = lmpack_check_session(L, 1); + result = mpack_rpc_notify(session->session, &b, &bl); + assert(result == MPACK_OK); + lua_pushlstring(L, buf, sizeof(buf) - bl); + return 1; +} + +static int lmpack_nil_tostring(lua_State* L) +{ + lua_pushfstring(L, NIL_NAME, lua_topointer(L, 1)); + return 1; +} + +static int lmpack_unpack(lua_State *L) +{ + int result; + size_t len; + const char *str; + Unpacker unpacker; + mpack_parser_t parser; + + if (lua_gettop(L) != 1) + return luaL_error(L, "expecting exactly 1 argument"); + + str = luaL_checklstring(L, 1, &len); + + /* initialize unpacker */ + lua_newtable(L); + unpacker.reg = luaL_ref(L, LUA_REGISTRYINDEX); + unpacker.ext = LUA_NOREF; + unpacker.parser = &parser; + mpack_parser_init(unpacker.parser, 0); + unpacker.parser->data.p = &unpacker; + unpacker.string_buffer = NULL; + unpacker.L = L; + + result = mpack_parse(&parser, &str, &len, lmpack_parse_enter, + lmpack_parse_exit); + + luaL_unref(L, LUA_REGISTRYINDEX, unpacker.reg); + + if (result == MPACK_NOMEM) + return luaL_error(L, "object was too deep to unpack"); + else if (result == MPACK_EOF) + return luaL_error(L, "incomplete msgpack string"); + else if (result == MPACK_ERROR) + return luaL_error(L, "invalid msgpack string"); + else if (result == MPACK_OK && len) + return luaL_error(L, "trailing data in msgpack string"); + + assert(result == MPACK_OK); + return 1; +} + +static int lmpack_pack(lua_State *L) +{ + char *b; + size_t bl; + int result; + Packer packer; + mpack_parser_t parser; + luaL_Buffer buffer; + + if (lua_gettop(L) != 1) + return luaL_error(L, "expecting exactly 1 argument"); + + /* initialize packer */ + lua_newtable(L); + packer.reg = luaL_ref(L, LUA_REGISTRYINDEX); + packer.ext = LUA_NOREF; + packer.parser = &parser; + mpack_parser_init(packer.parser, 0); + packer.parser->data.p = &packer; + packer.is_bin = 0; + packer.L = L; + packer.root = lmpack_ref(L, packer.reg); + + luaL_buffinit(L, &buffer); + b = luaL_prepbuffer(&buffer); + bl = LUAL_BUFFERSIZE; + + do { + size_t bl_init = bl; + result = mpack_unparse(packer.parser, &b, &bl, lmpack_unparse_enter, + lmpack_unparse_exit); + + if (result == MPACK_NOMEM) { + lmpack_unref(L, packer.reg, packer.root); + luaL_unref(L, LUA_REGISTRYINDEX, packer.reg); + return luaL_error(L, "object was too deep to pack"); + } + + luaL_addsize(&buffer, bl_init - bl); + + if (!bl) { + /* buffer empty, resize */ + b = luaL_prepbuffer(&buffer); + bl = LUAL_BUFFERSIZE; + } + } while (result == MPACK_EOF); + + lmpack_unref(L, packer.reg, packer.root); + luaL_unref(L, LUA_REGISTRYINDEX, packer.reg); + luaL_pushresult(&buffer); + return 1; +} + +static const luaL_reg unpacker_methods[] = { + {"__call", lmpack_unpacker_unpack}, + {"__gc", lmpack_unpacker_delete}, + {NULL, NULL} +}; + +static const luaL_reg packer_methods[] = { + {"__call", lmpack_packer_pack}, + {"__gc", lmpack_packer_delete}, + {NULL, NULL} +}; + +static const luaL_reg session_methods[] = { + {"receive", lmpack_session_receive}, + {"request", lmpack_session_request}, + {"reply", lmpack_session_reply}, + {"notify", lmpack_session_notify}, + {"__gc", lmpack_session_delete}, + {NULL, NULL} +}; + +static const luaL_reg mpack_functions[] = { + {"Unpacker", lmpack_unpacker_new}, + {"Packer", lmpack_packer_new}, + {"Session", lmpack_session_new}, + {"unpack", lmpack_unpack}, + {"pack", lmpack_pack}, + {NULL, NULL} +}; + +int luaopen_mpack(lua_State *L) +{ + /* Unpacker */ + luaL_newmetatable(L, UNPACKER_META_NAME); + lua_pushvalue(L, -1); + lua_setfield(L, -2, "__index"); + luaL_register(L, NULL, unpacker_methods); + /* Packer */ + luaL_newmetatable(L, PACKER_META_NAME); + lua_pushvalue(L, -1); + lua_setfield(L, -2, "__index"); + luaL_register(L, NULL, packer_methods); + /* Session */ + luaL_newmetatable(L, SESSION_META_NAME); + lua_pushvalue(L, -1); + lua_setfield(L, -2, "__index"); + luaL_register(L, NULL, session_methods); + /* NIL */ + /* Check if NIL is already stored in the registry */ + lua_getfield(L, LUA_REGISTRYINDEX, NIL_NAME); + /* If it isn't, create it */ + if (lua_isnil(L, -1)) { + /* Use a constant userdata to represent NIL */ + (void)lua_newuserdata(L, sizeof(void *)); + /* Create a metatable for NIL userdata */ + lua_createtable(L, 0, 1); + lua_pushstring(L, "__tostring"); + lua_pushcfunction(L, lmpack_nil_tostring); + lua_settable(L, -3); + /* Assign the metatable to the userdata object */ + lua_setmetatable(L, -2); + /* Save NIL on the registry so we can access it easily from other functions */ + lua_setfield(L, LUA_REGISTRYINDEX, NIL_NAME); + } + /* module */ + lua_newtable(L); + luaL_register(L, NULL, mpack_functions); + /* save NIL on the module */ + lua_getfield(L, LUA_REGISTRYINDEX, NIL_NAME); + lua_setfield(L, -2, "NIL"); + return 1; +} |