Add Lua heuristic dissector support

This adds the ability for Lua scripts to register heuristic dissectors
for any protocol that has registered a heuristic dissector list, such
as UDP, TCP, and ~50 others. The Lua function can also establish a
conversation tied to its Proto dissector, to avoid having to check the
heuristics for the same flow. The example dissector in the testsuite
has also been enhanced to include a heuristic dissector, to verify
the functionality and provide an example implementation.

Change-Id: Ie232602779f43d3418fe8db09c61d5fc0b59597a
Reviewed-on: https://code.wireshark.org/review/576
Reviewed-by: Anders Broman <a.broman58@gmail.com>

6 files changed, 293 insertions, 6 deletions

diff --git a/epan/wslua/init_wslua.c b/epan/wslua/init_wslua.c
index 2d778b9b68..c37602f82f 100644
--- a/epan/wslua/init_wslua.c
+++ b/epan/wslua/init_wslua.c
@@ -47,10 +47,13 @@ static wslua_plugin *wslua_plugin_list = NULL;
 
 static lua_State* L = NULL;
 
+/* XXX: global variables? Really?? Yuck. These could be done differently,
+   using the Lua registry */
 packet_info* lua_pinfo;
 struct _wslua_treeitem* lua_tree;
 tvbuff_t* lua_tvb;
-int lua_dissectors_table_ref;
+int lua_dissectors_table_ref = LUA_NOREF;
+int lua_heur_dissectors_table_ref = LUA_NOREF;
 
 static int proto_lua = -1;
 static expert_field ei_lua_error = EI_INIT;
@@ -133,6 +136,95 @@ int dissect_lua(tvbuff_t* tvb, packet_info* pinfo, proto_tree* tree, void* data
 
 }
 
