wttr/src/location/GeoIp.zig

const std = @import("std");
const Coordinates = @import("../Coordinates.zig");
const Ip2location = @import("Ip2location.zig");

const c = @cImport({
    @cInclude("maxminddb.h");
});

const GeoIP = @This();
const log = std.log.scoped(.geoip);

mmdb: *c.MMDB_s,
ip2location_client: *Ip2location,
allocator: std.mem.Allocator,

pub fn init(allocator: std.mem.Allocator, db_path: []const u8, api_key: ?[]const u8, cache_path: []const u8) !GeoIP {
    const path_z = try std.heap.c_allocator.dupeZ(u8, db_path);
    defer std.heap.c_allocator.free(path_z);

    const mmdb = try allocator.create(c.MMDB_s);
    errdefer allocator.destroy(mmdb);

    const status = c.MMDB_open(path_z.ptr, c.MMDB_MODE_MMAP, mmdb);
    if (status != c.MMDB_SUCCESS)
        return error.CannotOpenDatabase;

    const client: *Ip2location = try allocator.create(Ip2location);
    errdefer allocator.destroy(client);
    client.* = try Ip2location.init(allocator, api_key, cache_path);
    errdefer {
        client.deinit();
        allocator.destroy(client);
    }

    std.log.info(
        "IP2Location fallback: {s} (cache: {s})",
        .{ if (api_key) |_| "key provided, 50k/mo limit" else "no key, 1k/day limit", cache_path },
    );

    return GeoIP{
        .mmdb = mmdb,
        .ip2location_client = client,
        .allocator = allocator,
    };
}

pub fn deinit(self: *GeoIP) void {
    c.MMDB_close(self.mmdb);
    self.allocator.destroy(self.mmdb);
    self.ip2location_client.deinit();
    log.debug("destroying client", .{});
    self.allocator.destroy(self.ip2location_client);
}

pub fn lookup(self: *GeoIP, ip: []const u8) ?Coordinates {
    // Try MaxMind first
    const result = lookupInternal(self.mmdb, ip) catch return null;

    log.debug("lookup geoip db for ip {s}. Found: {}", .{ ip, result.found_entry });
    if (result.found_entry)
        if (self.extractCoordinates(ip, result)) |coords|
            return coords;

    // Fallback to IP2Location
    return self.ip2location_client.lookup(ip);
}

fn lookupInternal(mmdb: *c.MMDB_s, ip: []const u8) !c.MMDB_lookup_result_s {
    const ip_z = try std.heap.c_allocator.dupeZ(u8, ip);
    defer std.heap.c_allocator.free(ip_z);

    var gai_error: c_int = 0;
    var mmdb_error: c_int = 0;

    const result = c.MMDB_lookup_string(mmdb, ip_z.ptr, &gai_error, &mmdb_error);
    if (mmdb_error != 0) {
        log.warn("got error on MMDB_lookup_string for ip {s}. gai = {d}, mmdb_error = {d}", .{ ip, gai_error, mmdb_error });
        return error.MMDBLookupError;
    }
    return result;
}

pub fn isUSIp(self: *GeoIP, ip: []const u8) bool {
    var result = lookupInternal(self.mmdb, ip) catch return false;
    if (!result.found_entry) return false;

    // SAFETY: The C API will initialize this on the next line
    var country_data: c.MMDB_entry_data_s = undefined;
    const status = c.MMDB_get_value(&result.entry, &country_data, "country", "iso_code", @as([*c]const u8, null));

    if (status != c.MMDB_SUCCESS or !country_data.has_data) {
        log.info("lookup found result, but no country available in data for ip {s}. MMDB_get_value returned {d}", .{ ip, status });
        return false;
    }

    const country_code = country_data.unnamed_0.utf8_string[0..country_data.data_size];
    return std.mem.eql(u8, country_code, "US");
}

