[Minor] Lua_text: Allow comparison methods with plain strings

rspamd · Oct 26, 2020 · eb5f3f5 · eb5f3f5
1 parent 797a8e3
commit eb5f3f5
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 2 deletions.
diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h
@@ -101,6 +101,7 @@ struct rspamd_lua_ip {
 #define RSPAMD_TEXT_FLAG_MMAPED (1u << 1u)
 #define RSPAMD_TEXT_FLAG_WIPE (1u << 2u)
 #define RSPAMD_TEXT_FLAG_SYSMALLOC (1u << 3u)
+#define RSPAMD_TEXT_FLAG_FAKE (1u << 4u)
 struct rspamd_lua_text {
 	const gchar *start;
 	guint len;
@@ -254,6 +255,14 @@ void rspamd_lua_task_push (lua_State *L, struct rspamd_task *task);
 struct rspamd_lua_ip *lua_check_ip (lua_State *L, gint pos);
 
 struct rspamd_lua_text *lua_check_text (lua_State *L, gint pos);
+/**
+ * Checks for a text or a string. In case of string a pointer to static structure is returned.
+ * So it should not be reused or placed to Lua stack anyhow!
+ * @param L
+ * @param pos
+ * @return
+ */
+struct rspamd_lua_text *lua_check_text_or_string (lua_State *L, gint pos);
 /* Creates and *pushes* new rspamd text, data is copied if  RSPAMD_TEXT_FLAG_OWN is in flags*/
 struct rspamd_lua_text *lua_new_text (lua_State *L, const gchar *start,
 		gsize len, gboolean own);

diff --git a/src/lua/lua_text.c b/src/lua/lua_text.c
@@ -249,6 +249,29 @@ lua_check_text (lua_State * L, gint pos)
 	return ud ? (struct rspamd_lua_text *)ud : NULL;
 }
 
+struct rspamd_lua_text *
+lua_check_text_or_string (lua_State * L, gint pos)
+{
+	gint pos_type = lua_type (L, pos);
+
+	if (pos_type == LUA_TUSERDATA) {
+		void *ud = rspamd_lua_check_udata (L, pos, "rspamd{text}");
+		luaL_argcheck (L, ud != NULL, pos, "'text' expected");
+		return ud ? (struct rspamd_lua_text *) ud : NULL;
+	}
+	else if (pos_type == LUA_TSTRING) {
+		/* Fake static lua_text */
+		static struct rspamd_lua_text fake_text;
+
+		fake_text.start = lua_tolstring (L, pos, &fake_text.len);
+		fake_text.flags = RSPAMD_TEXT_FLAG_FAKE;
+
+		return &fake_text;
+	}
+
+	return NULL;
+}
+
 struct rspamd_lua_text *
 lua_new_text (lua_State *L, const gchar *start, gsize len, gboolean own)
 {
@@ -1082,6 +1105,8 @@ lua_text_gc (lua_State *L)
 	struct rspamd_lua_text *t = lua_check_text (L, 1);
 
 	if (t != NULL) {
+		g_assert (!(t->flags & RSPAMD_TEXT_FLAG_FAKE));
+
 		if (t->flags & RSPAMD_TEXT_FLAG_OWN) {
 			if (t->flags & RSPAMD_TEXT_FLAG_WIPE) {
 				rspamd_explicit_memzero ((guchar *)t->start, t->len);
@@ -1110,7 +1135,7 @@ lua_text_eq (lua_State *L)
 {
 	LUA_TRACE_POINT;
 	struct rspamd_lua_text *t1 = lua_check_text (L, 1),
-			*t2 = lua_check_text (L, 2);
+			*t2 = lua_check_text_or_string (L, 2);
 
 	if (t1->len == t2->len) {
 		lua_pushboolean (L, memcmp (t1->start, t2->start, t1->len) == 0);
@@ -1127,7 +1152,7 @@ lua_text_lt (lua_State *L)
 {
 	LUA_TRACE_POINT;
 	struct rspamd_lua_text *t1 = lua_check_text (L, 1),
-			*t2 = lua_check_text (L, 2);
+			*t2 = lua_check_text_or_string (L, 2);
 
 	if (t1 && t2) {
 		if (t1->len == t2->len) {