Wireshark protocol dissector for Riak

riak.lua

-- Riak protocol decoder
-- Based on copypasta from https://wiki.wireshark.org/Protobuf.md
-- To use,
--   * Plase in ~/.local/lib/wireshark/plugins/riak.lua
--   * Check out https://github.com/protocolbuffers/protobuf somewhere
--   * Check out https://github.com/basho/riak_pb somewhere
--   * In the "ProtoBuf" protocol preferences, add protobuf/src and riak_pb/src
--     to the search paths. Make sure "load all" is checked for riak_pb, but NOT
--     for protobuf.
--   * right click a stream -> decode as -> "RIAK_PROTO"
do
    local pb_type_table = {
        [0] = "RpbErrorResp",
        [1] = "RpbPingReq",
        [2] = "RpbPingResp",
        [3] = "RpbGetClientIdReq",
        [4] = "RpbGetClientIdResp",
        [5] = "RpbSetClientIdReq",
        [6] = "RpbSetClientIdResp",
        [7] = "RpbGetServerInfoReq",
        [8] = "RpbGetServerInfoResp",
        [9] = "RpbGetReq",
        [10] = "RpbGetResp",
        [11] = "RpbPutReq",
        [12] = "RpbPutResp",
        [13] = "RpbDelReq",
        [14] = "RpbDelResp",
        [15] = "RpbListBucketsReq",
        [16] = "RpbListBucketsResp",
        [17] = "RpbListKeysReq",
        [18] = "RpbListKeysResp",
        [19] = "RpbGetBucketReq",
        [20] = "RpbGetBucketResp",
        [21] = "RpbSetBucketReq",
        [22] = "RpbSetBucketResp",
        [23] = "RpbMapRedReq",
        [24] = "RpbMapRedResp",
        [25] = "RpbIndexReq",
        [26] = "RpbIndexResp",
        [27] = "RpbSearchQueryReq",
        [28] = "RpbSearchQueryResp",
        [29] = "RpbResetBucketReq",
        [30] = "RpbResetBucketResp",
        [31] = "RpbGetBucketTypeReq",
        [32] = "RpbSetBucketTypeReq",
        [33] = "RpbGetBucketKeyPreflistReq",
        [34] = "RpbGetBucketKeyPreflistResp",
        [40] = "RpbCSBucketReq",
        [41] = "RpbCSBucketResp",
        [42] = "RpbIndexBodyResp",
        [50] = "RpbCounterUpdateReq",
        [51] = "RpbCounterUpdateResp",
        [52] = "RpbCounterGetReq",
        [53] = "RpbCounterGetResp",
        [54] = "RpbYokozunaIndexGetReq",
        [55] = "RpbYokozunaIndexGetResp",
        [56] = "RpbYokozunaIndexPutReq",
        [57] = "RpbYokozunaIndexDeleteReq",
        [58] = "RpbYokozunaSchemaGetReq",
        [59] = "RpbYokozunaSchemaGetResp",
        [60] = "RpbYokozunaSchemaPutReq",
        [70] = "RpbCoverageReq",
        [71] = "RpbCoverageResp",
        [80] = "DtFetchReq",
        [81] = "DtFetchResp",
        [82] = "DtUpdateReq",
        [83] = "DtUpdateResp",
        [90] = "TsQueryReq",
        [91] = "TsQueryResp",
        [92] = "TsPutReq",
        [93] = "TsPutResp",
        [94] = "TsDelReq",
        [95] = "TsDelResp",
        [96] = "TsGetReq",
        [97] = "TsGetResp",
        [98] = "TsListKeysReq",
        [99] = "TsListKeysResp",
        [100] = "TsCoverageReq",
        [101] = "TsCoverageResp",
        [102] = "TsCoverageEntry",
        [103] = "TsRange",
        [104] = "TsTtbMsg",
        [128] = "RpbReplGetReq",
        [129] = "RpbReplGetClusterIdReq",
        [130] = "RpbReplGetClusterIdResp",
        [200] = "RpbRTEReq",
        [201] = "RpbRTEResp",
        [202] = "RpbFetchReq",
        [203] = "RpbFetchResp",
        [204] = "RpbPushReq",
        [205] = "RpbPushResp",
        [206] = "RpbMembershipReq",
        [207] = "RpbMembershipResp",
        [210] = "RpbAaeFoldMergeRootNValReq",
        [211] = "RpbAaeFoldMergeBranchNValReq",
        [212] = "RpbAaeFoldFetchClocksNValReq",
        [213] = "RpbAaeFoldMergeTreesRangeReq",
        [214] = "RpbAaeFoldFetchClocksRangeReq",
        [215] = "RpbAaeFoldFindKeysReq",
        [216] = "RpbAaeFoldObjectStatsReq",
        [217] = "RpbAaeFoldReplKeysReq",
        [218] = "RpbAaeFoldFindTombsReq",
        [219] = "RpbAaeFoldRepairKeysReq",
        [220] = "RpbAaeFoldTreeResp",
        [221] = "RpbAaeFoldKeyValueResp",
        [222] = "RpbAaeFoldKeyCountResp",
        [223] = "RpbAaeFoldListBucketsReq",
        [224] = "RpbAaeFoldListBucketsResp",
        [230] = "RpbAaeFoldReapTombsReq",
        [231] = "RpbAaeFoldEraseKeysReq",
        [253] = "RpbAuthReq",
        [254] = "RpbAuthResp",
        [255] = "RpbStartTls",
    }

    local pb_dissector = Dissector.get('protobuf')
    local proto = Proto("riak_proto", "Riak protobuf")
    local f_length = ProtoField.uint32("riak_proto.length", "Length", base.DEC)
    local f_code = ProtoField.uint8("riak_proto.code", "Code", base.DEC, pb_type_table)
    proto.fields = { f_length, f_code }
    proto.dissector = function(tvb, pinfo, tree)
        local subtree = tree:add(proto, tvb())
        local offset = 0
        local remaining_len = tvb:len()
        while remaining_len > 0 do
            if remaining_len < 5 then -- head not enough
                pinfo.desegment_offset = offset
                pinfo.desegment_len = DESEGMENT_ONE_MORE_SEGMENT
                return -1
            end
            local data_len = tvb(offset, 4):uint()
            if remaining_len - 4 < data_len then -- data not enough
                pinfo.desegment_offset = offset
                pinfo.desegment_len = data_len - (remaining_len - 4)
                return -1
            end
            local code = tvb(offset + 4, 1):uint()
            subtree:add(f_length, tvb(offset, 4))
            subtree:add(f_code, tvb(offset + 4, 1))

            local pb_type = pb_type_table[code]
            if pb_type ~= nil then
                pinfo.private["pb_msg_type"] = "message," .. pb_type
            end

            pcall(Dissector.call, pb_dissector, tvb(offset + 5, data_len - 1):tvb(), pinfo, subtree)
            offset = offset + 4 + data_len
            remaining_len = remaining_len - 4 - data_len
        end
        pinfo.columns.protocol:set("riak_proto")
    end
    DissectorTable.get("tcp.port"):add(0, proto)
    return proto
end