mirrors/neovim
1
0
Fork 0
mirror of https://github.com/neovim/neovim.git synced 2025-09-22 11:18:19 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
4ef3c2c2c121862a9d2b3d1d7578e1f7f1f0f60c
neovim/src/mpack/lmpack.c

1180 lines
32 KiB
C
Raw Blame History

/*
* 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;
}
Reference in New Issue View Git Blame Copy Permalink