Allow chained calls with Lua TreeItem functions, and fix a couple of minor errors.

A common Lua idiom is to use chained calls, i.e. tree:foo():bar():choo(). This actually
works for tree:add() because it returns the new child tree item which is then the one
being applied to the next chained call. But it doesn't work beyond that for things like
set_generated() and so on.  So this commit fixes that.

This also fixes the Lua tree:add() function for the FT_BOOL type to let it be a Lua boolean value.
And it reverts a previous change to Struct.tohex() to allow coercion of the argument.

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

5 files changed, 121 insertions, 23 deletions

diff --git a/epan/wslua/wslua.h b/epan/wslua/wslua.h
index 1ec3087233..5909027273 100644
--- a/epan/wslua/wslua.h
+++ b/epan/wslua/wslua.h
@@ -661,7 +661,10 @@ WSLUA_DECLARE_FUNCTIONS()
 extern lua_State* wslua_state(void);
 
 extern int wslua__concat(lua_State* L);
+extern gboolean wslua_toboolean(lua_State* L, int n);
+extern gboolean wslua_checkboolean(lua_State* L, int n);
 extern gboolean wslua_optbool(lua_State* L, int n, gboolean def);
+extern lua_Integer wslua_tointeger(lua_State* L, int n);
 extern int wslua_optboolint(lua_State* L, int n, int def);
 extern const char* wslua_checklstring_only(lua_State* L, int n, size_t *l);
 extern const char* wslua_checkstring_only(lua_State* L, int n);
diff --git a/epan/wslua/wslua_internals.c b/epan/wslua/wslua_internals.c
index c18bdc6441..0d0e61388c 100644
--- a/epan/wslua/wslua_internals.c
+++ b/epan/wslua/wslua_internals.c
@@ -43,6 +43,36 @@ WSLUA_API int wslua__concat(lua_State* L) {
     return 1;
 }
 
