118 lines
3.8 KiB
Zig
118 lines
3.8 KiB
Zig
const std = @import("std");
|
|
const Coordinates = @import("../Coordinates.zig");
|
|
|
|
const c = @cImport({
|
|
@cInclude("maxminddb.h");
|
|
});
|
|
|
|
const GeoIP = @This();
|
|
|
|
mmdb: c.MMDB_s,
|
|
|
|
pub fn init(db_path: []const u8) !GeoIP {
|
|
const path_z = try std.heap.c_allocator.dupeZ(u8, db_path);
|
|
defer std.heap.c_allocator.free(path_z);
|
|
|
|
// SAFETY: The C API will initialize this on the next line
|
|
var mmdb: c.MMDB_s = undefined;
|
|
const status = c.MMDB_open(path_z.ptr, c.MMDB_MODE_MMAP, &mmdb);
|
|
if (status != c.MMDB_SUCCESS)
|
|
return error.CannotOpenDatabase;
|
|
|
|
return GeoIP{ .mmdb = mmdb };
|
|
}
|
|
|
|
pub fn deinit(self: *GeoIP) void {
|
|
c.MMDB_close(&self.mmdb);
|
|
}
|
|
|
|
pub fn lookup(self: *GeoIP, ip: []const u8) !?Coordinates {
|
|
const result = lookupInternal(&self.mmdb, ip) catch return null;
|
|
|
|
if (!result.found_entry) return null;
|
|
|
|
return try self.extractCoordinates(result);
|
|
}
|
|
|
|
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 (gai_error != 0 or mmdb_error != 0) return error.MMDBLookupError;
|
|
return result;
|
|
}
|
|
|
|
pub fn isUSIP(self: *GeoIP, ip: []const u8) bool {
|
|
const result = lookupInternal(&self.mmdb, ip) catch return false;
|
|
if (!result.found_entry) return false;
|
|
|
|
var entry_mut = result.entry;
|
|
const null_term: [*:0]const u8 = @ptrCast(&[_]u8{0});
|
|
// 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(&entry_mut, &country_data, "country\x00", "iso_code\x00", null_term);
|
|
|
|
if (status != c.MMDB_SUCCESS or !country_data.has_data)
|
|
return false;
|
|
|
|
const country_code = std.mem.span(country_data.unnamed_0.utf8_string);
|
|
return std.mem.eql(u8, country_code, "US");
|
|
}
|
|
|
|
fn extractCoordinates(self: *GeoIP, result: c.MMDB_lookup_result_s) !Coordinates {
|
|
_ = self;
|
|
var entry_mut = result.entry;
|
|
|
|
// SAFETY: The C API will initialize below
|
|
var latitude_data: c.MMDB_entry_data_s = undefined;
|
|
// SAFETY: The C API will initialize below
|
|
var longitude_data: c.MMDB_entry_data_s = undefined;
|
|
|
|
const lat_status = c.MMDB_get_value(&entry_mut, &latitude_data, "location", "latitude", @as([*:0]const u8, @ptrCast(&[_]u8{0})));
|
|
const lon_status = c.MMDB_get_value(&entry_mut, &longitude_data, "location", "longitude", @as([*:0]const u8, @ptrCast(&[_]u8{0})));
|
|
|
|
if (lat_status != c.MMDB_SUCCESS or lon_status != c.MMDB_SUCCESS or !latitude_data.has_data or !longitude_data.has_data)
|
|
return error.CoordinatesNotFound;
|
|
|
|
return .{
|
|
.latitude = latitude_data.unnamed_0.double_value,
|
|
.longitude = longitude_data.unnamed_0.double_value,
|
|
};
|
|
}
|
|
|
|
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("/nonexistent/path.mmdb");
|
|
try std.testing.expectError(error.CannotOpenDatabase, result);
|
|
}
|
|
|
|
test "isUSIP detects US IPs" {
|
|
const build_options = @import("build_options");
|
|
const db_path = "./GeoLite2-City.mmdb";
|
|
|
|
if (build_options.download_geoip) {
|
|
const GeoLite2 = @import("GeoLite2.zig");
|
|
try GeoLite2.ensureDatabase(std.testing.allocator, db_path);
|
|
}
|
|
|
|
var geoip = GeoIP.init(db_path) catch {
|
|
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);
|
|
}
|