Use existing Lua strings for script output; avoid creating copies.

2025-12-06 04:31:29 +00:00 · 2022-11-16 15:55:54 +00:00
parent 56f59de131
commit f59d546c8f
2 changed files with 73 additions and 77 deletions
--- 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;
-
+  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 (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_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)
--- 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;
    }
 };