Skip to content

Commit b196bab

Browse files
trinktrink
committed
Issue 223 - Add support to add_geoip for the full Heka field schema
1 parent 347c6a4 commit b196bab

File tree

3 files changed

+130
-13
lines changed

3 files changed

+130
-13
lines changed

‎geoip/CMakeLists.txt.geoip‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
44

55
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
6-
project(geoip VERSION 0.2.1 LANGUAGES C)
6+
project(geoip VERSION 0.2.2 LANGUAGES C)
77

88
set(CPACK_PACKAGE_NAME luasandbox-${PROJECT_NAME})
99
set(CPACK_PACKAGE_DESCRIPTION_SUMMARY "Lua geoip Module")

‎geoip/io_modules/geoip/heka.lua‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ local geoip = require "geoip.city"
4949
local ostime = require "os".time
5050
local smatch = string.match
5151
local assert = assert
52+
local ipairs = ipairs
5253
local pairs = pairs
5354
local type = type
5455

@@ -91,10 +92,49 @@ else
9192
refresh_db()
9293
end
9394

95+
96+
local function add_array(msg, field_name, values)
97+
refresh_db()
98+
local found = false
99+
for k,v in pairs(cfg.lookup) do
100+
local fname = field_name .. v
101+
local items = {}
102+
for i,j in ipairs(values) do
103+
local geo
104+
if validate_ip(j) then
105+
geo = geodb:query_by_addr(j, k)
106+
if geo then found = true end
107+
end
108+
items[i] = geo or ""
109+
end
110+
if found then
111+
msg.Fields[fname] = items
112+
else
113+
break
114+
end
115+
end
116+
117+
if found and cfg.remove_original_field then
118+
msg.Fields[field_name] = nil
119+
end
120+
end
121+
122+
94123
function add_geoip(msg, field_name)
95124
if not msg.Fields then return end
96125
local value = msg.Fields[field_name]
97-
if type(value) ~= "string" or not validate_ip(value) then return end
126+
local vt = type(value)
127+
if vt == "table" then
128+
if value.value then -- Heka schema
129+
value = value.value
130+
end
131+
vt = type(value)
132+
if vt == "table" then -- array of values
133+
add_array(msg, field_name, value)
134+
return
135+
end
136+
end
137+
if vt ~= "string" or not validate_ip(value) then return end
98138

99139
refresh_db()
100140

‎geoip/tests/test.lua‎

Lines changed: 88 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -22,31 +22,108 @@ local tests = {
2222
{{Fields = {remote_addr = "foobar"}}, "remote_addr", {remote_addr = "foobar"}},
2323
{{Fields = {remote_addr = "192.168.1.1"}}, "remote_addr", {remote_addr = "192.168.1.1"}},
2424
{{Fields = {remote_addr = "192.168.1.2"}}, "other_addr", {remote_addr = "192.168.1.2"}},
25-
{{Fields = {remote_addr = "192.168.1.2"}}, "remote_addr", {remote_addr = "192.168.1.2", remote_addr_city = "Mountain View", remote_addr_country = "US"}},
25+
{{Fields = {remote_addr = "192.168.1.2"}}, "remote_addr", {
26+
remote_addr = "192.168.1.2",
27+
remote_addr_city = "Mountain View",
28+
remote_addr_country = "US"}
29+
},
2630
{{Fields = {remote_addr = 7}}, "remote_addr", {remote_addr = 7}},
31+
{{Fields = {remote_addr = {value = "192.168.1.2", representation = "ipv4"}}},
32+
"remote_addr",
33+
{
34+
remote_addr = {value = "192.168.1.2", representation = "ipv4"},
35+
remote_addr_city = "Mountain View",
36+
remote_addr_country = "US"
37+
}
38+
},
39+
{{Fields = {remote_addrs = {"192.168.1.2", "192.168.1.1"}}},
40+
"remote_addrs",
41+
{
42+
remote_addrs = {"192.168.1.2", "192.168.1.1"},
43+
remote_addrs_city = {"Mountain View", ""},
44+
remote_addrs_country = {"US", ""}
45+
}
46+
},
47+
{{Fields = {remote_addrs = {value = {"192.168.1.2", "192.168.1.1"}, representation = "ipv4"}}},
48+
"remote_addrs",
49+
{
50+
remote_addrs = {value = {"192.168.1.2", "192.168.1.1"}, representation = "ipv4"},
51+
remote_addrs_city = {"Mountain View", ""},
52+
remote_addrs_country = {"US", ""}
53+
}
54+
},
55+
{{Fields = {remote_addrs = {"192.168.1.3", "192.168.1.4"}}},
56+
"remote_addrs",
57+
{
58+
remote_addrs = {"192.168.1.3", "192.168.1.4"}
59+
}
60+
},
2761
}
2862

