wttr/src/location/Ip2location.zig

const std = @import("std");
const Allocator = std.mem.Allocator;
const Location = @import("resolver.zig").Location;
const srf = @import("srf");

const Self = @This();

const log = std.log.scoped(.ip2location);

allocator: Allocator,
/// Zig 0.16 requires an explicit `Io` for both filesystem and HTTP work. It is
/// stored alongside the allocator so only `init` signatures change rather than
/// every method that touches a file or the network.
io: std.Io,
api_key: ?[]const u8,
http_client: std.http.Client,
cache: *Cache,

pub fn init(allocator: Allocator, io: std.Io, api_key: ?[]const u8, cache_path: []const u8) !Self {
    const cache = try allocator.create(Cache);
    errdefer allocator.destroy(cache);
    cache.* = try .init(allocator, io, cache_path);
    return .{
        .allocator = allocator,
        .io = io,
        .api_key = if (api_key) |k| try allocator.dupe(u8, k) else null,
        .http_client = std.http.Client{ .allocator = allocator, .io = io },
        .cache = cache,
    };
}

pub fn deinit(self: *Self) void {
    self.cache.deinit();
    self.allocator.destroy(self.cache);
    self.http_client.deinit();
    if (self.api_key) |k|
        self.allocator.free(k);
}

/// An IP address packed into a `u128` cache key, with the address family it
/// came from (4 or 6).
pub const PackedIp = struct {
    key: u128,
    family: u8,
};

/// Packs a textual IP address into a `u128` cache key.
///
/// Zig 0.16 replaced `std.net.Address` with `std.Io.net.IpAddress`, which
/// exposes the address bytes directly instead of requiring casts through
/// `sockaddr`. The union is exhaustive, so there is no longer an unreachable
/// "unknown family" case to handle.
///
/// IPv4-mapped IPv6 addresses (`::ffff:1.2.3.4`) are folded down to IPv4.
/// `IpAddress.parse` does not do this itself, and without it one address has two
/// representations: cache entries would be stored twice, and a CIDR pin written
/// for an IPv4 range would not match a client that arrived as a mapped address.
/// That is not hypothetical -- with no `X-Forwarded-For` header the client
/// address comes from the socket, and a dual-stack listener reports IPv4 peers
/// in mapped form.
pub fn packIp(ip_str: []const u8) ?PackedIp {
    const parsed = std.Io.net.IpAddress.parse(ip_str, 0) catch return null;
    const addr = switch (parsed) {
        .ip4 => parsed,
        .ip6 => |a| std.Io.net.IpAddress.fromIp6(a),
    };
    return switch (addr) {
        .ip4 => |a| .{ .key = std.mem.readInt(u32, &a.bytes, .big), .family = 4 },
        .ip6 => |a| .{ .key = std.mem.readInt(u128, &a.bytes, .big), .family = 6 },
    };
}

pub fn lookup(self: *Self, ip_str: []const u8) ?Location {
    // Parse IP to u128 for cache lookup
    const parsed = packIp(ip_str) orelse return null;
    const ip_u128 = parsed.key;
    const family = parsed.family;

    // Check cache first
    if (self.cache.get(ip_u128)) |result|
        return result;

    // Fetch from API
    const result = self.fetch(ip_str) catch |err| {
        log.err("API lookup failed: {}", .{err});
        return null;
    };

    // Store in cache
    self.cache.put(ip_u128, family, result) catch |err| {
        log.warn("Failed to cache result: {}", .{err});
    };

    return result;
}

