wttr/src/cli/pin.zig

const std = @import("std");
const builtin = @import("builtin");
const Config = @import("../Config.zig");
const Pins = @import("../location/Pins.zig");
const GeoCache = @import("../location/GeoCache.zig");
const Resolver = @import("../location/resolver.zig").Resolver;

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

/// Operator commands for managing IP pins.
///
/// A pin is an IP range to location override consulted ahead of GeoLite2. It
/// exists because GeoLite2 can be confidently wrong: a corporate WAN address may
/// carry `accuracy_radius = 20` while pointing at a city 1400 miles away, and
/// because the online fallback only runs when GeoLite2 declines to answer,
/// nothing automatic ever corrects it.
///
/// After writing the file these commands signal any running server so the change
/// takes effect without a restart. The database is mapped and the pins are held
/// in memory, so editing files alone would not be noticed.
pub fn printUsage(io: std.Io, exe: []const u8) !u8 {
    var buf: [2048]u8 = undefined;
    var out = std.Io.File.stdout().writer(io, &buf);
    const w = &out.interface;

    const base = std.fs.path.basename(exe);
    try w.print(
        \\usage: {s} [command]
        \\
        \\With no command, runs the weather server.
        \\
        \\Commands:
        \\  pin <cidr> <location>   Override the location for an IP or range
        \\  unpin <cidr>            Remove an override
        \\  pins                    List overrides
        \\  help                    Show this message
        \\
        \\Pins take precedence over the GeoLite2 database, which is the point:
        \\they exist for addresses the database resolves confidently and wrongly.
        \\
        \\Examples:
        \\  {s} pin 12.94.132.0/24 San Francisco
        \\  {s} pin 12.94.132.170 Portland, Oregon
        \\  {s} pin 2001:db8::/32 London
        \\  {s} unpin 12.94.132.0/24
        \\
        \\A bare address is treated as a single host (/32 or /128). When ranges
        \\overlap, the longest prefix wins, so a host pin beats a range pin.
        \\
    , .{ base, base, base, base, base });
    try w.flush();
    return 0;
}

/// Joins the remaining argv into a Nominatim-friendly query.
///
/// Spaces become `+` because the geocoding path builds a URL query directly, so
/// both `pin ... San Francisco` and `pin ... "San Francisco"` work.
fn joinLocation(allocator: std.mem.Allocator, parts: []const []const u8) ![]u8 {
    var list: std.ArrayList(u8) = .empty;
    errdefer list.deinit(allocator);

    for (parts, 0..) |part, i| {
        if (i > 0) try list.append(allocator, '+');
        for (part) |ch| try list.append(allocator, if (ch == ' ') '+' else ch);
    }
    return list.toOwnedSlice(allocator);
}

pub fn runPin(
    allocator: std.mem.Allocator,
    io: std.Io,
    cfg: Config,
    args: []const []const u8,
) !u8 {
    var err_buf: [512]u8 = undefined;
    var err_out = std.Io.File.stderr().writer(io, &err_buf);
    const ew = &err_out.interface;

    if (args.len < 2) {
        try ew.writeAll("usage: wttr pin <cidr> <location>\n");
        try ew.flush();
        return 2;
    }

    const cidr = Pins.parseCidr(args[0]) catch |e| {
        try ew.print("invalid address or range {s}: {t}\n", .{ args[0], e });
        try ew.flush();
        return 2;
    };

    const query = try joinLocation(allocator, args[1..]);
    defer allocator.free(query);

    // Geocode now rather than storing a bare name, for two reasons: it validates
    // the location while the operator is present to see the error, and it keeps
    // the lookup path free of network calls.
    var geocache = try GeoCache.init(allocator, io, cfg.geocache_file);
    defer geocache.deinit();
    var resolver = Resolver.init(allocator, io, null, &geocache, null);

    const location = resolver.resolve(query) catch |e| {
        try ew.print("could not resolve location {s}: {t}\n", .{ query, e });
        try ew.flush();
        return 1;
    };
    defer location.deinit();

    var pins = try Pins.load(allocator, io, cfg.pins_file);
    defer pins.deinit();

    try pins.put(cidr, location.name, location.coords);
    try pins.save(io, cfg.pins_file);

    var out_buf: [512]u8 = undefined;
    var out = std.Io.File.stdout().writer(io, &out_buf);
    const w = &out.interface;
    try w.print("pinned {s} -> {s} ({d:.4}, {d:.4})\n", .{
        args[0],
        location.name,
        location.coords.latitude,
        location.coords.longitude,
    });
    try w.print("wrote {s}\n", .{cfg.pins_file});
    try reportSignal(w, notifyServers(io));
    try w.flush();
    return 0;
}

