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

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, 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));
-  }
   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, 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));
-  }
   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, 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));
-  }
   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(lua_gettop(L) - base == 2); // we must have 3 args
-  assert(output_ref == LUA_NOREF);
+  id = luaL_checkstring(L, base);
-  lua_pushvalue(L, pos);
+  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);
+  size_t len = 0;
-  output = std::string(lua_tostring(L, -2), luaL_checkinteger(L, -1));
+  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())
+  if (LUA_TSTRING == lua_getfield(L_NSE, -1, "str")) {
-    return output_str;
+    size_t len = 0;
-
+    const char *str = lua_tolstring(L_NSE, -1, &len);
-  /* Auto-formatted table output? */
+    output.assign(str, len);
-  if (output_ref != LUA_NOREF) {
+  }
-    lua_rawgeti(L_NSE, LUA_REGISTRYINDEX, output_ref);
+  else {
-    if (!lua_isnil(L_NSE, -1))
+    lua_pop(L_NSE, 1); // get rid of whatever that was (nil)
-      output = format_obj(L_NSE, -1);
+    /* Auto-formatted table output? */
-
+    if (LUA_TNIL != lua_getfield(L_NSE, -1, "tab"))
-    lua_pop(L_NSE, 1);
+      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_TNIL != lua_getfield(L_NSE, out_obj, "tab")) {
-  if (!lua_isnil(L_NSE, -1)) {
     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)

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() {
+    ScriptResult() : id(NULL), output_ref(LUA_NOREF) {}
-      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 { return id; }
-    const char *get_id (void) const;
     void write_xml() const;
     bool operator<(ScriptResult const &b) const {
-      return this->id.compare(b.id) < 0;
+      return strcmp(this->id, b.id) < 0;
     }
 };