+/** Type of a heuristic dissector, used in heur_dissector_add().
+ *
+ * @param tvb the tvbuff with the (remaining) packet data
+ * @param pinfo the packet info of this packet (additional info)
+ * @param tree the protocol tree to be build or NULL
+ * @return TRUE if the packet was recognized by the sub-dissector (stop dissection here)
+ */
+gboolean heur_dissect_lua(tvbuff_t* tvb, packet_info* pinfo, proto_tree* tree, void* data _U_) {
+    gboolean result = FALSE;
+    lua_tvb = tvb;
+    lua_pinfo = pinfo;
+
+    if (!tvb || !pinfo || !pinfo->heur_list_name || !pinfo->current_proto) {
+        report_failure("internal error in heur_dissect_lua: NULL packet info");
+        return FALSE;
+    }
+
+    /* heuristic functions are stored in a table in the registry; the registry has a
+     * table at reference lua_heur_dissectors_table_ref, and that table has keys for
+     * the heuristic listname (e.g., "udp", "tcp", etc.), and that key's value is a
+     * table of keys of the Proto->name, and their value is the function.
+     * So it's like registry[table_ref][heur_list_name][proto_name] = func
+     */
+
+    lua_settop(L,0);
+
+    /* get the table of all lua heuristic dissector lists */
+    lua_rawgeti(L, LUA_REGISTRYINDEX, lua_heur_dissectors_table_ref);
+
+    /* get the table inside that, for the lua heuristic dissectors of the requested heur list */
+    if (!wslua_get_table(L, -1, pinfo->heur_list_name)) {
+        /* this shouldn't happen */
+        lua_settop(L,0);
+        report_failure("internal error in heur_dissect_lua: no %s heur list table", pinfo->heur_list_name);
+        return FALSE;
+    }
+
+    /* get the table inside that, for the specific lua heuristic dissector */
+    if (!wslua_get_field(L,-1,pinfo->current_proto)) {
+        /* this shouldn't happen */
+        lua_settop(L,0);
+        report_failure("internal error in heur_dissect_lua: no %s heuristic dissector for list %s",
+                        pinfo->current_proto, pinfo->heur_list_name);
+        return FALSE;
+    }
+
+    /* remove the table of all lists (the one in the registry) */
+    lua_remove(L,1);
+    /* remove the heur_list_name heur list table */
+    lua_remove(L,1);
+
+    if (!lua_isfunction(L,-1)) {
+        /* this shouldn't happen */
+        lua_settop(L,0);
+        report_failure("internal error in heur_dissect_lua: %s heuristic dissector is not a function", pinfo->current_proto);
+        return FALSE;
+    }
+
+    lua_tree = (struct _wslua_treeitem *)g_malloc(sizeof(struct _wslua_treeitem));
+    lua_tree->tree = tree;
+    lua_tree->item = proto_tree_add_text(tree,tvb,0,0,"lua fake item");
+    lua_tree->expired = FALSE;
+    PROTO_ITEM_SET_HIDDEN(lua_tree->item);
+
+    push_Tvb(L,tvb);
+    push_Pinfo(L,pinfo);
+    push_TreeItem(L,lua_tree);
+
+    if  ( lua_pcall(L,3,1,0) ) {
+        report_failure(" error calling %s heuristic dissector: %s", pinfo->current_proto, lua_tostring(L,-1));
+        lua_settop(L,0);
+    } else {
+        if (lua_isboolean(L, -1) || lua_isnil(L, -1)) {
+            result = lua_toboolean(L, -1);
+        } else {
+            report_failure(" invalid return value from Lua %s heuristic dissector", pinfo->current_proto);
+        }
+        lua_pop(L, 1);
+    }
+
+    register_frame_end_routine(pinfo, lua_frame_end);
+
+    lua_pinfo = NULL;
+    lua_tree = NULL;
+    lua_tvb = NULL;
+
+    return result;
+}
+
 static void iter_table_and_call(lua_State* LS, const gchar* table_name, lua_CFunction error_handler) {
     lua_settop(LS,0);
 
@@ -454,6 +546,8 @@ int wslua_init(register_cb cb, gpointer client_data) {
     /* the dissectors table goes in the registry (not accessible) */
     lua_newtable (L);
     lua_dissectors_table_ref = luaL_ref(L, LUA_REGISTRYINDEX);
+    lua_newtable (L);
+    lua_heur_dissectors_table_ref = luaL_ref(L, LUA_REGISTRYINDEX);
 
     /* the preferences apply_cb table (accessible by the user) */
     lua_newtable (L);
diff --git a/epan/wslua/wslua.h b/epan/wslua/wslua.h
index 9a41a714f0..840c61a28c 100644
--- a/epan/wslua/wslua.h
+++ b/epan/wslua/wslua.h
@@ -552,6 +552,7 @@ extern tvbuff_t* lua_tvb;
 extern dissector_handle_t lua_data_handle;
 extern gboolean lua_initialized;
 extern int lua_dissectors_table_ref;
+extern int lua_heur_dissectors_table_ref;
 
 WSLUA_DECLARE_CLASSES()
 WSLUA_DECLARE_FUNCTIONS()
@@ -565,8 +566,11 @@ extern const gchar* lua_shiftstring(lua_State* L,int idx);
 extern void wslua_setfuncs(lua_State *L, const luaL_Reg *l, int nup);
 extern const gchar* wslua_typeof_unknown;
 extern const gchar* wslua_typeof(lua_State *L, int idx);
+extern gboolean wslua_get_table(lua_State *L, int idx, const gchar *name);
+extern gboolean wslua_get_field(lua_State *L, int idx, const gchar *name);
 extern void wslua_assert_table_field_new(lua_State *L, int idx, const gchar *name);
 extern int dissect_lua(tvbuff_t* tvb, packet_info* pinfo, proto_tree* tree, void* data);
+extern int heur_dissect_lua(tvbuff_t* tvb, packet_info* pinfo, proto_tree* tree, void* data);
 extern void wslua_prefs_changed(void);
 extern void proto_register_lua(void);
 extern GString* lua_register_all_taps(void);
diff --git a/epan/wslua/wslua_internals.c b/epan/wslua/wslua_internals.c
index 50cd8a3d40..15d62845ff 100644
--- a/epan/wslua/wslua_internals.c
+++ b/epan/wslua/wslua_internals.c
@@ -99,13 +99,17 @@ WSLUA_API void wslua_setfuncs(lua_State *L, const luaL_Reg *l, int nup) {
   lua_pop(L, nup);  /* remove upvalues */
 }
 
-/* identical to lua_getfield but without triggering metamethods */
+/* identical to lua_getfield but without triggering metamethods
+   warning: cannot be used directly with negative index (and shouldn't be changed to)
+   decrement your negative index if you want to use this */
 static void lua_rawgetfield(lua_State *L, int idx, const char *k) {
     lua_pushstring(L, k);
     lua_rawget(L, idx);
 }
 
-/* identical to lua_setfield but without triggering metamethods */
+/* identical to lua_setfield but without triggering metamethods
+   warning: cannot be used with negative index (and shouldn't be changed to)
+   decrement your negative index if you want to use this */
 static void lua_rawsetfield (lua_State *L, int idx, const char *k) {
     lua_pushstring(L, k);
     lua_insert(L, -2);
@@ -140,6 +144,38 @@ const gchar* wslua_typeof(lua_State *L, int idx) {
     return classname;
 }
 
+/* this gets a Lua table of the given name, from the table at the given
+ * location idx. If it does not get a table, it pops whatever it got
+ * and returns false.
+ * warning: cannot be used with pseudo-indeces like LUA_REGISTRYINDEX
+ */
+gboolean wslua_get_table(lua_State *L, int idx, const gchar *name) {
+    gboolean result = TRUE;
+    if (idx < 0) idx--;
+    lua_rawgetfield(L, idx, name);
+    if (!lua_istable(L,-1)) {
+        lua_pop(L,1);
+        result = FALSE;
+    }
+    return result;
+}
+
+/* this gets a table field of the given name, from the table at the given
+ * location idx. If it does not get a field, it pops whatever it got
+ * and returns false.
+ * warning: cannot be used with pseudo-indeces like LUA_REGISTRYINDEX
+ */
+gboolean wslua_get_field(lua_State *L, int idx, const gchar *name) {
+    gboolean result = TRUE;
+    if (idx < 0) idx--;
+    lua_rawgetfield(L, idx, name);
+    if (lua_isnil(L,-1)) {
+        lua_pop(L,1);
+        result = FALSE;
+    }
+    return result;
+}
+
 /* This verifies/asserts that field 'name' doesn't already exist in table at location idx.
  * If it does, this EXITS wireshark, because this is a fundamental programming error.
  * As such, this function is only useful for special circumstances, notably
diff --git a/epan/wslua/wslua_pinfo.c b/epan/wslua/wslua_pinfo.c
index 13a2a33a7a..c5cbaa3802 100644
--- a/epan/wslua/wslua_pinfo.c
+++ b/epan/wslua/wslua_pinfo.c
@@ -36,6 +36,7 @@
 #include "wslua.h"
 
 #include <epan/addr_resolv.h>
+#include <epan/conversation.h>
 #include <string.h>
 
 
@@ -1128,6 +1129,23 @@ static int Pinfo_get_lo(lua_State *L) {
     return 1;
 }
 
+/* WSLUA_ATTRIBUTE Pinfo_conversation WO sets the packet conversation to the given Proto object */
+static int Pinfo_set_conversation(lua_State *L) {
+    Pinfo pinfo = checkPinfo(L,1);
+    Proto proto = checkProto(L,2);
+    conversation_t  *conversation;
+
+    if (!proto->handle) {
+        luaL_error(L,"Proto %s has no registered dissector", proto->name? proto->name:"<UKNOWN>");
+        return 0;
+    }
+
+    conversation = find_or_create_conversation(pinfo->ws_pinfo);
+    conversation_set_dissector(conversation,proto->handle);
+
+    return 0;
+}
+
 /* Gets registered as metamethod automatically by WSLUA_REGISTER_CLASS/META */
 static int Pinfo__gc(lua_State* L) {
     Pinfo pinfo = toPinfo(L,1);
@@ -1183,6 +1201,7 @@ WSLUA_ATTRIBUTES Pinfo_attributes[] = {
     WSLUA_ATTRIBUTE_ROREG(Pinfo,fragmented),
     WSLUA_ATTRIBUTE_ROREG(Pinfo,match_uint),
     WSLUA_ATTRIBUTE_ROREG(Pinfo,match_string),
+    WSLUA_ATTRIBUTE_WOREG(Pinfo,conversation),
     { NULL, NULL, NULL }
 };
 
diff --git a/epan/wslua/wslua_proto.c b/epan/wslua/wslua_proto.c
index 1f1e320cc9..172e049c40 100644
--- a/epan/wslua/wslua_proto.c
+++ b/epan/wslua/wslua_proto.c
@@ -1458,6 +1458,76 @@ WSLUA_FUNCTION wslua_register_postdissector(lua_State* L) {
     return 0;
 }
 
+WSLUA_METHOD Proto_register_heuristic(lua_State* L) {
+    /* Registers a heristic dissector function for this protocol, for the given heristic list name.
+       When later called, the passed-in function will be given (1) a Tvb object, (2) a Pinfo object,
+       and (3) a TreeItem object. The function must return true if the payload is for it, else false.
+       The function should perform as much verification as possible to ensure the payload is for it,
+       and dissect the packet (including setting TreeItem info and such) only if the payload is for it,
+       before returning true or false.*/
+#define WSLUA_ARG_Proto_register_heuristic_LISTNAME 2 /* the heristic list name this function is a heuristic for (e.g., "udp" or "infiniband.payload") */
+#define WSLUA_ARG_Proto_register_heuristic_FUNC 3 /* a Lua function that will be invoked for heuristic dissection */
+    Proto proto = checkProto(L,1);
+    const gchar *listname = luaL_checkstring(L, WSLUA_ARG_Proto_register_heuristic_LISTNAME);
+    const gchar *proto_name = proto->name;
+    const int top = lua_gettop(L);
+
+    if (!proto_name || proto->hfid == -1) {
+        /* this shouldn't happen - internal bug if it does */
+        luaL_error(L,"Proto_register_heuristic: got NULL proto name or invalid hfid");
+        return 0;
+    }
+
+    /* verify listname has a heuristic list */
+    if (!has_heur_dissector_list(listname)) {
+        luaL_error(L, "there is no heuristic list for '%s'", listname);
+        return 0;
+    }
+
+    /* heuristic functions are stored in a table in the registry; the registry has a
+     * table at reference lua_heur_dissectors_table_ref, and that table has keys for
+     * the heuristic listname (e.g., "udp", "tcp", etc.), and that key's value is a
+     * table of keys of the Proto->name, and their value is the function.
+     * So it's like registry[table_ref][heur_list_name][proto_name] = func
+     */
+    if (lua_isfunction(L,WSLUA_ARG_Proto_register_heuristic_FUNC)) {
+        /* insert the heur dissector into the heur dissectors table */
+        lua_rawgeti(L, LUA_REGISTRYINDEX, lua_heur_dissectors_table_ref);
+        /* the heuristic lists table is now at -1 */
+        if (!lua_istable(L,-1)) {
+            /* this shouldn't be possible */
+            luaL_error(L,"Proto_register_heuristic: could not get lua_heur_dissectors table from registry");
+            return 0;
+        }
+
+        if (!wslua_get_table(L,-1,listname)) {
+            /* no one's registered a lua heuristic for this list, so make a new list table */
+            lua_newtable(L);
+            lua_pushvalue(L,-1); /* duplicate the table so we can set it as a field */
+            lua_setfield(L,-3,listname); /* sets this new list table into the lists table */
+        }
+        else if (wslua_get_field(L,-1,proto_name)) {
+            luaL_error(L,"A heuristic dissector for Proto '%s' is already registered for the '%s' list", proto_name, listname);
+            return 0;
+        }
+
+        /* copy the func, set it as the value for key proto_name in listname's table */
+        lua_pushvalue(L,WSLUA_ARG_Proto_register_heuristic_FUNC);
+        lua_setfield(L,-2,proto_name);
+
+        /* ok, we're done with lua stuff, pop what we added to the stack */
+        lua_pop(L,2); /* pop the lists table and the listname table */
+        g_assert(top == lua_gettop(L));
+
+        /* now register the single/common heur_dissect_lua function */
+        heur_dissector_add(listname, heur_dissect_lua, proto->hfid);
+
+    } else {
+        luaL_argerror(L,3,"The heuristic dissector must be a function");
+    }
+    return 0;
+}
+
 /* WSLUA_ATTRIBUTE Proto_dissector RW The protocol's dissector, a function you define.
    The called dissector function will be given three arguments of (1) a Tvb object, (2) a Pinfo object, and (3) a TreeItem object. */
 static int Proto_get_dissector(lua_State* L) {
@@ -1602,6 +1672,7 @@ WSLUA_ATTRIBUTES Proto_attributes[] = {
 
 WSLUA_METHODS Proto_methods[] = {
     WSLUA_CLASS_FNREG(Proto,new),
+    WSLUA_CLASS_FNREG(Proto,register_heuristic),
     { NULL, NULL }
 };
 
@@ -1780,6 +1851,14 @@ WSLUA_METHOD Dissector_call(lua_State* L) {
     return 0;
 }
 
+WSLUA_METAMETHOD Dissector__call(lua_State* L) {
+    /* Calls a dissector against a given packet (or part of it) */
+#define WSLUA_ARG_Dissector__call_TVB 2 /* The buffer to dissect */
+#define WSLUA_ARG_Dissector__call_PINFO 3 /* The packet info */
+#define WSLUA_ARG_Dissector__call_TREE 4 /* The tree on which to add the protocol items */
+    return Dissector_call(L);
+}
+
 WSLUA_METAMETHOD Dissector__tostring(lua_State* L) {
     /* Gets the Dissector's protocol short name */
     Dissector d = checkDissector(L,1);
@@ -1803,6 +1882,7 @@ WSLUA_METHODS Dissector_methods[] = {
 
 WSLUA_META Dissector_meta[] = {
     WSLUA_CLASS_MTREG(Dissector,tostring),
+    WSLUA_CLASS_MTREG(Dissector,call),
     { NULL, NULL }
 };
 
@@ -1856,6 +1936,58 @@ WSLUA_CONSTRUCTOR DissectorTable_new (lua_State *L) {
     return 0;
 }
 
+/* this struct is used for passing ourselves user_data through dissector_all_tables_foreach_table() */
+typedef struct dissector_tables_foreach_table_info {
+    int num;
+    lua_State *L;
+} dissector_tables_foreach_table_info_t;
+
+/* this is the DATFunc_table function used for dissector_all_tables_foreach_table()
+   so we can get all dissector_table names. This pushes the name into a table at stack index 1 */
+static void
+dissector_tables_list_func(const gchar *table_name, const gchar *ui_name _U_, gpointer user_data) {
+    dissector_tables_foreach_table_info_t *data = (dissector_tables_foreach_table_info_t*) user_data;
+    lua_pushstring(data->L, table_name);
+    lua_rawseti(data->L, 1, data->num);
+    data->num = data->num + 1;
+}
+
+WSLUA_CONSTRUCTOR DissectorTable_list (lua_State *L) {
+    /* Gets a Lua array table of all DissectorTable names - i.e., the string names you can
+       use for the first argument to DissectorTable.get().
+       NOTE: this is an expensive operation, and should only be used for troubleshooting. */
+    dissector_tables_foreach_table_info_t data = { 1, L };
+
+    lua_newtable(L);
+
+    dissector_all_tables_foreach_table(dissector_tables_list_func, (gpointer)&data, (GCompareFunc)compare_dissector_key_name);
+
+    WSLUA_RETURN(1); /* The array table of registered DissectorTable names */
+}
+
+/* this is the DATFunc_heur_table function used for dissector_all_heur_tables_foreach_table()
+   so we can get all heuristic dissector list names. This pushes the name into a table at stack index 1 */
+static void
+heur_dissector_tables_list_func(const gchar *table_name, gpointer table _U_, gpointer user_data) {
+    dissector_tables_foreach_table_info_t *data = (dissector_tables_foreach_table_info_t*) user_data;
+    lua_pushstring(data->L, table_name);
+    lua_rawseti(data->L, 1, data->num);
+    data->num = data->num + 1;
+}
+
+WSLUA_CONSTRUCTOR DissectorTable_heuristic_list (lua_State *L) {
+    /* Gets a Lua array table of all heuristic list names - i.e., the string names you can
+       use for the first argument in Proto:register_heuristic().
+       NOTE: this is an expensive operation, and should only be used for troubleshooting. */
+    dissector_tables_foreach_table_info_t data = { 1, L };
+
+    lua_newtable(L);
+
+    dissector_all_heur_tables_foreach_table(heur_dissector_tables_list_func, (gpointer)&data);
+
+    WSLUA_RETURN(1); /* The array table of registered heuristic list names */
+}
+
 WSLUA_CONSTRUCTOR DissectorTable_get (lua_State *L) {
     /*
      Obtain a reference to an existing dissector table.
@@ -2218,6 +2350,8 @@ static int DissectorTable__gc(lua_State* L _U_) {
 WSLUA_METHODS DissectorTable_methods[] = {
     WSLUA_CLASS_FNREG(DissectorTable,new),
     WSLUA_CLASS_FNREG(DissectorTable,get),
+    WSLUA_CLASS_FNREG(DissectorTable,list),
+    WSLUA_CLASS_FNREG(DissectorTable,heuristic_list),
     WSLUA_CLASS_FNREG(DissectorTable,add),
     WSLUA_CLASS_FNREG(DissectorTable,set),
     WSLUA_CLASS_FNREG(DissectorTable,remove),
diff --git a/epan/wslua/wslua_tvb.c b/epan/wslua/wslua_tvb.c
index c4746bc97f..1dae9c3461 100644
--- a/epan/wslua/wslua_tvb.c
+++ b/epan/wslua/wslua_tvb.c
@@ -495,7 +495,7 @@ static int Tvb__gc(lua_State* L) {
 }
 
 WSLUA_METHOD Tvb_reported_len(lua_State* L) {
-	/* Obtain the reported length of a TVB */
+	/* Obtain the reported (not captured) length of a TVB */
     Tvb tvb = checkTvb(L,1);
 
     lua_pushnumber(L,tvb_reported_length(tvb->ws_tvb));
@@ -503,7 +503,7 @@ WSLUA_METHOD Tvb_reported_len(lua_State* L) {
 }
 
 WSLUA_METHOD Tvb_len(lua_State* L) {
-	/* Obtain the length of a TVB */
+	/* Obtain the actual (captured) length of a TVB */
     Tvb tvb = checkTvb(L,1);
 
     lua_pushnumber(L,tvb_length(tvb->ws_tvb));
@@ -511,7 +511,7 @@ WSLUA_METHOD Tvb_len(lua_State* L) {
 }
 
 WSLUA_METHOD Tvb_reported_length_remaining(lua_State* L) {
-	/* Obtain the reported length of packet data to end of a TVB or -1 if the offset is beyond the end of the TVB */
+	/* Obtain the reported (not captured) length of packet data to end of a TVB or -1 if the offset is beyond the end of the TVB */
 #define Tvb_reported_length_remaining_OFFSET 2 /* offset */
     Tvb tvb = checkTvb(L,1);
     int offset = luaL_optint(L, Tvb_reported_length_remaining_OFFSET, 0);