haproxy_geoip2/
lib.rs

1use std::net::IpAddr;
2
3use haproxy_api::Core;
4use mlua::prelude::{IntoLua, Lua, LuaFunction, LuaResult, LuaTable, LuaValue};
5use mlua::Variadic;
6
7#[derive(Debug, Clone)]
8pub(crate) enum GeoValue<'a> {
9    Str(&'a str),
10    Float(f64),
11    UInt(u32),
12    Bool(bool),
13}
14
15impl IntoLua for GeoValue<'_> {
16    fn into_lua(self, lua: &Lua) -> LuaResult<LuaValue> {
17        match self {
18            GeoValue::Str(s) => s.into_lua(lua),
19            GeoValue::Float(f) => f.into_lua(lua),
20            GeoValue::UInt(u) => u.into_lua(lua),
21            GeoValue::Bool(b) => b.into_lua(lua),
22        }
23    }
24}
25
26/// Register GeoIP2 lookups in the haproxy
27///
28/// This function registers the GeoIP2 lookups converters in the haproxy Lua environment,
29/// including task to reload the databases at a given interval.
30///
31/// The following Lua options are supported:
32///
33/// - `reload_interval`: Interval in seconds to reload the databases. Default is 0.
34/// - `db.city`: Path to the MaxMind GeoIP2 City database.
35/// - `db.asn`: Path to the MaxMind GeoIP2 ASN database.
36///
37/// # Example
38/// ```lua
39/// geoip2.register({
40///    reload_interval = 86400, -- 1 day
41///    db = {
42///        city = "/path/to/GeoLite2-City.mmdb",
43///         asn = "/path/to/GeoLite2-ASN.mmdb",
44///     },
45/// })
46/// ```
47pub fn register(lua: &Lua, options: LuaTable) -> LuaResult<()> {
48    let core = Core::new(lua)?;
49
50    // Parse options
51    let reload_interval: u64 = options.get("reload_interval").unwrap_or(0);
52
53    // Register databases
54    if let Ok(db) = options.get::<LuaTable>("db") {
55        if let Ok(path) = db.get::<String>("city") {
56            city::DB.configure(path.into(), reload_interval);
57            register_converter(&core, &city::DB, "city", city::lookup)?;
58        }
59
60        if let Ok(asn_path) = db.get::<String>("asn") {
61            asn::DB.configure(asn_path.into(), reload_interval);
62            register_converter(&core, &asn::DB, "asn", asn::lookup)?;
63        }
64    }
65
66    Ok(())
67}
68
69// `F` is zero-sized, so it's safe to send it across threads
70fn register_converter<F>(
71    core: &Core,
72    db: &'static db::Database,
73    prefix: &str,
74    lookup: F,
75) -> LuaResult<()>
76where
77    F: Fn(&Lua, IpAddr, &[String]) -> Option<LuaValue> + Send + Copy + 'static,
78{
79    // Trigger dummy lookup within a worker to load the database
80    core.register_task(move |lua| {
81        lookup(lua, "0.0.0.0".parse()?, &[]);
82        Ok(())
83    })?;
84
85    // Register reload task
86    let interval = db.reload_interval();
87    if interval > 0 {
88        let trigger_reload = LuaFunction::wrap(|| {
89            db.trigger_reload();
90            Ok(())
91        });
92        core.register_lua_task(mlua::chunk! {
93            if core.thread <= 1 then
94                while true do
95                    core.sleep($interval)
96                    $trigger_reload()
97                end
98            end
99        })?;
100    }
101
102    core.register_converters(
103        &format!("geoip2-lookup-{prefix}"),
104        move |lua, (ip, props): (String, Variadic<String>)| {
105            let ip = ip.parse::<IpAddr>()?;
106            lookup(lua, ip, &props).map_or_else(|| "".into_lua(lua), Ok)
107        },
108    )
109}
110
111mod asn;
112mod city;
113mod db;
haproxy_geoip2/lib.rs

haproxy_geoip2/
lib.rs