fn fetch(self: *Self, ip_str: []const u8) !Location {
    log.info("Fetching geolocation for IP {s}", .{ip_str});

    if (@import("builtin").is_test) return error.LookupUnavailableInUnitTest;
    var buf: [256]u8 = undefined;
    var w = std.Io.Writer.fixed(&buf);
    // Build URL: https://api.ip2location.io/?key=XXX&ip=1.2.3.4
    try w.writeAll("https://api.ip2location.io/?ip=");
    try w.writeAll(ip_str);
    if (self.api_key) |key|
        try w.print("&key={s}", .{key});

    var response_buf: [4096]u8 = undefined;
    var writer = std.Io.Writer.fixed(&response_buf);
    const result = try self.http_client.fetch(.{
        .location = .{ .url = w.buffered() },
        .method = .GET,
        .response_writer = &writer,
    });

    if (result.status != .ok) {
        log.err("API returned status {}", .{result.status});
        return error.ApiError;
    }

    const response_body = response_buf[0..writer.end];

    // Parse JSON response
    const parsed = try std.json.parseFromSlice(
        std.json.Value,
        self.allocator,
        response_body,
        .{},
    );
    defer parsed.deinit();

    const obj = parsed.value.object;
    const lat = obj.get("latitude") orelse return error.MissingLatitude;
    const lon = obj.get("longitude") orelse return error.MissingLongitude;
    if (lat == .null) return error.MissingLatitude;
    if (lat != .float)
        log.err(
            "Latitude returned from ip2location.io for ip {s} is not a float: {f}",
            .{ ip_str, std.json.fmt(lat, .{}) },
        );
    if (lon == .null) return error.MissingLongitude;
    if (lon != .float)
        log.err(
            "Longitude returned from ip2location.io for ip {s} is not a float: {f}",
            .{ ip_str, std.json.fmt(lon, .{}) },
        );

    const city = getString(obj, "city_name");
    const region = getString(obj, "region_name");
    const country = getString(obj, "country_name");

    const display_name = Location.buildDisplayName(
        self.allocator,
        city,
        region,
        country,
        ip_str,
    );

    return Location{
        .allocator = self.allocator,
        .name = display_name,
        .coords = .{
            .latitude = @floatCast(lat.float),
            .longitude = @floatCast(lon.float),
        },
    };
}

inline fn getString(obj: std.json.ObjectMap, key: []const u8) []const u8 {
    const maybe_val = obj.get(key);
    if (maybe_val == null) return "";
    if (maybe_val.? != .string) return "";
    return maybe_val.?.string;
}

