1 files changed, 179 insertions, 96 deletions

diff --git a/src/rspamadm/lua_repl.c b/src/rspamadm/lua_repl.c
index 4cd68185d..b7fd9a4d3 100644
--- a/src/rspamadm/lua_repl.c
+++ b/src/rspamadm/lua_repl.c
@@ -24,7 +24,9 @@
 #include "lua/lua_thread_pool.h"
 #include "message.h"
 #include "unix-std.h"
-#include "linenoise.h"
+#ifdef WITH_LUA_REPL
+#include "replxx.h"
+#endif
 #include "worker_util.h"
 #ifdef WITH_LUAJIT
 #include <luajit.h>
@@ -38,11 +40,14 @@ static guint max_history = 2000;
 static gchar *serve = NULL;
 static gchar *exec_line = NULL;
 static gint batch = -1;
-static gboolean per_line = FALSE;
 extern struct rspamd_async_session *rspamadm_session;
 
 static const char *default_history_file = ".rspamd_repl.hist";
 
+#ifdef WITH_LUA_REPL
+static Replxx *rx_instance = NULL;
+#endif
+
 #ifdef WITH_LUAJIT
 #define MAIN_PROMPT LUAJIT_VERSION "> "
 #else
@@ -119,8 +124,6 @@ static GOptionEntry entries[] = {
 				"Serve http lua server", NULL},
 		{"batch", 'b', 0, G_OPTION_ARG_NONE, &batch,
 				"Batch execution mode", NULL},
-		{"per-line", 'p', 0, G_OPTION_ARG_NONE, &per_line,
-				"Pass each line of input to the specified lua script", NULL},
 		{"exec", 'e', 0, G_OPTION_ARG_STRING, &exec_line,
 				"Execute specified script", NULL},
 		{"args", 'a', 0, G_OPTION_ARG_STRING_ARRAY, &lua_args,
@@ -222,15 +225,12 @@ rspamadm_lua_load_script (lua_State *L, const gchar *path)
 		return FALSE;
 	}
 
-	if (!per_line) {
-
-		if (lua_repl_thread_call (thread, 0, (void *)path, lua_thread_str_error_cb) != 0) {
-			return FALSE;
-		}
-
-		lua_settop (L, 0);
+	if (lua_repl_thread_call (thread, 0, (void *)path, lua_thread_str_error_cb) != 0) {
+		return FALSE;
 	}
 
+	lua_settop (L, 0);
+
 	return TRUE;
 }
 
@@ -268,24 +268,22 @@ rspamadm_exec_input (lua_State *L, const gchar *input)
 
 	g_string_free (tb, TRUE);
 