fn extractCoordinates(self: *GeoIP, ip: []const u8, result: c.MMDB_lookup_result_s) ?Coordinates {
    _ = self;

    if (!result.found_entry) return null;

    var entry_copy = result.entry;

    // SAFETY: latitude_data set by MMDB_get_value
    var latitude_data: c.MMDB_entry_data_s = undefined;
    const lat_status = c.MMDB_get_value(&entry_copy, &latitude_data, "location", "latitude", @as([*c]const u8, null));

    if (lat_status != c.MMDB_SUCCESS or !latitude_data.has_data) {
        log.info("lookup found result, but no latitude available in data for ip {s}. MMDB_get_value returned {d}", .{ ip, lat_status });
        return null;
    }

    // SAFETY: longitude_data set by MMDB_get_value
    var longitude_data: c.MMDB_entry_data_s = undefined;
    const lon_status = c.MMDB_get_value(&entry_copy, &longitude_data, "location", "longitude", @as([*c]const u8, null));

    if (lon_status != c.MMDB_SUCCESS or !longitude_data.has_data) {
        log.info("lookup found result, but no longitude available in data for ip {s}. MMDB_get_value returned {d}", .{ ip, lon_status });
        return null;
    }

    var coords = [_]f64{ latitude_data.unnamed_0.double_value, longitude_data.unnamed_0.double_value };

    // Depending on how this is compiled, the byteswap may or may not be necessary
    // original c, compiled with zig, statically linked: byteSwap
    // pre=built, dynamically linked, do not byte swap
    // I'm not sure precisely what causes this
    std.mem.byteSwapAllElements(f64, &coords);

    const latitude = coords[0];
    const longitude = coords[1];

    return .{
        .latitude = latitude,
        .longitude = longitude,
    };
}

test "MMDB functions are callable" {
    const mmdb_error = c.MMDB_strerror(0);
    try std.testing.expect(mmdb_error[0] != 0);
}

test "GeoIP init with invalid path fails" {
    const result = GeoIP.init(std.testing.allocator, "/nonexistent/path.mmdb", null, "");
    try std.testing.expectError(error.CannotOpenDatabase, result);
}

test "isUSIp detects US IPs" {
    const allocator = std.testing.allocator;
    const Config = @import("../Config.zig");
    const config = try Config.load(allocator);
    defer config.deinit(allocator);
    const build_options = @import("build_options");
    const db_path = config.geolite_path;

    if (build_options.download_geoip) {
        const GeoLite2 = @import("GeoLite2.zig");
        try GeoLite2.ensureDatabase(std.testing.allocator, db_path);
    }

    var geoip = GeoIP.init(std.testing.allocator, db_path, null, config.ip2location_cache_file) catch
        return error.SkipZigTest;
    defer geoip.deinit();

    // Test that the function doesn't crash with various IPs
    try std.testing.expect(geoip.isUSIp("73.158.64.1"));

    // Test invalid IP returns false
    const invalid = geoip.isUSIp("invalid");
    try std.testing.expect(!invalid);
}
test "lookup works" {
    const allocator = std.testing.allocator;
    const Config = @import("../Config.zig");
    const config = try Config.load(allocator);
    defer config.deinit(allocator);
    const build_options = @import("build_options");
    const db_path = config.geolite_path;

    if (build_options.download_geoip) {
        const GeoLite2 = @import("GeoLite2.zig");
        try GeoLite2.ensureDatabase(std.testing.allocator, db_path);
    }

    var geoip = GeoIP.init(std.testing.allocator, db_path, null, config.ip2location_cache_file) catch
        return error.SkipZigTest;
    defer geoip.deinit();

    // Test that the function doesn't crash with various IPs
    const maybe_coords = geoip.lookup("73.158.64.1");

    try std.testing.expect(maybe_coords != null);

    const coords = maybe_coords.?;
    try std.testing.expectEqual(@as(f64, 37.5958), coords.latitude);
}