wttr/src/location/geoip.zig

const std = @import("std");

pub const Coordinates = struct {
    latitude: f64,
    longitude: f64,
};

pub const MMDB = extern struct {
    filename: [*:0]const u8,
    flags: u32,
    file_content: ?*anyopaque,
    file_size: usize,
    data_section: ?*anyopaque,
    data_section_size: u32,
    metadata_section: ?*anyopaque,
    metadata_section_size: u32,
    full_record_byte_size: u16,
    depth: u16,
    ipv4_start_node: extern struct {
        node_value: u32,
        netmask: u16,
    },
    metadata: extern struct {
        node_count: u32,
        record_size: u16,
        ip_version: u16,
        database_type: [*:0]const u8,
        languages: extern struct {
            count: usize,
            names: [*][*:0]const u8,
        },
        binary_format_major_version: u16,
        binary_format_minor_version: u16,
        build_epoch: u64,
        description: extern struct {
            count: usize,
            descriptions: [*]?*anyopaque,
        },
    },
};

pub const MMDBLookupResult = extern struct {
    found_entry: bool,
    entry: MMDBEntry,
    netmask: u16,
};

pub const MMDBEntry = extern struct {
    mmdb: *MMDB,
    offset: u32,
};

pub const MMDBEntryData = extern struct {
    has_data: bool,
    data_type: u32,
    offset: u32,
    offset_to_next: u32,
    data_size: u32,
    utf8_string: [*:0]const u8,
    double_value: f64,
    bytes: [*]const u8,
    uint16: u16,
    uint32: u32,
    int32: i32,
    uint64: u64,
    uint128: u128,
    boolean: bool,
    float_value: f32,
};

extern fn MMDB_open(filename: [*:0]const u8, flags: u32, mmdb: *MMDB) c_int;
extern fn MMDB_close(mmdb: *MMDB) void;
extern fn MMDB_lookup_string(mmdb: *MMDB, ipstr: [*:0]const u8, gai_error: *c_int, mmdb_error: *c_int) MMDBLookupResult;
extern fn MMDB_get_value(entry: *MMDBEntry, entry_data: *MMDBEntryData, ...) c_int;
extern fn MMDB_strerror(error_code: c_int) [*:0]const u8;

pub const GeoIP = struct {
    mmdb: MMDB,

    pub fn init(db_path: []const u8) !GeoIP {
        var mmdb: MMDB = undefined;
        const path_z = try std.heap.c_allocator.dupeZ(u8, db_path);
        defer std.heap.c_allocator.free(path_z);

        const status = MMDB_open(path_z.ptr, 0, &mmdb);
        if (status != 0) {
            return error.CannotOpenDatabase;
        }

        return GeoIP{ .mmdb = mmdb };
    }

    pub fn deinit(self: *GeoIP) void {
        MMDB_close(&self.mmdb);
    }

    pub fn lookup(self: *GeoIP, ip: []const u8) !?Coordinates {
        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 = MMDB_lookup_string(&self.mmdb, ip_z.ptr, &gai_error, &mmdb_error);

        if (gai_error != 0 or mmdb_error != 0) {
            return null;
        }

        if (!result.found_entry) {
            return null;
        }

        return try self.extractCoordinates(result.entry);
    }

    pub fn isUSIP(self: *GeoIP, ip: []const u8) bool {
        const ip_z = std.heap.c_allocator.dupeZ(u8, ip) catch return false;
        defer std.heap.c_allocator.free(ip_z);

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

        const result = MMDB_lookup_string(&self.mmdb, ip_z.ptr, &gai_error, &mmdb_error);

        if (gai_error != 0 or mmdb_error != 0 or !result.found_entry) {
            return false;
        }

        var entry_mut = result.entry;
        var country_data: MMDBEntryData = undefined;
        const null_term: [*:0]const u8 = @ptrCast(&[_]u8{0});
        const status = MMDB_get_value(&entry_mut, &country_data, "country\x00", "iso_code\x00", null_term);

        if (status != 0 or !country_data.has_data) {
            return false;
        }

        const country_code = std.mem.span(country_data.utf8_string);
        return std.mem.eql(u8, country_code, "US");
    }

    fn extractCoordinates(self: *GeoIP, entry: MMDBEntry) !Coordinates {
        _ = self;
        var entry_mut = entry;
        var latitude_data: MMDBEntryData = undefined;
        var longitude_data: MMDBEntryData = undefined;

        const lat_status = MMDB_get_value(&entry_mut, &latitude_data, "location", "latitude", @as([*:0]const u8, @ptrCast(&[_]u8{0})));
        const lon_status = MMDB_get_value(&entry_mut, &longitude_data, "location", "longitude", @as([*:0]const u8, @ptrCast(&[_]u8{0})));

        if (lat_status != 0 or lon_status != 0 or !latitude_data.has_data or !longitude_data.has_data) {
            return error.CoordinatesNotFound;
        }

        return Coordinates{
            .latitude = latitude_data.double_value,
            .longitude = longitude_data.double_value,
        };
    }
};

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

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

test "isUSIP detects US IPs" {
    var geoip = GeoIP.init("./GeoLite2-City.mmdb") catch {
        std.debug.print("Skipping test - GeoLite2-City.mmdb not found\n", .{});
        return error.SkipZigTest;
    };
    defer geoip.deinit();

    // Test that the function doesn't crash with various IPs
    _ = geoip.isUSIP("8.8.8.8");
    _ = geoip.isUSIP("1.1.1.1");

    // Test invalid IP returns false
    const invalid = geoip.isUSIP("invalid");
    try std.testing.expect(!invalid);
}