wttr/src/location/GeoIp.zig

const std = @import("std");
const Ip2location = @import("Ip2location.zig");
const IpWhoIs = @import("IpWhoIs.zig");
const Location = @import("resolver.zig").Location;
const Config = @import("../Config.zig");

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

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

// libmaxminddb writes into `MMDB_entry_data_s` values that we declare on the
// Zig stack, so Zig's view of that struct must match the C compiler's exactly.
//
// It did not: `MMDB_UINT128_USING_MODE` types `mmdb_uint128_t` as
// `unsigned int __attribute__((__mode__(TI)))`, which translate-c cannot
// represent. Zig saw 32 bytes / 8-byte alignment where clang saw 48 / 16, so
// every `MMDB_get_value` call wrote `entry_data->offset` (C offset 32) past the
// end of the variable Zig had reserved -- silent stack corruption whose blast
// radius depended on the adjacent stack slot.
//
// `build.zig` now selects `MMDB_UINT128_IS_BYTE_ARRAY`, under which both
// compilers agree. This assertion fails the build if that ever regresses,
// rather than letting the corruption return unnoticed.
comptime {
    std.debug.assert(@sizeOf(c.MMDB_entry_data_s) == 40);
    std.debug.assert(@alignOf(c.MMDB_entry_data_s) == 8);
}

const FallbackClient = union(enum) {
    ip2location: *Ip2location,
    ipwhois: *IpWhoIs,

    fn lookup(self: FallbackClient, ip: []const u8) ?Location {
        return switch (self) {
            .ip2location => |client| client.lookup(ip),
            .ipwhois => |client| client.lookup(ip),
        };
    }

    fn deinit(self: FallbackClient, allocator: std.mem.Allocator) void {
        switch (self) {
            .ip2location => |client| {
                client.deinit();
                allocator.destroy(client);
            },
            .ipwhois => |client| {
                client.deinit();
                allocator.destroy(client);
            },
        }
    }
};

mmdb: *c.MMDB_s,
fallback_client: FallbackClient,
allocator: std.mem.Allocator,

pub fn init(allocator: std.mem.Allocator, io: std.Io, db_path: []const u8, config: Config) !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 fallback_client: FallbackClient = switch (config.geoip_fallback) {
        .ip2location => blk: {
            const client = try allocator.create(Ip2location);
            errdefer allocator.destroy(client);
            client.* = try Ip2location.init(allocator, io, config.ip2location_api_key, config.ip2location_cache_file);
            std.log.info(
                "GeoIP fallback: IP2Location ({s}, cache: {s})",
                .{ if (config.ip2location_api_key) |_| "key provided, 50k/mo limit" else "no key, 1k/day limit", config.ip2location_cache_file },
            );
            break :blk .{ .ip2location = client };
        },
        .ipwhois => blk: {
            const client = try allocator.create(IpWhoIs);
            errdefer allocator.destroy(client);
            client.* = try IpWhoIs.init(allocator, io, config.ipwhois_cache_file);
            std.log.info("GeoIP fallback: ipwho.is (cache: {s})", .{config.ipwhois_cache_file});
            break :blk .{ .ipwhois = client };
        },
    };

    return GeoIP{
        .mmdb = mmdb,
        .fallback_client = fallback_client,
        .allocator = allocator,
    };
}

pub fn deinit(self: *GeoIP) void {
    c.MMDB_close(self.mmdb);
    self.allocator.destroy(self.mmdb);
    self.fallback_client.deinit(self.allocator);
}

pub fn lookup(self: *GeoIP, ip: []const u8) ?Location {
    // 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 configured online provider
    return self.fallback_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");
}

/// Maximum accuracy radius (in km) to trust from GeoLite2. Entries with a
/// radius above this are too coarse for weather lookups (e.g. backbone/transit
/// IPs that MaxMind maps to the wrong city) and should fall back to IP2Location.
const max_accuracy_radius_km = 200;