/// Permanent IP-to-location cache, shared by the online providers.
///
/// Append-only: an online lookup for an address is answered once and the result
/// is kept forever, because an address's city does not meaningfully change and
/// every avoided request is one that does not count against a rate limit.
///
/// Stored as SRF in compact form, one record per line. Compact is safe for place
/// names containing commas because SRF length-prefixes any string containing the
/// field delimiter, so no escaping rules are needed here.
pub const Cache = struct {
    allocator: Allocator,
    io: std.Io,
    path: []const u8,
    entries: std.AutoHashMap(u128, Location),
    file: ?std.Io.File,

    /// On-disk shape of one entry.
    ///
    /// The address is stored as text rather than the `u128` used for the in-memory
    /// key, so the file stays legible to whoever has to look at it.
    const Record = struct {
        ip: []const u8,
        lat: f64,
        lon: f64,
        name: []const u8,
    };

    pub fn init(allocator: Allocator, io: std.Io, path: []const u8) !Cache {
        var cache = Cache{
            .allocator = allocator,
            .io = io,
            .path = try allocator.dupe(u8, path),
            .entries = std.AutoHashMap(u128, Location).init(allocator),
            .file = null,
        };
        errdefer allocator.free(cache.path);
        errdefer cache.entries.deinit();

        // Try to open existing cache file
        if (std.Io.Dir.openFileAbsolute(io, path, .{ .mode = .read_write })) |file| {
            cache.file = file;
            cache.load() catch |err| {
                // A cache that cannot be read is worth strictly less than the
                // service staying up: start empty and let it refill.
                log.warn("could not read cache {s} ({t}); starting empty", .{ path, err });
            };
        } else |err| switch (err) {
            error.FileNotFound => {
                const dir = std.fs.path.dirname(path) orelse return error.InvalidPath;
                try std.Io.Dir.cwd().createDirPath(io, dir);
                cache.file = try std.Io.Dir.createFileAbsolute(io, path, .{ .read = true });
                try cache.writeDirectives();
            },
            else => return err,
        }

        return cache;
    }

    /// Writes the SRF front matter that opens a new cache file.
    ///
    /// Emitted once at creation; `put` appends bare records afterwards.
    fn writeDirectives(self: *Cache) !void {
        const file = self.file orelse return;
        var buf: [64]u8 = undefined;
        var fw = file.writer(self.io, &buf);
        try fw.interface.print("{f}", .{srf.fmt(Record, &.{}, .{})});
        try fw.interface.flush();
    }

    pub fn deinit(self: *Cache) void {
        if (self.file) |f| f.close(self.io);
        var it = self.entries.valueIterator();
        while (it.next()) |loc| {
            self.allocator.free(loc.name);
        }
        self.entries.deinit();
        self.allocator.free(self.path);
    }

    fn load(self: *Cache) !void {
        const file = self.file orelse return;
        const file_size = try file.length(self.io);
        if (file_size == 0) return;

        const content = try std.Io.Dir.cwd().readFileAlloc(
            self.io,
            self.path,
            self.allocator,
            // `Io.Limit` fails when the limit is *reached*, not merely exceeded,
            // so a limit of exactly `file_size` would reject a file of that
            // size. Allow one extra byte.
            .limited64(file_size + 1),
        );
        defer self.allocator.free(content);

        var reader = std.Io.Reader.fixed(content);
        var records = try srf.iterator(&reader, self.allocator, .{});
        defer records.deinit();

        var index: usize = 0;
        while (records.next() catch |err| {
            log.warn("stopped reading {s} after {d} entr(ies): {t}", .{ self.path, index, err });
            return;
        }) |fields| {
            index += 1;
            const record = fields.to(Record, .{}) catch continue;
            const parsed = packIp(record.ip) orelse continue;

            const name_copy = try self.allocator.dupe(u8, record.name);
            errdefer self.allocator.free(name_copy);

            // Replacing an existing key would leak the name it already owns.
            const existing = try self.entries.fetchPut(parsed.key, .{
                .allocator = self.allocator,
                .name = name_copy,
                .coords = .{ .latitude = record.lat, .longitude = record.lon },
            });
            if (existing) |old| self.allocator.free(old.value.name);
        }
    }

    pub fn get(self: *Cache, ip: u128) ?Location {
        const entry = self.entries.get(ip) orelse return null;
        return .{
            .allocator = self.allocator,
            .name = self.allocator.dupe(u8, entry.name) catch return null,
            .coords = entry.coords,
        };
    }

    pub fn put(self: *Cache, ip: u128, family: u8, loc: Location) !void {
        const name_copy = try self.allocator.dupe(u8, loc.name);
        errdefer self.allocator.free(name_copy);

        const existing = try self.entries.fetchPut(ip, .{
            .allocator = self.allocator,
            .name = name_copy,
            .coords = loc.coords,
        });
        if (existing) |old| self.allocator.free(old.value.name);

        const file = self.file orelse return;

        var ip_buf: [64]u8 = undefined;
        const ip_str = try formatKey(ip, family, &ip_buf);

        // Append one record with no front matter; the directives were written
        // when the file was created.
        var line_buf: [512]u8 = undefined;
        const line = try std.fmt.bufPrint(&line_buf, "{f}", .{srf.fmt(Record, &.{.{
            .ip = ip_str,
            .lat = loc.coords.latitude,
            .lon = loc.coords.longitude,
            .name = loc.name,
        }}, .{ .emit_directives = false })});

        // 0.16 has no seek+write; append at the current end offset.
        const end = try file.length(self.io);
        try file.writePositionalAll(self.io, line, end);
    }

    /// Renders a packed key back to text for storage.
    fn formatKey(ip: u128, family: u8, buf: []u8) ![]const u8 {
        if (family == 4) {
            return std.fmt.bufPrint(buf, "{d}.{d}.{d}.{d}", .{
                @as(u8, @truncate(ip >> 24)),
                @as(u8, @truncate(ip >> 16)),
                @as(u8, @truncate(ip >> 8)),
                @as(u8, @truncate(ip)),
            });
        }
        var bytes: [16]u8 = undefined;
        std.mem.writeInt(u128, &bytes, ip, .big);
        // `Ip6Address.format` would append ":port" and bracket the address;
        // `Unresolved` is the bare-address formatter.
        const bare: std.Io.net.Ip6Address.Unresolved = .{ .bytes = bytes, .interface_name = null };
        return std.fmt.bufPrint(buf, "{f}", .{&bare});
    }
};

