static int update_checkpoint(lua_State *lua) { static const char *func_name = "update_checkpoint"; lsb_lua_sandbox *lsb = lua_touserdata(lua, lua_upvalueindex(1)); if (!lsb) { return luaL_error(lua, "%s() invalid upvalueindex", func_name); } lsb_heka_sandbox *hsb = lsb_get_parent(lsb); int result = 0; int n = lua_gettop(lua); switch (n) { case 2: // async case luaL_checktype(lua, 2, LUA_TNUMBER); hsb->stats.pm_failures += (unsigned long long)lua_tonumber(lua, 2); // fall thru case 1: luaL_checktype(lua, 1, LUA_TLIGHTUSERDATA); result = hsb->cb.ucp(hsb->parent, lua_touserdata(lua, 1)); break; case 0: // batch case result = hsb->cb.ucp(hsb->parent, NULL); break; default: return luaL_error(lua, "%s() invalid number of args: %d", func_name, n); } if (result) { return luaL_error(lua, "%s() failed: rejected by the callback", func_name); } return result; }
static char* test_api_assertion() { lsb_lua_sandbox *sb = lsb_create(NULL, "lua/counter.lua", "", NULL); lsb_err_value ret = lsb_init(sb, NULL); mu_assert(!ret, "lsb_init() received: %s", ret); lsb_stop_sandbox(NULL); mu_assert(lsb_destroy(NULL) == NULL, "not null"); mu_assert(lsb_usage(NULL, 0, 0) == 0, "not 0"); mu_assert(lsb_usage(sb, LSB_UT_MAX, 0) == 0, "not 0"); mu_assert(lsb_usage(sb, 0, LSB_US_MAX) == 0, "not 0"); mu_assert(strcmp(lsb_get_error(NULL), "") == 0, "not empty"); lsb_set_error(NULL, "foo"); mu_assert(lsb_get_lua(NULL) == NULL, "not null"); mu_assert(lsb_get_lua_file(NULL) == NULL, "not null"); mu_assert(lsb_get_parent(NULL) == NULL, "not null"); mu_assert(lsb_get_logger(NULL) == NULL, "not null"); mu_assert(lsb_get_state(NULL) == LSB_UNKNOWN, "not unknown"); lsb_add_function(NULL, lsb_test_write_output, "foo"); lsb_add_function(sb, NULL, "foo"); lsb_add_function(sb, lsb_test_write_output, NULL); mu_assert(lsb_pcall_setup(NULL, "foo") == LSB_ERR_UTIL_NULL, "not null"); mu_assert(lsb_pcall_setup(sb, NULL) == LSB_ERR_UTIL_NULL, "not null"); lsb_add_function(NULL, NULL, NULL); lsb_pcall_teardown(NULL); lsb_terminate(NULL, NULL); lsb_terminate(sb, NULL); lsb_add_function(sb, lsb_test_write_output, "write_output"); e = lsb_destroy(sb); mu_assert(!e, "lsb_destroy() received: %s", e); return NULL; }
static int async_checkpoint_update(lua_State* lua) { lua_sandbox* lsb = lua_touserdata(lua, lua_upvalueindex(1)); if (!lsb) { return luaL_error(lua, "async_checkpoint_update() invalid upvalueindex"); } luaL_checktype(lua, 1, LUA_TLIGHTUSERDATA); luaL_checktype(lua, 2, LUA_TNUMBER); size_t sequence_id = (size_t)lua_touserdata(lua, 1); size_t failures = lua_tonumber(lua, 2); hs_output_plugin* p = (hs_output_plugin*)lsb_get_parent(lsb); int i = sequence_id % p->async_len; pthread_mutex_lock(&p->cp_lock); if (p->async_cp[i].input.id >= p->cp.input.id && p->async_cp[i].input.offset > p->cp.input.offset) { p->cp.input.id = p->async_cp[i].input.id; p->cp.input.offset = p->async_cp[i].input.offset; } if (p->async_cp[i].analysis.id >= p->cp.analysis.id && p->async_cp[i].analysis.offset > p->cp.analysis.offset) { p->cp.analysis.id = p->async_cp[i].analysis.id; p->cp.analysis.offset = p->async_cp[i].analysis.offset; } p->sb->stats.pm_failures += failures; pthread_mutex_unlock(&p->cp_lock); return 0; }
static int inject_message_analysis(lua_State *lua) { luaL_checktype(lua, 1, LUA_TTABLE); lsb_lua_sandbox *lsb = lua_touserdata(lua, lua_upvalueindex(1)); if (NULL == lsb) { return luaL_error(lua, "%s() invalid lightuserdata", im_func_name); } lsb_heka_sandbox *hsb = lsb_get_parent(lsb); lua_pushstring(lua, hsb->name); lua_setfield(lua, 1, LSB_LOGGER); lua_pushstring(lua, hsb->hostname); lua_setfield(lua, 1, LSB_HOSTNAME); if (heka_encode_message_table(lsb, 1)) { return luaL_error(lua, "%s() failed: %s", im_func_name, lsb_get_error(lsb)); } size_t output_len = 0; const char *output = lsb_get_output(lsb, &output_len); if (hsb->cb.aim(hsb->parent, output, output_len) != 0) { return luaL_error(lua, "%s() failed: rejected by the callback", im_func_name); } ++hsb->stats.im_cnt; hsb->stats.im_bytes += output_len; return 0; }
int heka_encode_message(lua_State *lua) { int n = lua_gettop(lua); bool framed = false; switch (n) { case 2: luaL_checktype(lua, 2, LUA_TBOOLEAN); framed = lua_toboolean(lua, 2); // fall thru case 1: luaL_checktype(lua, 1, LUA_TTABLE); break; default: return luaL_argerror(lua, n, "incorrect number of arguments"); } lsb_lua_sandbox *lsb = lua_touserdata(lua, lua_upvalueindex(1)); if (!lsb) { return luaL_error(lua, "encode_message() invalid upvalueindex"); } lsb_heka_sandbox *hsb = lsb_get_parent(lsb); set_missing_headers(lua, 1, hsb); lsb->output.pos = 0; lsb_err_value ret = heka_encode_message_table(lsb, 1); if (ret) { const char *err = lsb_get_error(lsb); if (strlen(err) == 0) err = ret; return luaL_error(lua, "encode_message() failed: %s", err); } size_t len = 0; const char *output = lsb_get_output(lsb, &len); lsb->usage[LSB_UT_OUTPUT][LSB_US_CURRENT] = len; if (framed) { char header[14] = "\x1e\x00\x08"; // up to 10 varint bytes and a \x1f int hlen = lsb_pb_output_varint(header + 3, len) + 1; lsb->usage[LSB_UT_OUTPUT][LSB_US_CURRENT] = len + hlen + LSB_HDR_FRAME_SIZE; header[1] = (char)hlen; header[hlen + 2] = '\x1f'; luaL_Buffer b; luaL_buffinit(lua, &b); luaL_addlstring(&b, header, hlen + LSB_HDR_FRAME_SIZE); luaL_addlstring(&b, output, len); luaL_pushresult(&b); } else { lua_pushlstring(lua, output, len); } if (lsb->usage[LSB_UT_OUTPUT][LSB_US_CURRENT] > lsb->usage[LSB_UT_OUTPUT][LSB_US_MAXIMUM]) { lsb->usage[LSB_UT_OUTPUT][LSB_US_MAXIMUM] = lsb->usage[LSB_UT_OUTPUT][LSB_US_CURRENT]; } return 1; }
static int read_message(lua_State *lua) { lsb_lua_sandbox *lsb = lua_touserdata(lua, lua_upvalueindex(1)); if (NULL == lsb) { return luaL_error(lua, "read_message() invalid lightuserdata"); } lsb_heka_sandbox *hsb = lsb_get_parent(lsb); if (!hsb->msg || !hsb->msg->raw.s) { lua_pushnil(lua); return 1; } return heka_read_message(lua, hsb->msg); }
static int read_message(lua_State* lua) { lua_sandbox* lsb = lua_touserdata(lua, lua_upvalueindex(1)); if (!lsb) { return luaL_error(lua, "read_message() invalid upvalueindex"); } hs_output_plugin* p = (hs_output_plugin*)lsb_get_parent(lsb); if (!p->matched || !p->msg) { lua_pushnil(lua); return 1; } return hs_read_message(lua, p->msg); }
static int read_message(lua_State* lua) { void* luserdata = lua_touserdata(lua, lua_upvalueindex(1)); if (NULL == luserdata) { return luaL_error(lua, "read_message() invalid lightuserdata"); } lua_sandbox* lsb = (lua_sandbox*)luserdata; hs_analysis_plugin* p = (hs_analysis_plugin*)lsb_get_parent(lsb); if (!p->at->matched || !p->at->msg) { lua_pushnil(lua); return 1; } return hs_read_message(lua, p->at->msg); }
static int inject_message(lua_State* L) { static unsigned char header[14]; void* luserdata = lua_touserdata(L, lua_upvalueindex(1)); if (NULL == luserdata) { return luaL_error(L, "inject_message() invalid lightuserdata"); } lua_sandbox* lsb = (lua_sandbox*)luserdata; hs_analysis_plugin* p = (hs_analysis_plugin*)lsb_get_parent(lsb); if (lua_type(L, 1) == LUA_TTABLE) { lua_pushstring(L, p->sb->name); lua_setfield(L, 1, "Logger"); lua_pushstring(L, p->at->plugins->cfg->hostname); lua_setfield(L, 1, "Hostname"); lua_pushinteger(L, p->at->plugins->cfg->pid); lua_setfield(L, 1, "Pid"); } if (lsb_output_protobuf(lsb, 1, 0) != 0) { return luaL_error(L, "inject_message() could not encode protobuf - %s", lsb_get_error(lsb)); } size_t output_len = 0; const char* output = lsb_get_output(lsb, &output_len); pthread_mutex_lock(&p->at->plugins->output.lock); int len = hs_write_varint(header + 3, output_len); int tlen = 4 + len + output_len; ++p->sb->stats.im_cnt; p->sb->stats.im_bytes += tlen; header[0] = 0x1e; header[1] = (char)(len + 1); header[2] = 0x08; header[3 + len] = 0x1f; fwrite(header, 4 + len, 1, p->at->plugins->output.fh); fwrite(output, output_len, 1, p->at->plugins->output.fh); p->at->plugins->output.offset += tlen; if (p->at->plugins->output.offset >= (size_t)p->at->plugins->cfg->output_size) { ++p->at->plugins->output.id; hs_open_output_file(&p->at->plugins->output); } pthread_mutex_unlock(&p->at->plugins->output.lock); return 0; }
static int inject_message_input(lua_State *lua) { lsb_lua_sandbox *lsb = lua_touserdata(lua, lua_upvalueindex(1)); if (NULL == lsb) { return luaL_error(lua, "%s() invalid lightuserdata", im_func_name); } const char *scp = NULL; double ncp = NAN; int t = lua_type(lua, 2); switch (t) { case LUA_TNUMBER: ncp = lua_tonumber(lua, 2); break; case LUA_TSTRING: scp = lua_tostring(lua, 2); break; case LUA_TNONE: case LUA_TNIL: break; default: return luaL_error(lua, "%s() unsupported checkpoint type: %s", im_func_name, lua_typename(lua, t)); } lsb_const_string output; t = lua_type(lua, 1); switch (t) { case LUA_TUSERDATA: { heka_stream_reader *hsr = luaL_checkudata(lua, 1, mozsvc_heka_stream_reader); if (hsr->msg.raw.s) { output.len = hsr->msg.raw.len; output.s = hsr->msg.raw.s; } else { return luaL_error(lua, "%s() attempted to inject a nil message", im_func_name); } } break; case LUA_TSTRING: { lsb_heka_message m; lsb_init_heka_message(&m, 8); output.s = lua_tolstring(lua, 1, &output.len); bool ok = lsb_decode_heka_message(&m, output.s, output.len, NULL); lsb_free_heka_message(&m); if (!ok) { return luaL_error(lua, "%s() attempted to inject a invalid protobuf " "string", im_func_name); } } break; case LUA_TTABLE: if (heka_encode_message_table(lsb, 1)) { const char *err = lsb_get_error(lsb); if (strlen(err) == 0) err = "exceeded output_limit"; return luaL_error(lua, "%s() failed: %s", im_func_name, err); } output.len = 0; output.s = lsb_get_output(lsb, &output.len); break; default: return luaL_error(lua, "%s() unsupported message type: %s", im_func_name, lua_typename(lua, t)); } lsb_heka_sandbox *hsb = lsb_get_parent(lsb); if (hsb->cb.iim(hsb->parent, output.s, output.len, ncp, scp) != 0) { return luaL_error(lua, "%s() failed: rejected by the callback", im_func_name); } ++hsb->stats.im_cnt; hsb->stats.im_bytes += output.len; return 0; }
static int hj_parse_message(lua_State *lua) { lsb_lua_sandbox *lsb = static_cast<lsb_lua_sandbox *> (lua_touserdata(lua, lua_upvalueindex(1))); if (!lsb) { return luaL_error(lua, "%s() invalid lightuserdata", __FUNCTION__); } int n = lua_gettop(lua); int idx = 1; lsb_heka_message *msg = NULL; lsb_heka_sandbox *hsb = static_cast<lsb_heka_sandbox *>(lsb_get_parent(lsb)); if (hsb->type == 'i') { luaL_argcheck(lua, n >= 2 && n <= 4, 0, "invalid number of arguments"); heka_stream_reader *hsr = static_cast<heka_stream_reader *> (luaL_checkudata(lua, 1, mozsvc_heka_stream_reader)); msg = &hsr->msg; idx = 2; } else { luaL_argcheck(lua, n >= 1 && n <= 3, 0, "invalid number of arguments"); if (!hsb->msg || !hsb->msg->raw.s) { return luaL_error(lua, "no active message"); } msg = hsb->msg; } lsb_const_string json = read_message(lua, idx, msg); if (!json.s) return luaL_error(lua, "field not found"); char *inflated = NULL; #ifdef HAVE_ZLIB // automatically handle gzipped strings (optimization for Mozilla telemetry // messages) if (json.len > 2) { if (json.s[0] == 0x1f && (unsigned char)json.s[1] == 0x8b) { size_t mms = (size_t)lua_tointeger(lua, lua_upvalueindex(2)); inflated = lsb_ungzip(json.s, json.len, mms, NULL); if (!inflated) return luaL_error(lua, "lsb_ungzip failed"); } } #endif heka_json *hj = static_cast<heka_json *>(lua_newuserdata(lua, sizeof*hj)); hj->doc = new rj::Document; hj->val = NULL; hj->insitu = inflated; hj->refs = new std::set<rj::Value *>; luaL_getmetatable(lua, mozsvc_heka_json); lua_setmetatable(lua, -2); if (!hj->doc || !hj->refs) { lua_pushstring(lua, "memory allocation failed"); return lua_error(lua); } bool err = false; if (hj->insitu) { if (hj->doc->ParseInsitu<rj::kParseStopWhenDoneFlag>(hj->insitu) .HasParseError()) { err = true; lua_pushfstring(lua, "failed to parse offset:%f %s", (lua_Number)hj->doc->GetErrorOffset(), rj::GetParseError_En(hj->doc->GetParseError())); } } else { rj::MemoryStream ms(json.s, json.len); if (hj->doc->ParseStream<0, rj::UTF8<> >(ms).HasParseError()) { err = true; lua_pushfstring(lua, "failed to parse offset:%f %s", (lua_Number)hj->doc->GetErrorOffset(), rj::GetParseError_En(hj->doc->GetParseError())); } } if (err) return lua_error(lua); hj->refs->insert(hj->doc); return 1; }