Implement a Lua 5.3 plugin loader plugin

Also implemented SOCKS detection in said language.

There are probably going to be some bugs.

The program is no longer Valgrind-clean, as that would require plugin
deinitialization, in which there is very little point.

8 files changed, 573 insertions, 6 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index ea214e2..57d04ba 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -73,7 +73,11 @@ set_target_properties (plugin-http PROPERTIES OUTPUT_NAME http PREFIX "")
 install (TARGETS plugin-http DESTINATION ${plugin_dir})
 
 # Build the other plugins
-foreach (plugin irc ssh)
+set (plugins irc ssh)
+if (WITH_LUA)
+	list (APPEND plugins lua-loader)
+endif (WITH_LUA)
+foreach (plugin ${plugins})
 	set (target plugin-${plugin})
 	add_library (${target} SHARED plugins/${plugin}.c plugin-api.h)
 	target_link_libraries (${target} ${project_libraries})
diff --git a/plugin-api.h b/plugin-api.h
index f74b069..109f729 100644
--- a/plugin-api.h
+++ b/plugin-api.h
@@ -38,11 +38,12 @@ struct service
 {
 	const char *name;                   ///< Name of the service
 	int flags;                          ///< Service flags
+	void *user_data;                    ///< User data
 
 	// scan_init -> on_data* -> [on_eof/on_error] -> on_aborted -> scan_free
 
 	/// Initialize a scan, returning a handle to it
-	void *(*scan_init) (struct unit *u);
+	void *(*scan_init) (struct service *self, struct unit *u);
 
 	/// Destroy the handle created for the scan
 	void (*scan_free) (void *handle);
@@ -67,6 +68,9 @@ struct plugin_api
 	/// Register the plugin for a service
 	void (*register_service) (void *ctx, struct service *info);
 
+	/// Retrieve an item from the configuration
+	const char *(*get_config) (void *ctx, const char *key);
+
 	/// Get the IP address of the target as a string
 	const char *(*unit_get_address) (struct unit *u);
 
diff --git a/plugins/http.c b/plugins/http.c
index 1aa2bd9..a6a58e2 100644
--- a/plugins/http.c
+++ b/plugins/http.c
@@ -94,8 +94,10 @@ on_headers_complete (http_parser *parser)
 }
 
 static void *