29-
local function verify_table(idx, t, expected)
30-
if t == nil and expected ~= nil then
63+
local function verify_field(idx, k, expected, received)
64+
local ev = expected
65+
local rv = received
66+
if ev.value then
67+
if ev.representation ~= rv.representation then
68+
error(string.format("test: %d field: %s representation expected: %s received: %s",
69+
idx, k, tostring(ev.representation), tostring(rv.represntation)))
70+
end
71+
ev = ev.value
72+
rv = rv.value
73+
end
74+
75+
if type(ev) == "table" then
76+
for i,v in ipairs(ev) do
77+
if v ~= rv[i] then
78+
error(string.format("test: %d field: %s idx: %d expected: %s received: %s",
79+
idx, k, i, tostring(v), tostring(rv[i])))
80+
end
81+
end
82+
elseif ev ~= rv then
83+
error(string.format("test: %d field: %s expected: %s received: %s",
84+
idx, k, tostring(ev), tostring(rv)))
85+
end
86+
end
87+
88+
89+
local function verify_msg_fields(idx, expected, received)
90+
if expected == nil and received ~= nil then
3191
error(string.format("test: %d failed Fields not nil", idx))
3292
end
33-
if not t then return end
93+
if not expected then return end
3494

35-
for k,r in pairs(t) do
36-
local e = expected[k]
37-
if e ~= r then
38-
error(string.format("test: %d field: %s expected: %s received: %s", idx, k, tostring(e), tostring(r)))
95+
for k,e in pairs(expected) do
96+
local r = received[k]
97+
local etyp = type(e)
98+
local rtyp = type(r)
99+
if etyp ~= rtyp then
100+
error(string.format("test: %d field: %s type expected: %s received: %s",
101+
idx, k, etyp, rtyp))
102+
end
103+
if etyp == "table" then
104+
verify_field(idx, k, e, r)
105+
elseif e ~= r then
106+
error(string.format("test: %d field: %s expected: %s received: %s",
107+
idx, k, tostring(e), tostring(r)))
39108
end
40109
end
110+
if idx > 0 then
111+
verify_msg_fields(-idx, received, expected)
112+
end
41113
end
42114

43115
for i,v in ipairs(tests) do
44116
gh.add_geoip(v[1], v[2])
45-
verify_table(i, v[1].Fields, v[3])
117+
verify_msg_fields(i, v[3], v[1].Fields)
46118
end
47119

120+
cfg.geoip_heka.lookup = {city = ""} -- replace the field
121+
local msg = {Fields = {remote_addr = "192.168.1.2"}}
122+
gh.add_geoip(msg, "remote_addr")
123+
verify_msg_fields(98, {remote_addr = "Mountain View"}, msg.Fields)
124+
125+
cfg.geoip_heka.lookup = {city = "_city"}
48126
cfg.geoip_heka.remove_original_field = true
49127
local msg = {Fields = {remote_addr = "192.168.1.2"}}
50128
gh.add_geoip(msg, "remote_addr")
51-
verify_table(99, msg.Fields, tests[6][3])
52-
assert(not msg.Fields.remote_addr)
129+
verify_msg_fields(99, {remote_addr_city = "Mountain View"}, msg.Fields)

0 commit comments

Comments
 (0)