test "Cache: round-trips IPv4 and IPv6 entries through the file" {
    const allocator = std.testing.allocator;
    const io = std.testing.io;

    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    var path_buf: [std.fs.max_path_bytes]u8 = undefined;
    const dir_len = try tmp.dir.realPath(io, &path_buf);
    const path = try std.fmt.allocPrint(allocator, "{s}/ip.cache", .{path_buf[0..dir_len]});
    defer allocator.free(path);

    const v4 = packIp("192.168.1.1").?;
    const v6 = packIp("2001:db8::1").?;

    {
        var cache = try Cache.init(allocator, io, path);
        defer cache.deinit();

        // A name with commas is the interesting case: compact SRF delimits on
        // commas, so this only survives via length prefixing.
        try cache.put(v4.key, v4.family, .{
            .allocator = allocator,
            .name = "San Francisco, California, United States",
            .coords = .{ .latitude = 37.5, .longitude = -122.5 },
        });
        try cache.put(v6.key, v6.family, .{
            .allocator = allocator,
            .name = "London, United Kingdom",
            .coords = .{ .latitude = 51.5, .longitude = -0.1 },
        });
    }

    var reopened = try Cache.init(allocator, io, path);
    defer reopened.deinit();

    const got4 = reopened.get(v4.key) orelse return error.TestUnexpectedResult;
    defer got4.deinit();
    try std.testing.expectEqualStrings("San Francisco, California, United States", got4.name);
    try std.testing.expectEqual(@as(f64, 37.5), got4.coords.latitude);
    try std.testing.expectEqual(@as(f64, -122.5), got4.coords.longitude);

    const got6 = reopened.get(v6.key) orelse return error.TestUnexpectedResult;
    defer got6.deinit();
    try std.testing.expectEqualStrings("London, United Kingdom", got6.name);
    try std.testing.expectEqual(@as(f64, 51.5), got6.coords.latitude);
}

test "Cache: a legacy or corrupt file starts empty instead of failing" {
    const allocator = std.testing.allocator;
    const io = std.testing.io;

    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    var path_buf: [std.fs.max_path_bytes]u8 = undefined;
    const dir_len = try tmp.dir.realPath(io, &path_buf);
    const path = try std.fmt.allocPrint(allocator, "{s}/ip.cache", .{path_buf[0..dir_len]});
    defer allocator.free(path);

    // The previous hand-rolled format. Deliberately not migrated: the entries are
    // re-fetchable and losing them costs a handful of API calls.
    try std.Io.Dir.cwd().writeFile(io, .{
        .sub_path = path,
        .data = "#Ip2location:v2\n192.168.1.1,37.5,-122.5,San Francisco\n",
    });

    var cache = try Cache.init(allocator, io, path);
    defer cache.deinit();

    try std.testing.expectEqual(@as(usize, 0), cache.entries.count());

    // Still usable afterwards, which is the point of not treating it as fatal.
    const v4 = packIp("10.0.0.1").?;
    try cache.put(v4.key, v4.family, .{
        .allocator = allocator,
        .name = "Somewhere",
        .coords = .{ .latitude = 1, .longitude = 2 },
    });
    const got = cache.get(v4.key) orelse return error.TestUnexpectedResult;
    defer got.deinit();
    try std.testing.expectEqualStrings("Somewhere", got.name);
}

test "Cache: formatKey renders both families without a port" {
    var buf: [64]u8 = undefined;
    const v4 = packIp("12.94.132.170").?;
    try std.testing.expectEqualStrings("12.94.132.170", try Cache.formatKey(v4.key, v4.family, &buf));

    const v6 = packIp("2001:db8::1").?;
    try std.testing.expectEqualStrings("2001:db8::1", try Cache.formatKey(v6.key, v6.family, &buf));
}