Lua: allow a Dissector object to be passed in for register_heuristic

Bug: 10695
Change-Id: I81181b2d00fcb5f0c25ab89bbe4968897f47a3a6
Reviewed-on: https://code.wireshark.org/review/6131
Petri-Dish: Hadriel Kaplan <hadrielk@yahoo.com>
Tested-by: Petri Dish Buildbot <buildbot-no-reply@wireshark.org>
Reviewed-by: Michael Mann <mmann78@netscape.net>

2 files changed, 32 insertions, 0 deletions

diff --git a/epan/wslua/init_wslua.c b/epan/wslua/init_wslua.c
index 40f03a276a..bee9fb6d3d 100644
--- a/epan/wslua/init_wslua.c
+++ b/epan/wslua/init_wslua.c
@@ -283,6 +283,8 @@ gboolean heur_dissect_lua(tvbuff_t* tvb, packet_info* pinfo, proto_tree* tree, v
     } else {
         if (lua_isboolean(L, -1) || lua_isnil(L, -1)) {
             result = lua_toboolean(L, -1);
+        } else if (lua_type(L, -1) == LUA_TNUMBER) {
+            result = lua_tointeger(L,-1) != 0 ? TRUE : FALSE;
         } else {
             proto_tree_add_expert_format(tree, pinfo, &ei_lua_error, tvb, 0, 0,
                     "Lua Error: invalid return value from Lua %s heuristic dissector", pinfo->current_proto);
diff --git a/epan/wslua/wslua_proto.c b/epan/wslua/wslua_proto.c
index 069ab373ff..afa4526bbe 100644
--- a/epan/wslua/wslua_proto.c
+++ b/epan/wslua/wslua_proto.c
@@ -1635,6 +1635,13 @@ WSLUA_METHOD Proto_register_heuristic(lua_State* L) {
        and dissect the packet (including setting TreeItem info and such) only if the payload is for it,
        before returning true or false.
 
+       Since version 1.99.1, this function also accepts a Dissector object as the second argument,
+       to allow re-using the same Lua code as the `function proto.dissector(...)`. In this case,
+       the Dissector must return a Lua number of the number of bytes consumed/parsed: if 0 is returned,
+       it will be treated the same as a `false` return for the heuristic; if a positive or negative
+       number is returned, then the it will be treated the same as a `true` return for the heuristic,
+       meaning the packet is for this protocol and no other heuristic will be tried.
+
        @since 1.11.3
      */
 #define WSLUA_ARG_Proto_register_heuristic_LISTNAME 2 /* The heuristic list name this function
@@ -1659,6 +1666,29 @@ WSLUA_METHOD Proto_register_heuristic(lua_State* L) {
         return 0;
     }
 
+    /* we'll check if the second form of this function was called: when the second arg is
+       a Dissector obejct. The truth is we don't need the Dissector object to do this
+       form of registration, but someday we might... so we're using it as a boolean arg
+       right now and in the future might use it for other things in this registration.
+     */
+    if (isDissector(L, WSLUA_ARG_Proto_register_heuristic_FUNC)) {
+        /* retrieve the Dissector's Lua function... first get the table of all dissector funcs */
+        lua_rawgeti(L, LUA_REGISTRYINDEX, lua_dissectors_table_ref);
+        /* then get the one for this Proto */
+        lua_getfield(L, -1, proto_name);
+
+        if (!lua_isfunction(L,-1)) {
+            /* this shouldn't be possible */
+            luaL_error(L,"Proto_register_heuristic: could not get lua function from lua_dissectors_table");
+            return 0;
+        }
+        /* replace the Dissector with the function */
+        lua_replace(L, WSLUA_ARG_Proto_register_heuristic_FUNC);
+        /* pop the lua_dissectors_table */
+        lua_pop(L, 1);
+        g_assert(top == lua_gettop(L));
+    }
+
     /* 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