fn extractCoordinates(self: *GeoIP, ip: []const u8, result: c.MMDB_lookup_result_s) ?Location {
    if (!result.found_entry) return null;

    var entry_copy = result.entry;

    // Check accuracy_radius first -- reject low-confidence entries so we
    // fall back to the IP2Location online lookup instead.
    // SAFETY: accuracy_data set by MMDB_get_value
    var accuracy_data: c.MMDB_entry_data_s = undefined;
    const acc_status = c.MMDB_get_value(&entry_copy, &accuracy_data, "location", "accuracy_radius", @as([*c]const u8, null));
    if (acc_status == c.MMDB_SUCCESS and accuracy_data.has_data) {
        const radius = accuracy_data.unnamed_0.uint16;
        if (radius > max_accuracy_radius_km) {
            log.info("GeoLite2 accuracy_radius for ip {s} is {d} km (>{d} km threshold), falling back to IP2Location", .{ ip, radius, max_accuracy_radius_km });
            return null;
        }
    }

    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];

    // Extract location name parts
    // SAFETY: value set by MMDB_get_value
    var city_data: c.MMDB_entry_data_s = undefined;
    entry_copy = result.entry;
    const city_status = c.MMDB_get_value(&entry_copy, &city_data, "city", "names", "en", @as([*c]const u8, null));
    const city = if (city_status == c.MMDB_SUCCESS and city_data.has_data)
        city_data.unnamed_0.utf8_string[0..city_data.data_size]
    else
        "";

    // SAFETY: value set by MMDB_get_value
    var subdivision_data: c.MMDB_entry_data_s = undefined;
    entry_copy = result.entry;
    const subdivision_status = c.MMDB_get_value(&entry_copy, &subdivision_data, "subdivisions", "0", "names", "en", @as([*c]const u8, null));
    const subdivision = if (subdivision_status == c.MMDB_SUCCESS and subdivision_data.has_data)
        subdivision_data.unnamed_0.utf8_string[0..subdivision_data.data_size]
    else
        "";

    // SAFETY: value set by MMDB_get_value
    var country_data: c.MMDB_entry_data_s = undefined;
    entry_copy = result.entry;
    const country_status = c.MMDB_get_value(&entry_copy, &country_data, "country", "names", "en", @as([*c]const u8, null));
    const country = if (country_status == c.MMDB_SUCCESS and country_data.has_data)
        country_data.unnamed_0.utf8_string[0..country_data.data_size]
    else
        "";

    const final_name = Location.buildDisplayName(self.allocator, city, subdivision, country, ip);

    return .{
        .allocator = self.allocator,
        .name = final_name,
        .coords = .{
            .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 config = try Config.loadForTest(std.testing.allocator);
    defer config.deinit(std.testing.allocator);
    const result = GeoIP.init(std.testing.allocator, std.testing.io, "/nonexistent/path.mmdb", config);
    try std.testing.expectError(error.CannotOpenDatabase, result);
}

test "isUSIp detects US IPs" {
    const allocator = std.testing.allocator;
    const config = try Config.loadForTest(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, std.testing.io, db_path);
    }

    var geoip = GeoIP.init(std.testing.allocator, std.testing.io, db_path, config) 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 = try Config.loadForTest(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, std.testing.io, db_path);
    }

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

    // Test that lookup returns a valid location for a well-known residential IP.
    // We don't assert exact values since the GeoLite2 database is fetched from
    // the latest upstream release and city/coordinate mappings change over time.
    const maybe_result = geoip.lookup("73.158.64.1");

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

    const result = maybe_result.?;
    defer result.deinit();
    try std.testing.expect(result.coords.latitude > 37.0 and result.coords.latitude < 38.0);
    try std.testing.expect(result.coords.longitude < -121.0 and result.coords.longitude > -123.0);
    try std.testing.expect(result.name.len > 0);
}