From f59d546c8fe7dd9129ed2f7492f52f9388bf3051 Mon Sep 17 00:00:00 2001 From: dmiller Date: Wed, 16 Nov 2022 15:55:54 +0000 Subject: [PATCH] Use existing Lua strings for script output; avoid creating copies. --- nse_main.cc | 134 +++++++++++++++++++++++++++------------------------- nse_main.h | 16 ++----- 2 files changed, 73 insertions(+), 77 deletions(-) diff --git a/nse_main.cc b/nse_main.cc index 152bece2a..56a084d6c 100644 --- a/nse_main.cc +++ b/nse_main.cc @@ -120,12 +120,7 @@ static int ports (lua_State *L) static int script_set_output (lua_State *L) { ScriptResult *sr = new ScriptResult; - sr->set_id(luaL_checkstring(L, 1)); - sr->set_output_tab(L, 2); - if (!lua_isnil(L, 3)) { - lua_len(L, 3); - sr->set_output_str(luaL_checkstring(L, 3), luaL_checkinteger(L,-1)); - } + sr->set_output_tab(L, 1); script_scan_results.insert(sr); return 0; } @@ -134,12 +129,7 @@ static int host_set_output (lua_State *L) { ScriptResult *sr = new ScriptResult; Target *target = nseU_gettarget(L, 1); - sr->set_id(luaL_checkstring(L, 2)); - sr->set_output_tab(L, 3); - if (!lua_isnil(L, 4)) { - lua_len(L, 4); - sr->set_output_str(luaL_checkstring(L, 4), luaL_checkinteger(L,-1)); - } + sr->set_output_tab(L, 2); target->scriptResults.insert(sr); return 0; } @@ -151,12 +141,7 @@ static int port_set_output (lua_State *L) ScriptResult *sr = new ScriptResult; Target *target = nseU_gettarget(L, 1); p = nseU_getport(L, target, &port, 2); - sr->set_id(luaL_checkstring(L, 3)); - sr->set_output_tab(L, 4); - if (!lua_isnil(L, 5)) { - lua_len(L, 5); - sr->set_output_str(luaL_checkstring(L, 5), luaL_checkinteger(L,-1)); - } + sr->set_output_tab(L, 3); target->ports.addScriptResult(p->portno, p->proto, sr); target->ports.numscriptresults++; return 0; @@ -406,33 +391,51 @@ void ScriptResult::clear (void) log_write(LOG_STDOUT, "ScriptResult::clear %d id %s\n", output_ref, get_id()); luaL_unref(L_NSE, LUA_REGISTRYINDEX, output_ref); output_ref = LUA_NOREF; - output_str.clear(); + id = NULL; } -void ScriptResult::set_output_tab (lua_State *L, int pos) +void ScriptResult::set_output_tab (lua_State *L, int base) { - // No reason to set output of a script twice unless you specifically cleared it. - assert(output_ref == LUA_NOREF); - lua_pushvalue(L, pos); + assert(lua_gettop(L) - base == 2); // we must have 3 args + id = luaL_checkstring(L, base); + lua_newtable(L); + lua_insert(L, base); + + // string output + if (lua_isnil(L, -1)) { + lua_pop(L, 1); + } + else { + if (!lua_isstring(L, -1)) { + luaL_error(L, "String output is not a string"); + return; + } + lua_setfield(L, base, "str"); + } + + // structured output + if (lua_isnil(L, -1)) { + lua_pop(L, 1); + } + else { + lua_setfield(L, base, "tab"); + } + + // script id. We copy it here to ensure ScriptResult::id stays valid + // (i.e. not garbage-collected) + if (!lua_isstring(L, -1)) { + luaL_error(L, "Script ID is not a string"); + return; + } + lua_setfield(L, base, "id"); + assert(lua_gettop(L) == base); output_ref = luaL_ref(L_NSE, LUA_REGISTRYINDEX); if (o.debugging > 3) log_write(LOG_STDOUT, "ScriptResult::set_output_tab %d id %s\n", output_ref, get_id()); } -void ScriptResult::set_output_str (const char *out) +static void format_obj(lua_State *L, int pos, std::string &output) { - output_str = std::string(out); -} - -void ScriptResult::set_output_str (const char *out, size_t len) -{ - output_str = std::string(out, len); -} - -static std::string format_obj(lua_State *L, int pos) -{ - std::string output; - pos = lua_absindex(L, pos); /* Look up the FORMAT_TABLE function from nse_main.lua and call it. */ @@ -441,7 +444,7 @@ static std::string format_obj(lua_State *L, int pos) log_write(LOG_STDOUT, "%s: Cannot find function _R[\"%s\"] that should be in nse_main.lua\n", SCRIPT_ENGINE, NSE_FORMAT_TABLE); lua_pop(L, 1); - return output; + return; } lua_pushvalue(L, pos); @@ -449,46 +452,41 @@ static std::string format_obj(lua_State *L, int pos) if (o.debugging) log_write(LOG_STDOUT, "%s: Error in FORMAT_TABLE: %s\n", SCRIPT_ENGINE, lua_tostring(L, -1)); lua_pop(L, 1); - return output; + return; } - lua_len(L, -1); - output = std::string(lua_tostring(L, -2), luaL_checkinteger(L, -1)); + size_t len = 0; + const char *str = lua_tolstring(L, -1, &len); + output.assign(str, len); lua_pop(L, 1); - - return output; } std::string ScriptResult::get_output_str (void) const { std::string output; + assert(output_ref != LUA_NOREF); + + int out_obj = lua_rawgeti(L_NSE, LUA_REGISTRYINDEX, output_ref); + assert(out_obj == LUA_TTABLE); + /* Explicit string output? */ - if (!output_str.empty()) - return output_str; - - /* Auto-formatted table output? */ - if (output_ref != LUA_NOREF) { - lua_rawgeti(L_NSE, LUA_REGISTRYINDEX, output_ref); - if (!lua_isnil(L_NSE, -1)) - output = format_obj(L_NSE, -1); - - lua_pop(L_NSE, 1); + if (LUA_TSTRING == lua_getfield(L_NSE, -1, "str")) { + size_t len = 0; + const char *str = lua_tolstring(L_NSE, -1, &len); + output.assign(str, len); + } + else { + lua_pop(L_NSE, 1); // get rid of whatever that was (nil) + /* Auto-formatted table output? */ + if (LUA_TNIL != lua_getfield(L_NSE, -1, "tab")) + format_obj(L_NSE, -1, output); } + lua_pop(L_NSE, 2); return output; } -void ScriptResult::set_id (const char *ident) -{ - id = std::string(ident); -} - -const char *ScriptResult::get_id (void) const -{ - return id.c_str(); -} - ScriptResults *get_script_scan_results_obj (void) { return &script_scan_results; @@ -519,6 +517,13 @@ static void format_xml(lua_State *L, int pos) void ScriptResult::write_xml() const { std::string output_str; + assert(output_ref != LUA_NOREF); + assert(id != NULL); + + int out_obj = lua_rawgeti(L_NSE, LUA_REGISTRYINDEX, output_ref); + assert(out_obj == LUA_TTABLE); + + out_obj = lua_gettop(L_NSE); xml_open_start_tag("script"); xml_attribute("id", "%s", get_id()); @@ -532,8 +537,7 @@ void ScriptResult::write_xml() const } /* Any table output? */ - lua_rawgeti(L_NSE, LUA_REGISTRYINDEX, output_ref); - if (!lua_isnil(L_NSE, -1)) { + if (LUA_TNIL != lua_getfield(L_NSE, out_obj, "tab")) { xml_close_start_tag(); format_xml(L_NSE, -1); xml_end_tag(); @@ -541,7 +545,7 @@ void ScriptResult::write_xml() const xml_close_empty_tag(); } - lua_pop(L_NSE, 1); + lua_settop(L_NSE, out_obj - 1); } /* int panic (lua_State *L) diff --git a/nse_main.h b/nse_main.h index 9bce74503..b65161f96 100644 --- a/nse_main.h +++ b/nse_main.h @@ -12,30 +12,22 @@ class ScriptResult { private: - std::string id; + const char *id; /* Structured output table, an integer ref in L_NSE[LUA_REGISTRYINDEX]. */ int output_ref; - /* Unstructured output string, for scripts that do not return a structured - table, or return a string in addition to a table. */ - std::string output_str; public: - ScriptResult() { - output_ref = LUA_NOREF; - } + ScriptResult() : id(NULL), output_ref(LUA_NOREF) {} ~ScriptResult() { // ensures Lua ref is released clear(); } void clear (void); void set_output_tab (lua_State *, int); - void set_output_str (const char *); - void set_output_str (const char *, size_t); std::string get_output_str (void) const; - void set_id (const char *); - const char *get_id (void) const; + const char *get_id (void) const { return id; } void write_xml() const; bool operator<(ScriptResult const &b) const { - return this->id.compare(b.id) < 0; + return strcmp(this->id, b.id) < 0; } };