+/* like lua_toboolean, except only coerces int, nil, and bool, and errors on other types.
+   note that normal lua_toboolean returns 1 for any Lua value different from false and
+   nil; otherwise it returns 0. So a string would give a 0, as would a number of 1.
+   This function errors if the arg is a string, and sets the boolean to 1 for any
+   number other than 0. Like toboolean, this returns FALSE if the arg was missing. */
+WSLUA_API gboolean wslua_toboolean(lua_State* L, int n) {
+    gboolean val = FALSE;
+
+    if ( lua_isboolean(L,n) ||  lua_isnil(L,n)  || lua_gettop(L) < n ) {
+        val = lua_toboolean(L,n);
+    } else if ( lua_type(L,n) == LUA_TNUMBER ) {
+        int num = luaL_checkint(L,n);
+        val = num != 0 ? TRUE : FALSE;
+    } else {
+        luaL_argerror(L,n,"must be a boolean or number");
+    }
+
+    return val;
+}
+
+/* like luaL_checkint, except for booleans - this does not coerce other types */
+WSLUA_API gboolean wslua_checkboolean(lua_State* L, int n) {
+
+    if (!lua_isboolean(L,n) ) {
+        luaL_argerror(L,n,"must be a boolean");
+    }
+
+    return lua_toboolean(L,n);;
+}
+
 WSLUA_API gboolean wslua_optbool(lua_State* L, int n, gboolean def) {
     gboolean val = FALSE;
 
@@ -57,6 +87,24 @@ WSLUA_API gboolean wslua_optbool(lua_State* L, int n, gboolean def) {
     return val;
 }
 
+/* like lua_tointeger, except only coerces int, nil, and bool, and errors on other types.
+   note that normal lua_tointeger does not coerce nil or bool, but does coerce strings. */
+WSLUA_API lua_Integer wslua_tointeger(lua_State* L, int n) {
+    lua_Integer val = 0;
+
+    if ( lua_type(L,n) == LUA_TNUMBER) {
+        val = lua_tointeger(L,n);
+    } else if ( lua_isboolean(L,n) ) {
+        val = (lua_Integer) (lua_toboolean(L,n));
+    } else if ( lua_isnil(L,n) ) {
+        val = 0;
+    } else {
+        luaL_argerror(L,n,"must be a integer, boolean or nil");
+    }
+
+    return val;
+}
+
 /* like luaL_optint, except converts/handles Lua booleans as well */
 WSLUA_API int wslua_optboolint(lua_State* L, int n, int def) {
     int val = 0;
diff --git a/epan/wslua/wslua_proto.c b/epan/wslua/wslua_proto.c
index 5ef77cddf2..0cb67e01e0 100644
--- a/epan/wslua/wslua_proto.c
+++ b/epan/wslua/wslua_proto.c
@@ -126,7 +126,7 @@ static int new_pref(lua_State* L, pref_type_t type) {
 
     switch(type) {
         case PREF_BOOL: {
-            gboolean def = lua_toboolean(L,2);
+            gboolean def = wslua_toboolean(L,2);
             pref->value.b = def;
             break;
         }
@@ -143,7 +143,7 @@ static int new_pref(lua_State* L, pref_type_t type) {
         case PREF_ENUM: {
             guint32 def = wslua_optgint32(L,2,0);
             enum_val_t *enum_val = get_enum(L,4);
-            gboolean radio = lua_toboolean(L,5);
+            gboolean radio = wslua_toboolean(L,5);
             pref->value.e = def;
             pref->info.enum_info.enumvals = enum_val;
             pref->info.enum_info.radio_buttons = radio;
diff --git a/epan/wslua/wslua_struct.c b/epan/wslua/wslua_struct.c
index d99b7ef059..06e5c203c7 100644
--- a/epan/wslua/wslua_struct.c
+++ b/epan/wslua/wslua_struct.c
@@ -618,8 +618,9 @@ WSLUA_CONSTRUCTOR Struct_tohex (lua_State *L) {
   gboolean lowercase = FALSE;
   const gchar* sep = NULL;
 
-  /* luaL_checklstring coerces the argument to a string, and we don't want to do that */
-  s = wslua_checklstring_only(L, WSLUA_ARG_Struct_tohex_BYTESTRING, &len);
+  /* luaL_checklstring coerces the argument to a string, and that's ok for tohex,
+     just not fromhex. In fact, we should accept/coerce a Int64/UInt64 here too someday. */
+  s = luaL_checklstring(L, WSLUA_ARG_Struct_tohex_BYTESTRING, &len);
 
   if (!s) {
     WSLUA_ARG_ERROR(Struct_tohex,BYTESTRING,"must be a Lua string");
diff --git a/epan/wslua/wslua_tree.c b/epan/wslua/wslua_tree.c
index f0c56106b0..5eee9e561e 100644
--- a/epan/wslua/wslua_tree.c
+++ b/epan/wslua/wslua_tree.c
@@ -121,7 +121,7 @@ WSLUA_METHOD TreeItem_add_packet_field(lua_State *L) {
 
     PUSH_TREEITEM(L,tree_item);
 
-    return 1;
+    WSLUA_RETURN(1); /* The new child TreeItem. */
 }
 
 static int TreeItem_add_item_any(lua_State *L, gboolean little_endian) {
@@ -173,7 +173,11 @@ static int TreeItem_add_item_any(lua_State *L, gboolean little_endian) {
                     lua_insert(L,1);
                     break;
                 case FT_BOOLEAN:
-                    item = proto_tree_add_boolean(tree_item->tree,hfid,tvbr->tvb->ws_tvb,tvbr->offset,tvbr->len,wslua_checkguint32(L,1));
+                    {
+                        /* this needs to use checkinteger so that it can accept a Lua boolean and coerce it to an int */
+                        guint32 val = (guint32) (wslua_tointeger(L,1));
+                        item = proto_tree_add_boolean(tree_item->tree,hfid,tvbr->tvb->ws_tvb,tvbr->offset,tvbr->len,val);
+                    }
                     break;
                 case FT_UINT8:
                 case FT_UINT16:
@@ -283,7 +287,7 @@ WSLUA_METHOD TreeItem_add(lua_State *L) {
 #define WSLUA_OPTARG_TreeItem_add_TVBRANGE 3 /* The TvbRange of bytes in the packet this tree item covers/represents. */
 #define WSLUA_OPTARG_TreeItem_add_VALUE 4 /* The field's value, instead of the ProtoField/Proto one. */
 #define WSLUA_OPTARG_TreeItem_add_LABEL 5 /* One or more strings to use for the tree item label, instead of the ProtoField/Proto one. */
-    WSLUA_RETURN(TreeItem_add_item_any(L,FALSE)); /* The new child TreeItem */
+    WSLUA_RETURN(TreeItem_add_item_any(L,FALSE)); /* The new child TreeItem. */
 }
 
 WSLUA_METHOD TreeItem_add_le(lua_State *L) {
@@ -300,40 +304,58 @@ WSLUA_METHOD TreeItem_add_le(lua_State *L) {
 #define WSLUA_OPTARG_TreeItem_add_le_TVBRANGE 3 /* The TvbRange of bytes in the packet this tree item covers/represents. */
 #define WSLUA_OPTARG_TreeItem_add_le_VALUE 4 /* The field's value, instead of the ProtoField/Proto one. */
 #define WSLUA_OPTARG_TreeItem_add_le_LABEL 5 /* One or more strings to use for the tree item label, instead of the ProtoField/Proto one. */
-    WSLUA_RETURN(TreeItem_add_item_any(L,TRUE)); /* The new child TreeItem */
+    WSLUA_RETURN(TreeItem_add_item_any(L,TRUE)); /* The new child TreeItem. */
 }
 
 WSLUA_METHOD TreeItem_set_text(lua_State *L) {
-    /* Sets the text of the label. */
+    /* Sets the text of the label.
+
+       This used to return nothing, but as of 1.11.3 it returns the same tree item to allow chained calls.
+    */
 #define WSLUA_ARG_TreeItem_set_text_TEXT 2 /* The text to be used. */
     TreeItem ti = checkTreeItem(L,1);
     const gchar* s = luaL_checkstring(L,WSLUA_ARG_TreeItem_set_text_TEXT);
 
     proto_item_set_text(ti->item,"%s",s);
 
-    return 0;
+    /* copy the TreeItem userdata so we give it back */
+    lua_pushvalue(L, 1);
+
+    WSLUA_RETURN(1); /* The same TreeItem. */
 }
 
 WSLUA_METHOD TreeItem_append_text(lua_State *L) {
-    /* Appends text to the label. */
+    /* Appends text to the label.
+
+       This used to return nothing, but as of 1.11.3 it returns the same tree item to allow chained calls.
+    */
 #define WSLUA_ARG_TreeItem_append_text_TEXT 2 /* The text to be appended. */
     TreeItem ti = checkTreeItem(L,1);
     const gchar* s = luaL_checkstring(L,WSLUA_ARG_TreeItem_append_text_TEXT);
 
     proto_item_append_text(ti->item,"%s",s);
 
-    return 0;
+    /* copy the TreeItem userdata so we give it back */
+    lua_pushvalue(L, 1);
+
+    WSLUA_RETURN(1); /* The same TreeItem. */
 }
 
 WSLUA_METHOD TreeItem_prepend_text(lua_State *L) {
-    /* Prepends text to the label. */
+    /* Prepends text to the label.
+
+       This used to return nothing, but as of 1.11.3 it returns the same tree item to allow chained calls.
+    */
 #define WSLUA_ARG_TreeItem_prepend_text_TEXT 2 /* The text to be prepended. */
     TreeItem ti = checkTreeItem(L,1);
     const gchar* s = luaL_checkstring(L,WSLUA_ARG_TreeItem_prepend_text_TEXT);
 
     proto_item_prepend_text(ti->item,"%s",s);
 
-    return 0;
+    /* copy the TreeItem userdata so we give it back */
+    lua_pushvalue(L, 1);
+
+    WSLUA_RETURN(1); /* The same TreeItem. */
 }
 
 WSLUA_METHOD TreeItem_add_expert_info(lua_State *L) {
@@ -366,7 +388,10 @@ WSLUA_METHOD TreeItem_add_expert_info(lua_State *L) {
         expert_add_info(lua_pinfo, ti->item, ei_info);
     }
 
-    return 0;
+    /* copy the TreeItem userdata so we give it back */
+    lua_pushvalue(L, 1);
+
+    WSLUA_RETURN(1); /* The same TreeItem. */
 }
 
 WSLUA_METHOD TreeItem_add_proto_expert_info(lua_State *L) {
@@ -394,7 +419,10 @@ WSLUA_METHOD TreeItem_add_proto_expert_info(lua_State *L) {
         expert_add_info(lua_pinfo, ti->item, &expert->ids);
     }
 
-    return 0;
+    /* copy the TreeItem userdata so we give it back */
+    lua_pushvalue(L, 1);
+
+    WSLUA_RETURN(1); /* The same TreeItem. */
 }
 
 WSLUA_METHOD TreeItem_add_tvb_expert_info(lua_State *L) {
@@ -442,37 +470,55 @@ WSLUA_METHOD TreeItem_add_tvb_expert_info(lua_State *L) {
                               tvbr->tvb->ws_tvb, tvbr->offset, tvbr->len);
     }
 
-    return 0;
+    /* copy the TreeItem userdata so we give it back */
+    lua_pushvalue(L, 1);
+
+    WSLUA_RETURN(1); /* The same TreeItem. */
 }
 
 WSLUA_METHOD TreeItem_set_generated(lua_State *L) {
-    /* Marks the `TreeItem` as a generated field (with data inferred but not contained in the packet). */
+    /* Marks the `TreeItem` as a generated field (with data inferred but not contained in the packet).
+
+       This used to return nothing, but as of 1.11.3 it returns the same tree item to allow chained calls.
+    */
     TreeItem ti = checkTreeItem(L,1);
 
     PROTO_ITEM_SET_GENERATED(ti->item);
 
-    return 0;
+    /* copy the TreeItem userdata so we give it back */
+    lua_pushvalue(L, 1);
+
+    WSLUA_RETURN(1); /* The same TreeItem. */
 }
 
 
 WSLUA_METHOD TreeItem_set_hidden(lua_State *L) {
-    /* Should not be used */
+    /* This function should not be used, and is provided for backwards-compatibility only. */
     TreeItem ti = checkTreeItem(L,1);
 
     PROTO_ITEM_SET_HIDDEN(ti->item);
 
-    return 0;
+    /* copy the TreeItem userdata so we give it back */
+    lua_pushvalue(L, 1);
+
+    WSLUA_RETURN(1); /* The same TreeItem. */
 }
 
 WSLUA_METHOD TreeItem_set_len(lua_State *L) {
-    /* Set `TreeItem`'s length inside tvb, after it has already been created. */
+    /* Set `TreeItem`'s length inside tvb, after it has already been created.
+
+       This used to return nothing, but as of 1.11.3 it returns the same tree item to allow chained calls.
+    */
 #define WSLUA_ARG_TreeItem_set_len_LEN 2 /* The length to be used. */
     TreeItem ti = checkTreeItem(L,1);
     gint len = luaL_checkint(L,WSLUA_ARG_TreeItem_set_len_LEN);
 
     proto_item_set_len(ti->item, len);
 
-    return 0;
+    /* copy the TreeItem userdata so we give it back */
+    lua_pushvalue(L, 1);
+
+    WSLUA_RETURN(1); /* The same TreeItem. */
 }
 
 /* Gets registered as metamethod automatically by WSLUA_REGISTER_CLASS/META */