pub fn runUnpin(
    allocator: std.mem.Allocator,
    io: std.Io,
    cfg: Config,
    args: []const []const u8,
) !u8 {
    var err_buf: [512]u8 = undefined;
    var err_out = std.Io.File.stderr().writer(io, &err_buf);
    const ew = &err_out.interface;

    if (args.len != 1) {
        try ew.writeAll("usage: wttr unpin <cidr>\n");
        try ew.flush();
        return 2;
    }

    const cidr = Pins.parseCidr(args[0]) catch |e| {
        try ew.print("invalid address or range {s}: {t}\n", .{ args[0], e });
        try ew.flush();
        return 2;
    };

    var pins = try Pins.load(allocator, io, cfg.pins_file);
    defer pins.deinit();

    if (!pins.remove(cidr)) {
        try ew.print("no pin found for {s}\n", .{args[0]});
        try ew.flush();
        return 1;
    }
    try pins.save(io, cfg.pins_file);

    var out_buf: [512]u8 = undefined;
    var out = std.Io.File.stdout().writer(io, &out_buf);
    const w = &out.interface;
    try w.print("removed pin {s}\n", .{args[0]});
    try reportSignal(w, notifyServers(io));
    try w.flush();
    return 0;
}

pub fn runList(allocator: std.mem.Allocator, io: std.Io, cfg: Config) !u8 {
    var pins = try Pins.load(allocator, io, cfg.pins_file);
    defer pins.deinit();

    var out_buf: [4096]u8 = undefined;
    var out = std.Io.File.stdout().writer(io, &out_buf);
    const w = &out.interface;

    if (pins.entries.items.len == 0) {
        try w.print("no pins configured ({s})\n", .{cfg.pins_file});
        try w.flush();
        return 0;
    }

    var cidr_buf: [64]u8 = undefined;
    for (pins.entries.items) |e| {
        const text = try Pins.formatCidr(
            .{ .family = e.family, .network = e.network, .prefix_len = e.prefix_len },
            &cidr_buf,
        );
        try w.print("{s: <20} {s} ({d:.4}, {d:.4})\n", .{
            text, e.name, e.coords.latitude, e.coords.longitude,
        });
    }
    try w.flush();
    return 0;
}

fn reportSignal(w: *std.Io.Writer, signaled: usize) !void {
    switch (signaled) {
        // Worth saying plainly: the most likely cause is running this outside the
        // container the server lives in, where the pins file being written is not
        // the one the server reads.
        0 => try w.writeAll(
            \\signaled 0 running servers (change applies on next start)
            \\  if the server runs in a container, run this inside it, e.g.
            \\  docker exec <container> /wttr ...
            \\
        ),
        1 => try w.writeAll("signaled 1 running server\n"),
        else => try w.print("signaled {d} running servers\n", .{signaled}),
    }
}

/// Sends SIGHUP to every running server, returning how many were signalled.
///
/// Scanning `/proc` rather than tracking a pidfile keeps the server free of
/// pidfile lifecycle handling (staleness, crash leftovers, one-per-cache-dir
/// assumptions). It also does the right thing inside a container, where the
/// server is PID 1 and shares the namespace with `docker exec`.
///
/// On a system without `/proc` the directory simply fails to open and the count
/// is zero, which the caller reports as "applies on next start".
fn notifyServers(io: std.Io) usize {
    if (!@import("../Signals.zig").supported) return 0;
    if (builtin.os.tag != .linux) return 0;

    var dir = std.Io.Dir.cwd().openDir(io, "/proc", .{ .iterate = true }) catch return 0;
    defer dir.close(io);

    // Compare against our own name rather than a hardcoded "wttr" so that running
    // under a different filename (a canary build, a symlinked unit) still works.
    // `comm` is compared on both sides so the kernel's 15-character truncation
    // applies identically and cannot cause a spurious mismatch.
    var self_name_buf: [64]u8 = undefined;
    const self_name = readComm(io, &dir, "self", &self_name_buf) orelse return 0;

    const self_pid = std.os.linux.getpid();
    var signaled: usize = 0;

    var it = dir.iterate();
    while (it.next(io) catch null) |entry| {
        if (entry.kind != .directory) continue;
        const pid = std.fmt.parseInt(std.posix.pid_t, entry.name, 10) catch continue;
        if (pid == self_pid) continue;

        var name_buf: [64]u8 = undefined;
        const name = readComm(io, &dir, entry.name, &name_buf) orelse continue;
        if (!std.mem.eql(u8, name, self_name)) continue;

        // Only the server reloads on SIGHUP; another command has nothing to
        // reload. Skipping them keeps this from depending on those processes
        // ignoring the signal.
        if (!isServing(io, &dir, entry.name)) continue;

        std.posix.kill(pid, .HUP) catch |err| {
            log.debug("could not signal pid {d}: {t}", .{ pid, err });
            continue;
        };
        signaled += 1;
    }

    return signaled;
}

