Add Country and ASN records with basic tests

marselester · marselester · commit 0b0cf765c848 · 2025-02-22T20:36:49.000-05:00
diff --git a/src/geolite2.zig b/src/geolite2.zig
@@ -1,58 +1,97 @@
 const std = @import("std");
 
-// City represents a record in the GeoLite2-City database, for example,
-// https://github.com/maxmind/MaxMind-DB/blob/main/source-data/GeoLite2-City-Test.json.
-pub const City = struct {
-    city: struct {
-        geoname_id: u32 = 0,
-        names: ?std.hash_map.StringHashMap([]const u8) = null,
-    },
-    continent: struct {
+pub const StringMap = std.hash_map.StringHashMap([]const u8);
+
+/// Country represents a record in the GeoLite2-Country database, for example,
+/// https://github.com/maxmind/MaxMind-DB/blob/main/source-data/GeoLite2-Country-Test.json.
+/// It can be used for geolocation at the country-level for analytics, content customization,
+/// or compliance use cases in territories that are not disputed.
+pub const Country = struct {
+    continent: Self.Continent,
+    country: Self.Country,
+    registered_country: Self.Country,
+    represented_country: Self.RepresentedCountry,
+
+    _arena: std.heap.ArenaAllocator,
+
+    const Self = @This();
+    pub const Continent = struct {
         code: []const u8 = "",
         geoname_id: u32 = 0,
-        names: ?std.hash_map.StringHashMap([]const u8) = null,
-    },
-    country: struct {
+        names: ?StringMap = null,
+    };
+    pub const Country = struct {
         geoname_id: u32 = 0,
         is_in_european_union: bool = false,
         iso_code: []const u8 = "",
-        names: ?std.hash_map.StringHashMap([]const u8) = null,
-    },
-    location: struct {
+        names: ?StringMap = null,
+    };
+    pub const RepresentedCountry = struct {
+        geoname_id: u32 = 0,
+        is_in_european_union: bool = false,
+        iso_code: []const u8 = "",
+        names: ?StringMap = null,
+        type: []const u8 = "",
+    };
+
+    pub fn init(allocator: std.mem.Allocator) Self {
+        const arena = std.heap.ArenaAllocator.init(allocator);
+
+        return .{
+            .continent = .{},
+            .country = .{},
+            .registered_country = .{},
+            .represented_country = .{},
+
+            ._arena = arena,
+        };
+    }
+
+    pub fn deinit(self: *const Self) void {
+        self._arena.deinit();
+    }
+};
+
+/// City represents a record in the GeoLite2-City database, for example,
+/// https://github.com/maxmind/MaxMind-DB/blob/main/source-data/GeoLite2-City-Test.json.
+/// It can be used for geolocation down to the city or postal code for analytics and content customization.
+pub const City = struct {
+    city: Self.City,
+    continent: Country.Continent,
+    country: Country.Country,
+    location: Self.Location,
+    postal: Self.Postal,
+    registered_country: Country.Country,
+    represented_country: Country.RepresentedCountry,
+    subdivisions: ?std.ArrayList(Self.Subdivision) = null,
+
+    _arena: std.heap.ArenaAllocator,
+
+    const Self = @This();
+    pub const City = struct {
+        geoname_id: u32 = 0,
+        names: ?StringMap = null,
+    };
+    pub const Location = struct {
         accuracy_radius: u16 = 0,
         latitude: f64 = 0,
         longitude: f64 = 0,
         metro_code: u16 = 0,
         time_zone: []const u8 = "",
-    },
-    postal: struct {
+    };
+    pub const Postal = struct {
         code: []const u8 = "",
-    },
-    registered_country: struct {
+    };
+    pub const Subdivision = struct {
         geoname_id: u32 = 0,
-        is_in_european_union: bool = false,
         iso_code: []const u8 = "",
-        names: ?std.hash_map.StringHashMap([]const u8) = null,
-    },
-    represented_country: struct {
-        geoname_id: u32 = 0,
-        is_in_european_union: bool = false,
-        iso_code: []const u8 = "",
-        names: ?std.hash_map.StringHashMap([]const u8) = null,
-        type: []const u8 = "",
-    },
-    subdivisions: ?std.ArrayList(struct {
-        geoname_id: u32 = 0,
-        iso_code: []const u8 = "",
-        names: ?std.hash_map.StringHashMap([]const u8) = null,
-    }) = null,
-
-    _arena: std.heap.ArenaAllocator,
+        names: ?StringMap = null,
+    };
 
-    pub fn init(allocator: std.mem.Allocator) City {
+    pub fn init(allocator: std.mem.Allocator) Self {
         const arena = std.heap.ArenaAllocator.init(allocator);
 
-        return City{
+        return .{
             .city = .{},
             .continent = .{},
             .country = .{},
@@ -65,7 +104,14 @@ pub const City = struct {
         };
     }
 
-    pub fn deinit(self: *const City) void {
+    pub fn deinit(self: *const Self) void {
         self._arena.deinit();
     }
 };
+
+/// Provides the autonomous system number and organization for IP addresses for analytics,
+/// e.g., https://github.com/maxmind/MaxMind-DB/blob/main/source-data/GeoLite2-ASN-Test.json.
+pub const ASN = struct {
+    autonomous_system_number: u32 = 0,
+    autonomous_system_organization: []const u8 = "",
+};
diff --git a/src/maxminddb.zig b/src/maxminddb.zig
@@ -1,3 +1,5 @@
+const std = @import("std");
+
 const reader = @import("reader.zig");
 const decoder = @import("decoder.zig");
 const net = @import("net.zig");
@@ -10,5 +12,150 @@ pub const Metadata = reader.Metadata;
 pub const Network = net.Network;
 
 test {
-    @import("std").testing.refAllDecls(@This());
+    std.testing.refAllDecls(@This());
+}
+
+fn expectEqualMaps(
+    map: geolite2.StringMap,
+    keys: []const []const u8,
+    values: []const []const u8,
+) !void {
+    try std.testing.expectEqual(map.count(), keys.len);
+
+    for (keys, values) |key, want_value| {
+        const got_value = map.get(key) orelse {
+            std.debug.print("map key=\"{s}\" was not found\n", .{key});
+            return error.MapKeyNotFound;
+        };
+        try std.testing.expectEqualStrings(want_value, got_value);
+    }
+}
+
+const expectEqual = std.testing.expectEqual;
+const expectEqualStrings = std.testing.expectEqualStrings;
+
+test "GeoLite2 Country" {
+    var db = try Reader.open_mmap(
+        std.testing.allocator,
+        "test-data/test-data/GeoLite2-Country-Test.mmdb",
+    );
+    defer db.close();
+
+    const ip = try std.net.Address.parseIp("89.160.20.128", 0);
+    const got = try db.lookup(geolite2.Country, &ip);
+    defer got.deinit();
+
+    try expectEqualStrings("EU", got.continent.code);
+    try expectEqual(6255148, got.continent.geoname_id);
+    try expectEqualMaps(
+        got.continent.names.?,
+        &.{ "de", "en", "es", "fr", "ja", "pt-BR", "ru", "zh-CN" },
+        &.{ "Europa", "Europe", "Europa", "Europe", "ヨーロッパ", "Europa", "Европа", "欧洲" },
+    );
+
+    try expectEqual(2661886, got.country.geoname_id);
+    try expectEqual(true, got.country.is_in_european_union);
+    try expectEqualStrings("SE", got.country.iso_code);
+    try expectEqualMaps(
+        got.country.names.?,
+        &.{ "de", "en", "es", "fr", "ja", "pt-BR", "ru", "zh-CN" },
+        &.{ "Schweden", "Sweden", "Suecia", "Suède", "スウェーデン王国", "Suécia", "Швеция", "瑞典" },
+    );
+
+    try expectEqual(2921044, got.registered_country.geoname_id);
+    try expectEqual(true, got.registered_country.is_in_european_union);
+    try expectEqualStrings("DE", got.registered_country.iso_code);
+    try expectEqualMaps(
+        got.registered_country.names.?,
+        &.{ "de", "en", "es", "fr", "ja", "pt-BR", "ru", "zh-CN" },
+        &.{ "Deutschland", "Germany", "Alemania", "Allemagne", "ドイツ連邦共和国", "Alemanha", "Германия", "德国" },
+    );
+
+    try std.testing.expectEqualDeep(geolite2.Country.RepresentedCountry{}, got.represented_country);
+}
+
+test "GeoLite2 City" {
+    var db = try Reader.open_mmap(
+        std.testing.allocator,
+        "test-data/test-data/GeoLite2-City-Test.mmdb",
+    );
+    defer db.close();
+
+    const ip = try std.net.Address.parseIp("89.160.20.128", 0);
+    const got = try db.lookup(geolite2.City, &ip);
+    defer got.deinit();
+
+    try expectEqual(2694762, got.city.geoname_id);
+    try expectEqualMaps(
+        got.city.names.?,
+        &.{ "de", "en", "fr", "ja", "zh-CN" },
+        &.{ "Linköping", "Linköping", "Linköping", "リンシェーピング", "林雪平" },
+    );
+
+    try expectEqualStrings("EU", got.continent.code);
+    try expectEqual(6255148, got.continent.geoname_id);
+    try expectEqualMaps(
+        got.continent.names.?,
+        &.{ "de", "en", "es", "fr", "ja", "pt-BR", "ru", "zh-CN" },
+        &.{ "Europa", "Europe", "Europa", "Europe", "ヨーロッパ", "Europa", "Европа", "欧洲" },
+    );
+
+    try expectEqual(2661886, got.country.geoname_id);
+    try expectEqual(true, got.country.is_in_european_union);
+    try expectEqualStrings("SE", got.country.iso_code);
+    try expectEqualMaps(
+        got.country.names.?,
+        &.{ "de", "en", "es", "fr", "ja", "pt-BR", "ru", "zh-CN" },
+        &.{ "Schweden", "Sweden", "Suecia", "Suède", "スウェーデン王国", "Suécia", "Швеция", "瑞典" },
+    );
+
+    try std.testing.expectEqualDeep(
+        geolite2.City.Location{
+            .accuracy_radius = 76,
+            .latitude = 58.4167,
+            .longitude = 15.6167,
+            .time_zone = "Europe/Stockholm",
+        },
+        got.location,
+    );
+
+    try std.testing.expectEqualDeep(geolite2.City.Postal{}, got.postal);
+
+    try expectEqual(2921044, got.registered_country.geoname_id);
+    try expectEqual(true, got.registered_country.is_in_european_union);
+    try expectEqualStrings("DE", got.registered_country.iso_code);
+    try expectEqualMaps(
+        got.registered_country.names.?,
+        &.{ "de", "en", "es", "fr", "ja", "pt-BR", "ru", "zh-CN" },
+        &.{ "Deutschland", "Germany", "Alemania", "Allemagne", "ドイツ連邦共和国", "Alemanha", "Германия", "德国" },
+    );
+
+    try std.testing.expectEqualDeep(geolite2.Country.RepresentedCountry{}, got.represented_country);
+
+    try expectEqual(1, got.subdivisions.?.items.len);
+    const sub = got.subdivisions.?.getLast();
+    try expectEqual(2685867, sub.geoname_id);
+    try expectEqualStrings("E", sub.iso_code);
+    try expectEqualMaps(
+        sub.names.?,
+        &.{ "en", "fr" },
+        &.{ "Östergötland County", "Comté d'Östergötland" },
+    );
+}
+
+test "GeoLite2 ASN" {
+    var db = try Reader.open_mmap(
+        std.testing.allocator,
+        "test-data/test-data/GeoLite2-ASN-Test.mmdb",
+    );
+    defer db.close();
+
+    const ip = try std.net.Address.parseIp("89.160.20.128", 0);
+    const got = try db.lookup(geolite2.ASN, &ip);
+
+    const want = geolite2.ASN{
+        .autonomous_system_number = 29518,
+        .autonomous_system_organization = "Bredband2 AB",
+    };
+    try std.testing.expectEqualDeep(want, got);
 }
