diff --git a/lualib-src/lua-socket.c b/lualib-src/lua-socket.c index 2428914e..ad717f94 100644 --- a/lualib-src/lua-socket.c +++ b/lualib-src/lua-socket.c @@ -188,10 +188,29 @@ pop_lstring(lua_State *L, struct socket_buffer *sb, int sz, int skip) { luaL_pushresult(&b); } +static int +lheader(lua_State *L) { + size_t len; + const uint8_t * s = (const uint8_t *)luaL_checklstring(L, 1, &len); + if (len > 4 || len < 1) { + return luaL_error(L, "Invalid read %s", s); + } + int i; + size_t sz = 0; + for (i=0;i<(int)len;i++) { + sz <<= 8; + sz |= s[i]; + } + + lua_pushunsigned(L, sz); + + return 1; +} + /* userdata send_buffer table pool - integer sz or string (big endian) + integer sz */ static int lpopbuffer(lua_State *L) { @@ -200,23 +219,7 @@ lpopbuffer(lua_State *L) { return luaL_error(L, "Need buffer object at param 1"); } luaL_checktype(L,2,LUA_TTABLE); - int type = lua_type(L,3); - int sz; - if (type == LUA_TNUMBER) { - sz = lua_tointeger(L,3); - } else { - size_t len; - const uint8_t * s = (const uint8_t *)luaL_checklstring(L, 3, &len); - if (len > 4 || len < 1) { - return luaL_error(L, "Invalid read %s", s); - } - int i; - sz = 0; - for (i=0;i<(int)len;i++) { - sz <<= 8; - sz |= s[i]; - } - } + int sz = luaL_checkinteger(L,3); if (sb->size < sz || sz == 0) { lua_pushnil(L); } else { @@ -485,6 +488,7 @@ luaopen_socketdriver(lua_State *L) { { "clear", lclearbuffer }, { "readline", lreadline }, { "str2p", lstr2p }, + { "header", lheader }, { "unpack", lunpack }, { NULL, NULL }, diff --git a/lualib/socket.lua b/lualib/socket.lua index 93b9bf4c..f5319ccf 100644 --- a/lualib/socket.lua +++ b/lualib/socket.lua @@ -266,6 +266,7 @@ end socket.write = assert(driver.send) socket.lwrite = assert(driver.lsend) +socket.header = assert(driver.header) function socket.invalid(id) return socket_pool[id] == nil diff --git a/service/clusterd.lua b/service/clusterd.lua index ef6bff68..563869ac 100644 --- a/service/clusterd.lua +++ b/service/clusterd.lua @@ -10,7 +10,7 @@ local node_session = {} local command = {} local function read_response(sock) - local sz = sock:read(2) + local sz = socket.header(sock:read(2)) local msg = sock:read(sz) return cluster.unpackresponse(msg) -- session, ok, data end