-	if (!per_line) {
 
-		top = lua_gettop (L);
+	top = lua_gettop (L);
 
-		if (lua_repl_thread_call (thread, 0, NULL, NULL) == 0) {
-			/* Print output */
-			for (i = top; i <= lua_gettop (L); i++) {
-				if (lua_isfunction (L, i)) {
-					lua_pushvalue (L, i);
-					cbref = luaL_ref (L, LUA_REGISTRYINDEX);
+	if (lua_repl_thread_call (thread, 0, NULL, NULL) == 0) {
+		/* Print output */
+		for (i = top; i <= lua_gettop (L); i++) {
+			if (lua_isfunction (L, i)) {
+				lua_pushvalue (L, i);
+				cbref = luaL_ref (L, LUA_REGISTRYINDEX);
 
-					rspamd_printf ("local function: %d\n", cbref);
-				} else {
-					memset (&tr, 0, sizeof (tr));
-					lua_logger_out_type (L, i, outbuf, sizeof (outbuf), &tr,
-							LUA_ESCAPE_UNPRINTABLE);
-					rspamd_printf ("%s\n", outbuf);
-				}
+				rspamd_printf ("local function: %d\n", cbref);
+			} else {
+				memset (&tr, 0, sizeof (tr));
+				lua_logger_out_type (L, i, outbuf, sizeof (outbuf) - 1, &tr,
+						LUA_ESCAPE_UNPRINTABLE);
+				rspamd_printf ("%s\n", outbuf);
 			}
 		}
 	}
@@ -510,17 +508,141 @@ rspamadm_lua_try_dot_command (lua_State *L, const gchar *input)
 	return FALSE;
 }
 
+#ifdef WITH_LUA_REPL
+static gint lex_ref_idx = -1;
+
 static void
-rspamadm_lua_run_repl (lua_State *L)
+lua_syntax_highlighter (const char *str, ReplxxColor *colours, int size, void *ud)
+{
+	lua_State *L = (lua_State *)ud;
+
+	if (lex_ref_idx == -1) {
+		if (!rspamd_lua_require_function (L, "lua_lexer", "lex_to_table")) {
+			fprintf (stderr, "cannot require lua_lexer!\n");
+
+			exit (EXIT_FAILURE);
+		}
+
+		lex_ref_idx = luaL_ref (L, LUA_REGISTRYINDEX);
+	}
+
+	lua_rawgeti (L, LUA_REGISTRYINDEX, lex_ref_idx);
+	lua_pushstring (L, str);
+
+	if (lua_pcall (L, 1, 1, 0) != 0) {
+		fprintf (stderr, "cannot lex a string!\n");
+	}
+	else {
+		/* Process what we have after lexing */
+		gsize nelts = rspamd_lua_table_size (L, -1);
+
+		for (gsize i = 0; i < nelts; i ++) {
+			/*
+			 * Indexes in the table:
+			 * 1 - type of element (string)
+			 * 2 - text (string)
+			 * 3 - line num (int), always 1...
+			 * 4 - column num (must be less than size)
+			 */
+			const gchar *what, *text;
+			gsize column, tlen, cur_top, elt_pos;
+			ReplxxColor elt_color = REPLXX_COLOR_DEFAULT;
+
+			cur_top = lua_gettop (L);
+			lua_rawgeti (L, -1, i + 1);
+			elt_pos = lua_gettop (L);
+			lua_rawgeti (L, elt_pos, 1);
+			what = lua_tostring (L, -1);
+			lua_rawgeti (L, elt_pos, 2);
+			text = lua_tolstring (L, -1, &tlen);
+			lua_rawgeti (L, elt_pos, 4);
+			column = lua_tointeger (L, -1);
+
+			g_assert (column > 0);
+			column --; /* Start from 0 */
+
+			if (column + tlen > size) {
+				/* Likely utf8 case, too complicated to match */
+				lua_settop (L, cur_top);
+				continue;
+			}
+
+			/* Check what and adjust color */
+			if (strcmp (what, "identifier") == 0) {
+				elt_color = REPLXX_COLOR_NORMAL;
+			}
+			else if (strcmp (what, "number") == 0) {
+				elt_color = REPLXX_COLOR_BLUE;
+			}
+			else if (strcmp (what, "string") == 0) {
+				elt_color = REPLXX_COLOR_GREEN;
+			}
+			else if (strcmp (what, "keyword") == 0) {
+				elt_color = REPLXX_COLOR_WHITE;
+			}
+			else if (strcmp (what, "constant") == 0) {
+				elt_color = REPLXX_COLOR_WHITE;
+			}
+			else if (strcmp (what, "operator") == 0) {
+				elt_color = REPLXX_COLOR_CYAN;
+			}
+			else if (strcmp (what, "comment") == 0) {
+				elt_color = REPLXX_COLOR_BRIGHTGREEN;
+			}
+			else if (strcmp (what, "error") == 0) {
+				elt_color = REPLXX_COLOR_ERROR;
+			}
+
+			for (gsize j = column; j < column + tlen; j ++) {
+				colours[j] = elt_color;
+			}
+
+			/* Restore stack */
+			lua_settop (L, cur_top);
+		}
+	}
+
+	lua_settop (L, 0);
+}
+#endif
+
+static void
+rspamadm_lua_run_repl (lua_State *L, bool is_batch)
 {
 	gchar *input;
 	gboolean is_multiline = FALSE;
 	GString *tb = NULL;
-	guint i;
+	gsize i;
 
 	for (;;) {
+#ifndef WITH_LUA_REPL
+		size_t linecap = 0;
+		ssize_t linelen;
+
+		fprintf (stdout, "%s ", MAIN_PROMPT);
+
+		linelen = getline (&input, &linecap, stdin);
+
+		if (linelen > 0) {
+			if (input[linelen - 1] == '\n') {
+				linelen --;
+			}
+
+			rspamadm_exec_input (L, input);
+		}
+		else {
+			break;
+		}
+
+		lua_settop (L, 0);
+#else
+		if (!is_batch) {
+			replxx_set_highlighter_callback (rx_instance, lua_syntax_highlighter,
+					L);
+		}
+
 		if (!is_multiline) {
-			input = linenoise (MAIN_PROMPT);
+			input = (gchar *)replxx_input (rx_instance, MAIN_PROMPT);
 
 			if (input == NULL) {
 				return;
@@ -528,26 +650,27 @@ rspamadm_lua_run_repl (lua_State *L)
 
 			if (input[0] == '.') {
 				if (rspamadm_lua_try_dot_command (L, input)) {
-					linenoiseHistoryAdd (input);
-					linenoiseFree (input);
+					if (!is_batch) {
+						replxx_history_add (rx_instance, input);
+					}
 					continue;
 				}
 			}
 
 			if (strcmp (input, "{{") == 0) {
 				is_multiline = TRUE;
-				linenoiseFree (input);
 				tb = g_string_sized_new (8192);
 				continue;
 			}
 
 			rspamadm_exec_input (L, input);
-			linenoiseHistoryAdd (input);
-			linenoiseFree (input);
+			if (!is_batch) {
+				replxx_history_add (rx_instance, input);
+			}
 			lua_settop (L, 0);
 		}
 		else {
-			input = linenoise (MULTILINE_PROMPT);
+			input = (gchar *)replxx_input (rx_instance, MULTILINE_PROMPT);
 
 			if (input == NULL) {
 				g_string_free (tb, TRUE);
@@ -556,7 +679,6 @@ rspamadm_lua_run_repl (lua_State *L)
 
 			if (strcmp (input, "}}") == 0) {
 				is_multiline = FALSE;
-				linenoiseFree (input);
 				rspamadm_exec_input (L, tb->str);
 
 				/* Replace \n with ' ' for sanity */
@@ -566,15 +688,17 @@ rspamadm_lua_run_repl (lua_State *L)
 					}
 				}
 
-				linenoiseHistoryAdd (tb->str);
+				if (!is_batch) {
+					replxx_history_add (rx_instance, tb->str);
+				}
 				g_string_free (tb, TRUE);
 			}
 			else {
 				g_string_append (tb, input);
 				g_string_append (tb, " \n");
-				linenoiseFree (input);
 			}
 		}
+#endif
 	}
 }
 
@@ -745,9 +869,12 @@ rspamadm_lua (gint argc, gchar **argv, const struct rspamadm_command *cmd)
 	if (!g_option_context_parse (context, &argc, &argv, &error)) {
 		fprintf (stderr, "option parsing failed: %s\n", error->message);
 		g_error_free (error);
+		g_option_context_free (context);
 		exit (1);
 	}
 
+	g_option_context_free (context);
+
 	if (batch == -1) {
 		if (isatty (STDIN_FILENO)) {
 			batch = 0;
@@ -876,63 +1003,19 @@ rspamadm_lua (gint argc, gchar **argv, const struct rspamadm_command *cmd)
 		g_hash_table_insert (cmds_hash, (gpointer)cmds[i].name, &cmds[i]);
 	}
 
-	if (per_line) {
-		GIOChannel *in;
-		GString *buf;
-		gsize end_pos;
-		GIOStatus ret;
-		gint old_top;
-		GError *err = NULL;
-
-		in = g_io_channel_unix_new (STDIN_FILENO);
-		buf = g_string_sized_new (BUFSIZ);
 
-again:
-		while ((ret = g_io_channel_read_line_string (in, buf, &end_pos, &err)) ==
-				G_IO_STATUS_NORMAL) {
-			old_top = lua_gettop (L);
-			lua_pushvalue (L, -1);
-			lua_pushlstring (L, buf->str, MIN (buf->len, end_pos));
-			lua_setglobal (L, "input");
-
-			struct thread_entry *thread = lua_thread_pool_get_for_config (rspamd_main->cfg);
-			L = thread->lua_state;
-
-			lua_repl_thread_call (thread, 0, NULL, NULL);
-
-			lua_settop (L, old_top);
-		}
-
-		if (ret == G_IO_STATUS_AGAIN) {
-			goto again;
-		}
-
-		g_string_free (buf, TRUE);
-		g_io_channel_shutdown (in, FALSE, NULL);
-
-		if (ret == G_IO_STATUS_EOF) {
-			if (err) {
-				g_error_free (err);
-			}
-		}
-		else {
-			rspamd_fprintf (stderr, "IO error: %e\n", err);
-
-			if (err) {
-				g_error_free (err);
-			}
-
-			exit (-errno);
-		}
-	}
-	else {
-		if (!batch) {
-			linenoiseHistorySetMaxLen (max_history);
-			linenoiseHistoryLoad (histfile);
-			rspamadm_lua_run_repl (L);
-			linenoiseHistorySave (histfile);
-		} else {
-			rspamadm_lua_run_repl (L);
-		}
-	}
+#ifdef WITH_LUA_REPL
+	rx_instance = replxx_init ();
+#endif
+	if (!batch) {
+		replxx_set_max_history_size (rx_instance, max_history);
+		replxx_history_load (rx_instance, histfile);
+		rspamadm_lua_run_repl (L, false);
+		replxx_history_save (rx_instance, histfile);
+	} else {
+		rspamadm_lua_run_repl (L, true);
+	}
+#ifdef WITH_LUA_REPL
+	replxx_end (rx_instance);
+#endif
 }