-scan_init (struct unit *u)
+scan_init (struct service *service, struct unit *u)
 {
+	(void) service;
+
 	struct str hello;
 	str_init (&hello);
 	str_append_printf (&hello, "GET / HTTP/1.0\r\n"
diff --git a/plugins/irc.c b/plugins/irc.c
index 554c00c..4db36ba 100644
--- a/plugins/irc.c
+++ b/plugins/irc.c
@@ -214,8 +214,10 @@ struct scan_data
 };
 
 static void *
-scan_init (struct unit *u)
+scan_init (struct service *service, struct unit *u)
 {
+	(void) service;
+
 	char nick[IRC_MAX_NICKNAME + 1];
 	size_t i;
 	for (i = 0; i < sizeof nick - 1; i++)
diff --git a/plugins/lua-loader.c b/plugins/lua-loader.c
new file mode 100644
index 0000000..341327b
--- /dev/null
+++ b/plugins/lua-loader.c
@@ -0,0 +1,446 @@
+/*
+ * lua-loader.c: Lua plugin loader plugin
+ *
+ * Copyright (c) 2015, Přemysl Janouch <p.janouch@gmail.com>
+ * All rights reserved.
+ *
+ * Permission to use, copy, modify, and/or distribute this software for any
+ * purpose with or without fee is hereby granted, provided that the above
+ * copyright notice and this permission notice appear in all copies.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
+ * SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
+ * OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
+ * CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+ *
+ */
+
+// I can't really recommend using this interface as it adds a lot of overhead
+
+#include "../utils.c"
+#include "../plugin-api.h"
+
+#include <dirent.h>
+
+#include <lua.h>
+#include <lualib.h>
+#include <lauxlib.h>
+
+// --- Utilities ---------------------------------------------------------------
+
+static struct plugin_data
+{
+	void *ctx;                          ///< Application context
+	struct plugin_api *api;             ///< Plugin API vtable
+}
+g_data;
+
+static void *
+xlua_alloc (void *ud, void *ptr, size_t o_size, size_t n_size)
+{
+	(void) ud;
+	(void) o_size;
+
+	if (n_size)
+		return realloc (ptr, n_size);
+
+	free (ptr);
+	return NULL;
+}
+
+static int
+xlua_panic (lua_State *L)
+{
+	// XXX: we might be able to do something better
+	print_fatal ("Lua panicked: %s", lua_tostring (L, -1));
+	lua_close (L);
+	exit (EXIT_FAILURE);
+	return 0;
+}
+
+// --- Unit wrapper ------------------------------------------------------------
+
+struct unit_wrapper
+{
+	struct unit *unit;                  ///< The underlying unit
+};
+
+#define UNIT_METATABLE "unit"
+
+static int
+xlua_unit_get_address (lua_State *L)
+{
+	struct unit_wrapper *data = luaL_checkudata (L, 1, UNIT_METATABLE);
+	lua_pushstring (L, g_data.api->unit_get_address (data->unit));
+	return 1;
+}
+
+static int
+xlua_unit_write (lua_State *L)
+{
+	struct unit_wrapper *data = luaL_checkudata (L, 1, UNIT_METATABLE);
+	size_t buffer_len;
+	const char *buffer = luaL_checklstring (L, 2, &buffer_len);
+	lua_pushinteger (L, g_data.api->unit_write
+		(data->unit, buffer, buffer_len));
+	return 1;
+}
+
+static int
+xlua_unit_set_success (lua_State *L)
+{
+	struct unit_wrapper *data = luaL_checkudata (L, 1, UNIT_METATABLE);
+	bool success = lua_toboolean (L, 2);
+	g_data.api->unit_set_success (data->unit, success);
+	return 0;
+}
+
+static int
+xlua_unit_add_info (lua_State *L)
+{
+	struct unit_wrapper *data = luaL_checkudata (L, 1, UNIT_METATABLE);
+	const char *info = luaL_checkstring (L, 2);
+	g_data.api->unit_add_info (data->unit, info);
+	return 0;
+}
+
+static int
+xlua_unit_abort (lua_State *L)
+{
+	struct unit_wrapper *data = luaL_checkudata (L, 1, UNIT_METATABLE);
+	g_data.api->unit_abort (data->unit);
+	return 0;
+}
+
+static int
+xlua_unit_destroy (lua_State *L)
+{
+	// TODO: when creating the wrapper object, increase the reference
+	//   count for the unit and decrease it in here again.  If we don't do
+	//   this, the Lua plugin may cause a use-after-free error.
+	(void) L;
+	return 0;
+}
+
+static luaL_Reg xlua_unit_table[] =
+{
+	{ "get_address",      xlua_unit_get_address },
+	{ "write",            xlua_unit_write       },
+	{ "set_success",      xlua_unit_set_success },
+	{ "add_info",         xlua_unit_add_info    },
+	{ "abort",            xlua_unit_abort       },
+	{ "__gc",             xlua_unit_destroy     },
+	{ NULL,               NULL                  }
+};
+
+// --- Scan wrapper ------------------------------------------------------------
+
+struct service_data
+{
+	struct service *service;            ///< The corresponding service
+	lua_State *L;                       ///< Lua state
+	int new_scan_cb_ref;                ///< Reference to "new_scan" callback
+};
+
+struct scan_data
+{
+	struct service *service;            ///< The corresponding service
+	struct unit *unit;                  ///< The corresponding unit
+	lua_State *L;                       ///< Lua state
+	int scan_ref;                       ///< Reference to scan data in Lua
+};
+
+static void *
+scan_init (struct service *self, struct unit *unit)
+{
+	struct service_data *service = self->user_data;
+	lua_geti (service->L, LUA_REGISTRYINDEX, service->new_scan_cb_ref);
+
+	// Wrap the unit in Lua userdata so that Lua code can use it
+	struct unit_wrapper *wrapper =
+		lua_newuserdata (service->L, sizeof *wrapper);
+	wrapper->unit = unit;
+	luaL_setmetatable (service->L, UNIT_METATABLE);
+
+	// Ask for a Lua object (table) to use for protocol detection
+	if (lua_pcall (service->L, 1, 1, 0))
+	{
+		print_error ("Lua: service `%s': new_scan: %s",
+			service->service->name, lua_tostring (service->L, -1));
+		lua_pop (service->L, 1);
+		return NULL;
+	}
+
+	if (!lua_istable (service->L, -1))
+	{
+		print_error ("Lua: service `%s': new_scan must return a table",
+			service->service->name);
+		return NULL;
+	}
+
+	// Return a scan handle
+	struct scan_data *data = xmalloc (sizeof *data);
+	data->service = self;
+	data->L = service->L;
+	data->scan_ref = luaL_ref (service->L, LUA_REGISTRYINDEX);
+	data->unit = unit;
+	return data;
+}
+
+static void
+scan_free (void *handle)
+{
+	if (!handle)
+		return;
+
+	struct scan_data *data = handle;
+	luaL_unref (data->L, LUA_REGISTRYINDEX, data->scan_ref);
+	free (handle);
+}
+
+static void
+handle_scan_method_failure (struct scan_data *data)
+{
+	print_error ("Lua: service `%s': %s", data->service->name,
+		lua_tostring (data->L, -1));
+	g_data.api->unit_abort (data->unit);
+	lua_pop (data->L, 1);
+}
+
+static bool
+prepare_scan_method (struct scan_data *data, const char *method_name)
+{
+	if (!data)
+		return false;
+
+	lua_geti (data->L, LUA_REGISTRYINDEX, data->scan_ref);
+	lua_getfield (data->L, -1, method_name);
+	if (lua_isnil (data->L, -1))
+	{
+		lua_pop (data->L, 2);
+		return false;
+	}
+
+	if (!lua_isfunction (data->L, -1))
+	{
+		lua_pop (data->L, 2);
+		lua_pushfstring (data->L, "`%s' must be a function", method_name);
+		handle_scan_method_failure (data);
+		return false;
+	}
+
+	// Swap the first two values on the stack, so that the function
+	// is first and the object we're calling it on is second
+	lua_insert (data->L, -2);
+	return true;
+}
+
+static void
+on_data (void *handle, const void *network_data, size_t len)
+{
+	struct scan_data *data = handle;
+	if (!prepare_scan_method (data, "on_data"))
+		return;
+
+	lua_pushlstring (data->L, network_data, len);
+	if (lua_pcall (data->L, 2, 0, 0))
+		handle_scan_method_failure (data);
+}
+
+static void
+on_eof (void *handle)
+{
+	struct scan_data *data = handle;
+	if (!prepare_scan_method (data, "on_eof"))
+		return;
+	if (lua_pcall (data->L, 1, 0, 0))
+		handle_scan_method_failure (data);
+}
+
+static void
+on_error (void *handle)
+{
+	struct scan_data *data = handle;
+	if (!prepare_scan_method (data, "on_error"))
+		return;
+	if (lua_pcall (data->L, 1, 0, 0))
+		handle_scan_method_failure (data);
+}
+
+static void
+on_aborted (void *handle)
+{
+	struct scan_data *data = handle;
+	if (!prepare_scan_method (data, "on_aborted"))
+		return;
+	if (lua_pcall (data->L, 1, 0, 0))
+		handle_scan_method_failure (data);
+}
+
+// --- Top-level interface -----------------------------------------------------
+
+static int
+xlua_register_service (lua_State *L)
+{
+	// Validate and decode the arguments
+	luaL_checktype (L, 1, LUA_TTABLE);
+
+	lua_getfield (L, 1, "name");
+	if (!lua_isstring (L, -1))
+		return luaL_error (L, "service registration failed: "
+			"invalid or missing `%s'", "name");
+	const char *name = lua_tostring (L, -1);
+	lua_pop (L, 1);
+
+	lua_getfield (L, 1, "flags");
+	lua_Unsigned flags;
+	if (lua_isnil (L, -1))
+		flags = 0;
+	else if (lua_isinteger (L, -1))
+		flags = lua_tointeger (L, -1);
+	else
+		return luaL_error (L, "service registration failed: "
+			"invalid or missing `%s'", "flags");
+	lua_pop (L, 1);
+
+	lua_getfield (L, 1, "new_scan");
+	if (!lua_isfunction (L, -1))
+		return luaL_error (L, "service registration failed: "
+			"invalid or missing `%s'", "new_scan");
+
+	// Reference the "new_scan" function for later use
+	struct service_data *data = xcalloc (1, sizeof *data);
+	data->L = L;
+	data->new_scan_cb_ref = luaL_ref (L, LUA_REGISTRYINDEX);
+
+	// Register a new service that proxies calls to Lua code
+	struct service *s = data->service = xcalloc (1, sizeof *s);
+	s->name       = xstrdup (name);
+	s->flags      = flags;
+	s->user_data  = data;
+
+	s->scan_init  = scan_init;
+	s->scan_free  = scan_free;
+	s->on_data    = on_data;
+	s->on_eof     = on_eof;
+	s->on_error   = on_error;
+	s->on_aborted = on_aborted;
+
+	g_data.api->register_service (g_data.ctx, s);
+	return 0;
+}
+
+static int
+xlua_get_config (lua_State *L)
+{
+	const char *key = luaL_checkstring (L, 1);
+	lua_pushstring (L, g_data.api->get_config (g_data.ctx, key));
+	return 0;
+}
+
+static luaL_Reg xlua_library[] =
+{
+	{ "register_service", xlua_register_service },
+	{ "get_config",       xlua_get_config       },
+	{ NULL,               NULL                  }
+};
+
+static bool
+load_one_plugin (lua_State *L, const char *name, const char *path)
+{
+	int ret;
+	if (!(ret = luaL_loadfile (L, path))
+	 && !(ret = lua_pcall (L, 0, 0, 0)))
+		return true;
+
+	print_error ("Lua: could not load `%s': %s", name, lua_tostring (L, -1));
+	lua_pop (L, 1);
+	return false;
+}
+
+static bool
+initialize (void *ctx, struct plugin_api *api)
+{
+	g_data = (struct plugin_data) { .ctx = ctx, .api = api };
+
+	if (sizeof (lua_Integer) < 8)
+	{
+		print_error ("%s: %s", "Lua",
+			"at least 64-bit Lua integers are required");
+		return false;
+	}
+
+	const char *plugin_dir = api->get_config (ctx, "plugin_dir");
+	if (!plugin_dir)
+	{
+		print_fatal ("%s: %s", "Lua", "no plugin directory defined");
+		return false;
+	}
+
+	DIR *dir = opendir (plugin_dir);
+	if (!dir)
+	{
+		print_fatal ("%s: %s: %s", "Lua",
+			"cannot open plugin directory", strerror (errno));
+		return false;
+	}
+
+	bool success = false;
+	lua_State *L;
+	if (!(L = lua_newstate (xlua_alloc, NULL)))
+	{
+		print_fatal ("%s: %s", "Lua", "initialization failed");
+		goto end;
+	}
+
+	lua_atpanic (L, xlua_panic);
+	luaL_openlibs (L);
+
+	// Register the ponymap library
+	luaL_newlib (L, xlua_library);
+	lua_pushinteger (L, SERVICE_SUPPORTS_TLS);
+	lua_setfield (L, -2, "SERVICE_SUPPORTS_TLS");
+	lua_setglobal (L, PROGRAM_NAME);
+
+	// Create a metatable for the units
+	luaL_newmetatable (L, UNIT_METATABLE);
+	lua_pushvalue (L, -1);
+	lua_setfield (L, -2, "__index");
+	luaL_setfuncs (L, xlua_unit_table, 0);
+
+	struct dirent buf, *iter;
+	while (true)
+	{
+		if (readdir_r (dir, &buf, &iter))
+		{
+			print_fatal ("%s: %s: %s", "Lua", "readdir_r", strerror (errno));
+			break;
+		}
+		if (!iter)
+		{
+			success = true;
+			break;
+		}
+
+		char *dot = strrchr (iter->d_name, '.');
+		if (!dot || strcmp (dot, ".lua"))
+			continue;
+
+		char *path = xstrdup_printf ("%s/%s", plugin_dir, iter->d_name);
+		(void) load_one_plugin (L, iter->d_name, path);
+		free (path);
+	}
+
+end:
+	closedir (dir);
+	return success;
+}
+
+struct plugin_info ponymap_plugin_info =
+{
+	.api_version  = API_VERSION,
+	.initialize   = initialize
+};
diff --git a/plugins/socks.lua b/plugins/socks.lua
new file mode 100644
index 0000000..a03f72d
--- /dev/null
+++ b/plugins/socks.lua
@@ -0,0 +1,99 @@
+--
+-- socks.lua: SOCKS service detection plugin
+--
+-- Copyright (c) 2015, Přemysl Janouch <p.janouch@gmail.com>
+-- All rights reserved.
+--
+-- Permission to use, copy, modify, and/or distribute this software for any
+-- purpose with or without fee is hereby granted, provided that the above
+-- copyright notice and this permission notice appear in all copies.
+--
+-- THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+-- WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+-- MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
+-- SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+-- WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
+-- OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
+-- CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+--
+
+-- This plugin serves as an example of how to write Lua plugins for ponymap
+
+-- SOCKS 4
+
+local Socks4 = {}
+Socks4.__index = Socks4
+
+function Socks4.new (unit)
+	unit:write (string.pack ("> I1 I1 I2 I1 I1 I1 I1 z",
+		4, 1, 80, 127, 0, 0, 1, ""))
+	return setmetatable ({ unit = unit, buf = "" }, Socks4)
+end
+
+function Socks4:on_data (data)
+	self.buf = self.buf .. data
+	if #self.buf >= 8 then
+		null, code = string.unpack ("> I1 I1", self.buf)
+		if null == 0 and code >= 90 and code <= 93 then
+			self.unit:set_success (true)
+		end
+		self.unit:abort ()
+	end
+end
+
+-- SOCKS 4A
+
+local Socks4A = {}
+Socks4A.__index = Socks4A
+
+function Socks4A.new (unit)
+	unit:write (string.pack ("> I1 I1 I2 I4 z z",
+		4, 1, 80, 1, "", "google.com"))
+	return setmetatable ({ unit = unit, buf = "" }, Socks4A)
+end
+
+Socks4A.on_data = Socks4.on_data
+
+-- SOCKS 5
+
+local Socks5 = {}
+Socks5.__index = Socks5
+
+function Socks5.new (unit)
+	unit:write (string.pack ("> I1 I1 I1", 5, 1, 0))
+	return setmetatable ({ unit = unit, buf = "" }, Socks5)
+end
+
+function Socks5:on_data (data)
+	self.buf = self.buf .. data
+	if #self.buf >= 2 then
+		version, result = string.unpack ("> I1 I1", self.buf)
+		if version == 5 and (result == 0 or result == 255) then
+			if result == 0 then
+				self.unit:add_info ("no authentication required")
+			end
+			self.unit:set_success (true)
+		end
+		self.unit:abort ()
+	end
+end
+
+-- Register everything
+
+ponymap.register_service ({
+	name = "SOCKS4",
+	flags = 0,
+	new_scan = Socks4.new
+})
+
+ponymap.register_service ({
+	name = "SOCKS4A",
+	flags = 0,
+	new_scan = Socks4A.new
+})
+
+ponymap.register_service ({
+	name = "SOCKS5",
+	flags = 0,
+	new_scan = Socks5.new
+})
diff --git a/plugins/ssh.c b/plugins/ssh.c
index 8a1cc3a..68f8d52 100644
--- a/plugins/ssh.c
+++ b/plugins/ssh.c
@@ -37,8 +37,10 @@ struct scan_data
 };
 
 static void *
-scan_init (struct unit *u)
+scan_init (struct service *service, struct unit *u)
 {
+	(void) service;
+
 	struct scan_data *scan = xcalloc (1, sizeof *scan);
 	str_init (&scan->input);
 	scan->u = u;
diff --git a/ponymap.c b/ponymap.c
index c1cca02..1ab521a 100644
--- a/ponymap.c
+++ b/ponymap.c
@@ -623,7 +623,7 @@ unit_start_scan (struct unit *u)
 	u->scan_started = true;
 	poller_timer_set (&u->timeout_event, u->target->ctx->scan_timeout * 1000);
 
-	u->service_data = u->service->scan_init (u);
+	u->service_data = u->service->scan_init (u->service, u);
 	u->fd_event.dispatcher = (poller_fd_fn) on_unit_ready;
 	unit_update_poller (u, NULL);
 }
@@ -792,6 +792,13 @@ plugin_api_register_service (void *app_context, struct service *info)
 }
 
 static const char *
+plugin_api_get_config (void *app_context, const char *key)
+{
+	struct app_context *ctx = app_context;
+	return str_map_find (&ctx->config, key);
+}
+
+static const char *
 plugin_api_unit_get_address (struct unit *u)
 {
 	return u->target->ip_string;
@@ -828,6 +835,7 @@ plugin_api_unit_abort (struct unit *u)
 static struct plugin_api g_plugin_vtable =
 {
 	.register_service  = plugin_api_register_service,
+	.get_config        = plugin_api_get_config,
 	.unit_get_address  = plugin_api_unit_get_address,
 	.unit_write        = plugin_api_unit_write,
 	.unit_set_success  = plugin_api_unit_set_success,