/// Reads `/proc/<pid>/comm` into `buf`, returning the trimmed name.
fn readComm(io: std.Io, proc_dir: *std.Io.Dir, pid_name: []const u8, buf: []u8) ?[]const u8 {
    var path_buf: [64]u8 = undefined;
    const path = std.fmt.bufPrint(&path_buf, "{s}/comm", .{pid_name}) catch return null;

    const file = proc_dir.openFile(io, path, .{}) catch return null;
    defer file.close(io);

    // Read streaming rather than via `readFileAlloc`: procfs reports `st_size`
    // as 0, so anything that sizes the read from the file's length reads nothing
    // and every process looks nameless.
    const n = file.readStreaming(io, &.{buf}) catch return null;
    if (n == 0) return null;

    return std.mem.trim(u8, buf[0..n], " \n");
}

/// Whether a process is running the server rather than a command.
///
/// The server takes no arguments; every command has a subcommand as its first.
fn isServing(io: std.Io, proc_dir: *std.Io.Dir, pid_name: []const u8) bool {
    var path_buf: [64]u8 = undefined;
    const path = std.fmt.bufPrint(&path_buf, "{s}/cmdline", .{pid_name}) catch return false;

    const file = proc_dir.openFile(io, path, .{}) catch return false;
    defer file.close(io);

    var buf: [4096]u8 = undefined;
    const n = file.readStreaming(io, &.{&buf}) catch return false;

    return cmdlineIsBare(buf[0..n]);
}

/// Whether a `/proc/<pid>/cmdline` payload holds exactly one argument.
///
/// Split from `isServing` so the parsing is testable without a live process.
/// `cmdline` is NUL-separated with a trailing NUL after the final argument, so a
/// bare invocation is one entry and anything with a subcommand has two or more.
fn cmdlineIsBare(raw: []const u8) bool {
    if (raw.len == 0) return false;
    const args = std.mem.trimEnd(u8, raw, "\x00");
    if (args.len == 0) return false;
    return std.mem.indexOfScalar(u8, args, 0) == null;
}

test "joinLocation joins argv with plus and converts spaces" {
    const allocator = std.testing.allocator;

    const a = try joinLocation(allocator, &.{ "San", "Francisco" });
    defer allocator.free(a);
    try std.testing.expectEqualStrings("San+Francisco", a);

    // A quoted argument arrives as one part containing a space.
    const b = try joinLocation(allocator, &.{"San Francisco"});
    defer allocator.free(b);
    try std.testing.expectEqualStrings("San+Francisco", b);

    const c = try joinLocation(allocator, &.{ "Portland,", "Oregon" });
    defer allocator.free(c);
    try std.testing.expectEqualStrings("Portland,+Oregon", c);
}

test "joinLocation on a single word is unchanged" {
    const allocator = std.testing.allocator;
    const s = try joinLocation(allocator, &.{"London"});
    defer allocator.free(s);
    try std.testing.expectEqualStrings("London", s);
}

test "cmdlineIsBare: a bare invocation is the server" {
    try std.testing.expect(cmdlineIsBare("/wttr\x00"));
    // Some kernels omit the trailing NUL on short reads.
    try std.testing.expect(cmdlineIsBare("/wttr"));
    try std.testing.expect(cmdlineIsBare("/usr/local/bin/wttr\x00"));
}

test "cmdlineIsBare: anything with a subcommand is not the server" {
    try std.testing.expect(!cmdlineIsBare("/wttr\x00pin\x00"));
    try std.testing.expect(!cmdlineIsBare("/wttr\x00pins\x00"));
    try std.testing.expect(!cmdlineIsBare("/wttr\x00pin\x0012.0.0.0/8\x00London\x00"));
    // Trailing NULs must not hide a second argument.
    try std.testing.expect(!cmdlineIsBare("/wttr\x00help\x00\x00\x00"));
}

test "cmdlineIsBare: empty or all-NUL input is not a server" {
    try std.testing.expect(!cmdlineIsBare(""));
    try std.testing.expect(!cmdlineIsBare("\x00"));
    try std.testing.expect(!cmdlineIsBare("\x00\x00\x00"));
}