diff --git a/build.zig b/build.zig index 4ca331c..1e7d5d2 100644 --- a/build.zig +++ b/build.zig @@ -23,6 +23,11 @@ pub fn build(b: *std.Build) void { .optimize = optimize, }); + const srf = b.dependency("srf", .{ + .target = target, + .optimize = optimize, + }); + const openflights = b.dependency("openflights", .{}); const maxminddb_upstream = b.dependency("maxminddb", .{}); @@ -126,6 +131,7 @@ pub fn build(b: *std.Build) void { }); root_module.addImport("httpz", httpz.module("httpz")); root_module.addImport("zeit", zeit.module("zeit")); + root_module.addImport("srf", srf.module("srf")); root_module.addAnonymousImport("airports.dat", .{ .root_source_file = openflights.path("data/airports.dat"), }); diff --git a/build.zig.zon b/build.zig.zon index 02c1d1c..515813e 100644 --- a/build.zig.zon +++ b/build.zig.zon @@ -20,6 +20,10 @@ }, .phoon = .{ .path = "libs/phoon_14Aug2014" }, .sunriset = .{ .path = "libs/sunriset" }, + .srf = .{ + .url = "git+https://git.lerch.org/lobo/srf#4a3e5f00f15b0e0ba79d06ffe69dbcfa052baa5b", + .hash = "srf-0.0.0-qZj572nkAQAAz3zEg6fdD8A7PJnQ9je3zCeAOJS5PoZj", + }, }, .fingerprint = 0x710c2b57e81aa678, .minimum_zig_version = "0.16.0", diff --git a/src/Config.zig b/src/Config.zig index 78bcc94..dfa4819 100644 --- a/src/Config.zig +++ b/src/Config.zig @@ -48,6 +48,14 @@ ip2location_cache_file: []const u8, /// Cache file for ipwho.is lookups ipwhois_cache_file: []const u8, +/// Manual IP-range to location overrides, consulted before GeoLite2. +/// +/// Not a cache: GeoLite2 can answer confidently and wrongly (a corporate WAN +/// address with `accuracy_radius = 20` pointing 1400 miles away), and the +/// fallback provider only runs when GeoLite2 declines to answer. Overrides +/// therefore have to sit in front of the database. +pins_file: []const u8, + /// Loads configuration from the process environment. /// /// Zig 0.16 removed `std.process.getEnvMap`; the environment map is now @@ -109,6 +117,12 @@ pub fn load(allocator: std.mem.Allocator, env: *const std.process.Environ.Map) ! } break :blk try std.fmt.allocPrint(allocator, "{s}/ip2location.cache", .{env.get("WTTR_CACHE_DIR") orelse default_cache_dir}); }, + .pins_file = blk: { + if (env.get("WTTR_PINS_FILE")) |v| { + break :blk try allocator.dupe(u8, v); + } + break :blk try std.fmt.allocPrint(allocator, "{s}/pins.srf", .{env.get("WTTR_CACHE_DIR") orelse default_cache_dir}); + }, .ipwhois_cache_file = blk: { if (env.get("IPWHOIS_CACHE_FILE")) |v| { break :blk try allocator.dupe(u8, v); @@ -145,6 +159,7 @@ pub fn deinit(self: Config, allocator: std.mem.Allocator) void { if (self.ip2location_api_key) |k| allocator.free(k); allocator.free(self.ip2location_cache_file); allocator.free(self.ipwhois_cache_file); + allocator.free(self.pins_file); } test "config loads defaults" { diff --git a/src/Signals.zig b/src/Signals.zig new file mode 100644 index 0000000..0b1af81 --- /dev/null +++ b/src/Signals.zig @@ -0,0 +1,106 @@ +const std = @import("std"); +const builtin = @import("builtin"); + +const log = std.log.scoped(.signals); + +/// SIGHUP-driven reload of on-disk state. +/// +/// The handler does nothing but an atomic store. That is deliberate: a signal +/// handler may run on any thread that has not blocked the signal (including an +/// httpz worker mid-request), and almost nothing is safe to call from one. An +/// atomic store is, so the handler records the request and a watcher task does +/// the actual work on a normal thread. +/// +/// A `signalfd` or self-pipe would remove the polling, but both add file +/// descriptor handling for no practical gain here: reloads are operator-driven +/// and rare, so noticing within a second is indistinguishable from instant. +/// +/// `SA_RESTART` matters more than it looks. Without it, delivering SIGHUP would +/// interrupt whatever blocking syscall the receiving thread was in (an `accept` +/// or `read` in the HTTP server) with `EINTR`, turning an operator reload into +/// spurious request failures. +var reload_requested: std.atomic.Value(bool) = .init(false); + +fn onHup(_: std.posix.SIG) callconv(.c) void { + reload_requested.store(true, .release); +} + +/// Installs the SIGHUP handler. Safe to call when unsupported: it does nothing. +/// +/// Installing a real handler (rather than blocking the signal and draining it +/// elsewhere) also settles a container question. The server runs as PID 1 under +/// Docker, and the kernel discards signals sent to PID 1 when they would take +/// their default action and no handler is installed. With a handler present, +/// delivery is unambiguous. +pub fn install() void { + if (!supported) { + log.debug("SIGHUP reload is not supported on this target", .{}); + return; + } + + var act: std.posix.Sigaction = .{ + .handler = .{ .handler = onHup }, + .mask = std.posix.sigemptyset(), + .flags = std.posix.SA.RESTART, + }; + std.posix.sigaction(.HUP, &act, null); +} + +/// Ignores SIGHUP, for processes that are not the server. +/// +/// SIGHUP's default disposition is to terminate. Without this, a short-lived +/// command that happens to be running while another invocation looks for servers +/// to signal would simply be killed: process discovery matches on the program, +/// and a command in flight (geocoding a location, say) looks much like a server. +/// +/// Making non-server modes ignore the signal removes that hazard at the source, +/// so correctness no longer depends on discovery being perfectly precise. +pub fn ignore() void { + if (!supported) return; + + var act: std.posix.Sigaction = .{ + .handler = .{ .handler = std.posix.SIG.IGN }, + .mask = std.posix.sigemptyset(), + .flags = 0, + }; + std.posix.sigaction(.HUP, &act, null); +} + +pub const supported = switch (builtin.os.tag) { + .windows, .wasi => false, + else => true, +}; + +/// Consumes a pending reload request, returning whether one was set. +pub fn takeReloadRequest() bool { + return reload_requested.swap(false, .acq_rel); +} + +/// Poll interval for the watcher. Short enough that an operator running +/// `wttr pin` sees the effect immediately in human terms, long enough to be +/// free. +pub const poll_interval_ms = 1000; + +test "takeReloadRequest consumes the flag exactly once" { + // Reset in case another test in this process set it. + _ = takeReloadRequest(); + + try std.testing.expect(!takeReloadRequest()); + + onHup(.HUP); + try std.testing.expect(takeReloadRequest()); + try std.testing.expect(!takeReloadRequest()); +} + +test "repeated signals before a poll collapse into one reload" { + _ = takeReloadRequest(); + + onHup(.HUP); + onHup(.HUP); + onHup(.HUP); + + // Coalescing is intended: three rapid reload requests need one reload, not + // three sequential re-reads of the same files. + try std.testing.expect(takeReloadRequest()); + try std.testing.expect(!takeReloadRequest()); +} diff --git a/src/cli/pin.zig b/src/cli/pin.zig new file mode 100644 index 0000000..fd0198d --- /dev/null +++ b/src/cli/pin.zig @@ -0,0 +1,354 @@ +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 Override the location for an IP or range + \\ unpin 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 \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 \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 /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//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//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")); +} diff --git a/src/location/GeoIp.zig b/src/location/GeoIp.zig index 1fbd889..6016bd9 100644 --- a/src/location/GeoIp.zig +++ b/src/location/GeoIp.zig @@ -3,6 +3,7 @@ const Ip2location = @import("Ip2location.zig"); const IpWhoIs = @import("IpWhoIs.zig"); const Location = @import("resolver.zig").Location; const Config = @import("../Config.zig"); +const Pins = @import("Pins.zig"); const c = @cImport({ @cInclude("maxminddb.h"); @@ -60,10 +61,17 @@ allocator: std.mem.Allocator, io: std.Io, /// Retained so `reload` can reopen the database after it is replaced on disk. db_path: []const u8, -/// Guards `mmdb`. Lookups hold this shared; `reload` holds it exclusively to -/// swap the handle. Request threads and the background refresher both touch -/// `mmdb`, so the swap cannot be unsynchronized. +/// Guards `mmdb` and `pins`. Lookups hold this shared; the reload paths hold it +/// exclusively to swap state. Request threads, the background refresher, and the +/// SIGHUP handler all touch this state, so the swaps cannot be unsynchronized. lock: std.Io.RwLock, +/// Manual IP-range overrides, consulted before the database. Separate from the +/// fallback cache on purpose: the fallback only runs when GeoLite2 declines to +/// answer, and the addresses that need overriding are exactly the ones GeoLite2 +/// answers confidently and wrongly. +pins: Pins, +/// Path pins are loaded from, retained so SIGHUP can re-read them. +pins_path: []const u8, pub fn init(allocator: std.mem.Allocator, io: std.Io, db_path: []const u8, config: Config) !GeoIP { const mmdb = try openDatabase(allocator, db_path); @@ -75,6 +83,19 @@ pub fn init(allocator: std.mem.Allocator, io: std.Io, db_path: []const u8, confi const db_path_copy = try allocator.dupe(u8, db_path); errdefer allocator.free(db_path_copy); + const pins_path_copy = try allocator.dupe(u8, config.pins_file); + errdefer allocator.free(pins_path_copy); + + // A malformed or unreadable pins file must not stop the server: overrides are + // an enhancement, and losing them is far better than refusing to serve. + var pins = Pins.load(allocator, io, config.pins_file) catch |err| blk: { + std.log.warn("could not load pins from {s} ({t}); continuing without overrides", .{ config.pins_file, err }); + break :blk Pins.init(allocator); + }; + errdefer pins.deinit(); + if (pins.entries.items.len > 0) + std.log.info("loaded {d} IP pin(s) from {s}", .{ pins.entries.items.len, config.pins_file }); + const fallback_client: FallbackClient = switch (config.geoip_fallback) { .ip2location => blk: { const client = try allocator.create(Ip2location); @@ -102,6 +123,8 @@ pub fn init(allocator: std.mem.Allocator, io: std.Io, db_path: []const u8, confi .io = io, .db_path = db_path_copy, .lock = .init, + .pins = pins, + .pins_path = pins_path_copy, }; } @@ -148,9 +171,27 @@ pub fn deinit(self: *GeoIP) void { c.MMDB_close(self.mmdb); self.allocator.destroy(self.mmdb); self.allocator.free(self.db_path); + self.allocator.free(self.pins_path); + self.pins.deinit(); self.fallback_client.deinit(self.allocator); } +/// Re-reads the pins file, replacing the in-memory overrides. +/// +/// Parsing happens before the exclusive lock is taken so lookups only block for +/// the swap. A parse failure leaves the existing pins installed. +pub fn reloadPins(self: *GeoIP) !void { + const fresh = try Pins.load(self.allocator, self.io, self.pins_path); + + self.lock.lockUncancelable(self.io); + var old = self.pins; + self.pins = fresh; + self.lock.unlock(self.io); + + old.deinit(); + log.info("reloaded {d} IP pin(s) from {s}", .{ self.pins.entries.items.len, self.pins_path }); +} + pub fn lookup(self: *GeoIP, ip: []const u8) ?Location { // Try MaxMind first. The shared lock has to cover `extractCoordinates` as // well as the lookup itself: `MMDB_lookup_result_s.entry` holds a pointer @@ -159,6 +200,14 @@ pub fn lookup(self: *GeoIP, ip: []const u8) ?Location { self.lock.lockSharedUncancelable(self.io); defer self.lock.unlockShared(self.io); + // Manual overrides win outright. They exist precisely for addresses the + // database resolves confidently and incorrectly, so consulting the + // database first would defeat the purpose. + if (self.pins.lookup(self.allocator, ip)) |pinned| { + log.debug("pin matched for ip {s} -> {s}", .{ ip, pinned.name }); + break :blk pinned; + } + const result = lookupInternal(self.mmdb, ip) catch break :blk null; log.debug("lookup geoip db for ip {s}. Found: {}", .{ ip, result.found_entry }); @@ -410,3 +459,107 @@ test "reload swaps the database while lookups keep working" { // catches a reload that left the lock in a bad state. try std.testing.expect(geoip.isUSIp(test_ip)); } + +test "a pin overrides a confident but wrong database answer" { + const allocator = std.testing.allocator; + const io = std.testing.io; + const build_options = @import("build_options"); + + var config = try Config.loadForTest(allocator); + defer config.deinit(allocator); + + if (build_options.download_geoip) { + const GeoLite2 = @import("GeoLite2.zig"); + try GeoLite2.ensureDatabase(allocator, io, config.geolite_path); + } + + // This is the bug this feature exists for: GeoLite2 places this AT&T address + // in Fort Worth, Texas with accuracy_radius = 20. The radius is low enough + // that the low-confidence fallback never triggers, so nothing but an explicit + // override can correct it. + const wrong_ip = "12.94.132.170"; + + 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 pins_path = try std.fmt.allocPrint(allocator, "{s}/pins", .{path_buf[0..dir_len]}); + defer allocator.free(pins_path); + + // Start with no pins and confirm the database's wrong answer is what we get. + allocator.free(config.pins_file); + config.pins_file = try allocator.dupe(u8, pins_path); + + { + var geoip = GeoIP.init(allocator, io, config.geolite_path, config) catch + return error.SkipZigTest; + defer geoip.deinit(); + + const unpinned = geoip.lookup(wrong_ip) orelse return error.SkipZigTest; + defer unpinned.deinit(); + // Guard against a future database that fixes this on its own; if that + // happens this test needs a different address rather than a silent pass. + if (std.mem.indexOf(u8, unpinned.name, "Texas") == null) return error.SkipZigTest; + } + + // Write a pin covering the whole /24 the address sits in. + { + var pins: Pins = .init(allocator); + defer pins.deinit(); + try pins.put( + try Pins.parseCidr("12.94.132.0/24"), + "San Francisco, California, United States", + .{ .latitude = 37.7749, .longitude = -122.4194 }, + ); + try pins.save(io, pins_path); + } + + var geoip = GeoIP.init(allocator, io, config.geolite_path, config) catch + return error.SkipZigTest; + defer geoip.deinit(); + + const pinned = geoip.lookup(wrong_ip) orelse return error.TestUnexpectedResult; + defer pinned.deinit(); + + try std.testing.expectEqualStrings("San Francisco, California, United States", pinned.name); + try std.testing.expectEqual(@as(f64, 37.7749), pinned.coords.latitude); + try std.testing.expectEqual(@as(f64, -122.4194), pinned.coords.longitude); +} + +test "reloadPins picks up a pin written after startup" { + const allocator = std.testing.allocator; + const io = std.testing.io; + + var config = try Config.loadForTest(allocator); + defer config.deinit(allocator); + + 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 pins_path = try std.fmt.allocPrint(allocator, "{s}/pins", .{path_buf[0..dir_len]}); + defer allocator.free(pins_path); + + allocator.free(config.pins_file); + config.pins_file = try allocator.dupe(u8, pins_path); + + var geoip = GeoIP.init(allocator, io, config.geolite_path, config) catch + return error.SkipZigTest; + defer geoip.deinit(); + + try std.testing.expectEqual(@as(usize, 0), geoip.pins.entries.items.len); + + { + var pins: Pins = .init(allocator); + defer pins.deinit(); + try pins.put(try Pins.parseCidr("203.0.113.0/24"), "Somewhere", .{ .latitude = 1, .longitude = 2 }); + try pins.save(io, pins_path); + } + + try geoip.reloadPins(); + try std.testing.expectEqual(@as(usize, 1), geoip.pins.entries.items.len); + + const hit = geoip.lookup("203.0.113.7") orelse return error.TestUnexpectedResult; + defer hit.deinit(); + try std.testing.expectEqualStrings("Somewhere", hit.name); +} diff --git a/src/location/Ip2location.zig b/src/location/Ip2location.zig index 2c312b1..639e8a9 100644 --- a/src/location/Ip2location.zig +++ b/src/location/Ip2location.zig @@ -49,8 +49,20 @@ pub const PackedIp = struct { /// 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 addr = std.Io.net.IpAddress.parse(ip_str, 0) catch return null; + 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 }, diff --git a/src/location/Pins.zig b/src/location/Pins.zig new file mode 100644 index 0000000..bb13de3 --- /dev/null +++ b/src/location/Pins.zig @@ -0,0 +1,565 @@ +const std = @import("std"); +const Coordinates = @import("../Coordinates.zig"); +const Location = @import("resolver.zig").Location; +const packIp = @import("Ip2location.zig").packIp; +const srf = @import("srf"); + +const Pins = @This(); +const log = std.log.scoped(.pins); + +/// Manual IP-range to location overrides, consulted *before* GeoLite2. +/// +/// This exists because GeoLite2 can be confidently wrong. A corporate WAN +/// address may carry `accuracy_radius = 20` (a strong confidence signal) while +/// pointing at a city 1400 miles from the client, which means no +/// confidence-based heuristic and no online fallback will ever correct it: the +/// fallback only runs when GeoLite2 declines to answer. +/// +/// Pins are therefore a higher-priority store rather than another cache. The +/// existing ipwho.is cache cannot serve this purpose, as it is only reached when +/// GeoLite2 has already given up. +allocator: std.mem.Allocator, +entries: std.ArrayList(Entry), + +pub const Entry = struct { + /// 4 or 6. A pin only matches addresses of the same family. + family: u8, + /// Network address with all bits below `prefix_len` cleared. + network: u128, + /// Prefix length in bits (0-32 for IPv4, 0-128 for IPv6). + prefix_len: u8, + name: []const u8, + coords: Coordinates, +}; + +pub fn init(allocator: std.mem.Allocator) Pins { + return .{ .allocator = allocator, .entries = .empty }; +} + +pub fn deinit(self: *Pins) void { + for (self.entries.items) |e| self.allocator.free(e.name); + self.entries.deinit(self.allocator); +} + +/// Total bits in an address of `family`. +fn familyBits(family: u8) u8 { + return if (family == 4) 32 else 128; +} + +/// Masks `addr` down to its first `prefix_len` bits. +/// +/// Shifting by the full width is undefined in Zig, so a prefix that covers the +/// whole address is returned unchanged. +fn maskToPrefix(addr: u128, family: u8, prefix_len: u8) u128 { + const bits = familyBits(family); + if (prefix_len >= bits) return addr; + const host_bits: u7 = @intCast(bits - prefix_len); + const mask = ~@as(u128, 0) << host_bits; + // For IPv4 the value only occupies the low 32 bits, so the mask needs + // narrowing to avoid clearing bits that were never part of the address. + const family_mask: u128 = if (bits == 128) ~@as(u128, 0) else (@as(u128, 1) << 32) - 1; + return addr & mask & family_mask; +} + +pub const ParseError = error{ + InvalidAddress, + InvalidPrefixLength, +}; + +pub const Cidr = struct { + family: u8, + network: u128, + prefix_len: u8, +}; + +/// Parses `1.2.3.0/24`, `2001:db8::/32`, or a bare address. +/// +/// A bare address is treated as a host route (`/32` or `/128`), so the common +/// case of pinning one address needs no prefix. +pub fn parseCidr(text: []const u8) ParseError!Cidr { + const slash = std.mem.indexOfScalar(u8, text, '/'); + const addr_text = if (slash) |i| text[0..i] else text; + + const packed_ip = packIp(addr_text) orelse return error.InvalidAddress; + const bits = familyBits(packed_ip.family); + + const prefix_len: u8 = if (slash) |i| blk: { + const n = std.fmt.parseInt(u8, text[i + 1 ..], 10) catch return error.InvalidPrefixLength; + if (n > bits) return error.InvalidPrefixLength; + break :blk n; + } else bits; + + return .{ + .family = packed_ip.family, + .network = maskToPrefix(packed_ip.key, packed_ip.family, prefix_len), + .prefix_len = prefix_len, + }; +} + +/// Formats a `Cidr` back into `address/prefix` form. +pub fn formatCidr(cidr: Cidr, buf: []u8) ![]const u8 { + if (cidr.family == 4) { + const v: u32 = @intCast(cidr.network); + return std.fmt.bufPrint(buf, "{d}.{d}.{d}.{d}/{d}", .{ + @as(u8, @truncate(v >> 24)), + @as(u8, @truncate(v >> 16)), + @as(u8, @truncate(v >> 8)), + @as(u8, @truncate(v)), + cidr.prefix_len, + }); + } + var bytes: [16]u8 = undefined; + std.mem.writeInt(u128, &bytes, cidr.network, .big); + // `Ip6Address.format` renders "[addr]:port", which would not parse back. + // `Unresolved` is the bare-address formatter. + const bare: std.Io.net.Ip6Address.Unresolved = .{ .bytes = bytes, .interface_name = null }; + return std.fmt.bufPrint(buf, "{f}/{d}", .{ &bare, cidr.prefix_len }); +} + +/// Longest-prefix match for `ip_str`, or null when no pin applies. +/// +/// Longest-prefix ordering means an exact host pin naturally wins over a range +/// covering it, without needing a separate precedence rule. +pub fn lookup(self: *const Pins, allocator: std.mem.Allocator, ip_str: []const u8) ?Location { + const packed_ip = packIp(ip_str) orelse return null; + + var best: ?*const Entry = null; + for (self.entries.items) |*e| { + if (e.family != packed_ip.family) continue; + if (maskToPrefix(packed_ip.key, e.family, e.prefix_len) != e.network) continue; + if (best == null or e.prefix_len > best.?.prefix_len) best = e; + } + + const entry = best orelse return null; + return .{ + .allocator = allocator, + .name = allocator.dupe(u8, entry.name) catch return null, + .coords = entry.coords, + }; +} + +/// Inserts a pin, replacing any existing pin for the same network. +pub fn put(self: *Pins, cidr: Cidr, name: []const u8, coords: Coordinates) !void { + const name_copy = try self.allocator.dupe(u8, name); + errdefer self.allocator.free(name_copy); + + for (self.entries.items) |*e| { + if (e.family == cidr.family and e.prefix_len == cidr.prefix_len and e.network == cidr.network) { + self.allocator.free(e.name); + e.name = name_copy; + e.coords = coords; + return; + } + } + + try self.entries.append(self.allocator, .{ + .family = cidr.family, + .network = cidr.network, + .prefix_len = cidr.prefix_len, + .name = name_copy, + .coords = coords, + }); +} + +/// Removes the pin for `cidr`, returning whether one was present. +pub fn remove(self: *Pins, cidr: Cidr) bool { + for (self.entries.items, 0..) |e, i| { + if (e.family == cidr.family and e.prefix_len == cidr.prefix_len and e.network == cidr.network) { + self.allocator.free(e.name); + _ = self.entries.orderedRemove(i); + return true; + } + } + return false; +} + +/// On-disk shape of one pin. +/// +/// Stored in SRF rather than a hand-rolled line format. This file is new, so +/// there is no legacy data to preserve, and inventing a bespoke format here +/// would only create something to migrate later. SRF also handles the awkward +/// part for free: place names routinely contain commas ("San Francisco, +/// California, United States"), which a comma-delimited format has to special +/// case. +const Record = struct { + cidr: []const u8, + lat: f64, + lon: f64, + name: []const u8, +}; + +/// Loads pins from `path`. A missing file yields an empty set: having no +/// overrides is the normal state, not an error. +pub fn load(allocator: std.mem.Allocator, io: std.Io, path: []const u8) !Pins { + var pins: Pins = .init(allocator); + errdefer pins.deinit(); + + const content = std.Io.Dir.cwd().readFileAlloc(io, path, allocator, .limited(1024 * 1024)) catch |err| switch (err) { + error.FileNotFound => return pins, + else => return err, + }; + defer allocator.free(content); + if (std.mem.trim(u8, content, " \r\n\t").len == 0) return pins; + + var reader = std.Io.Reader.fixed(content); + // Lenient number parsing because this file is meant to be operator-editable: + // a hand-written `lat::37.7` (untyped) should work as well as `lat:num:37.7`. + var records = srf.iterator(&reader, allocator, .{ .strict_number_parsing = false }) catch |err| { + log.warn("could not parse pins file {s} ({t}); continuing with no overrides", .{ path, err }); + return pins; + }; + defer records.deinit(); + + var index: usize = 0; + while (records.next() catch |err| { + log.warn("stopped reading {s} after {d} pin(s): {t}", .{ path, index, err }); + return pins; + }) |fields| { + index += 1; + const record = fields.to(Record, .{ .strings_to_numbers = true }) catch |err| { + log.warn("skipping malformed pin #{d} in {s}: {t}", .{ index, path, err }); + continue; + }; + + const cidr = parseCidr(record.cidr) catch |err| { + log.warn("skipping pin #{d} in {s}: invalid range {s} ({t})", .{ index, path, record.cidr, err }); + continue; + }; + + // `put` copies the name, so it does not matter that the parsed strings + // belong to the iterator's arena. + try pins.put(cidr, record.name, .{ .latitude = record.lat, .longitude = record.lon }); + } + + return pins; +} + +/// Writes all pins to `path`, replacing it atomically so a crash mid-write +/// cannot leave a half-written override file that fails to parse on next boot. +pub fn save(self: *const Pins, io: std.Io, path: []const u8) !void { + const tmp_path = try std.fmt.allocPrint(self.allocator, "{s}.new", .{path}); + defer self.allocator.free(tmp_path); + + if (std.fs.path.dirname(path)) |dir| { + try std.Io.Dir.cwd().createDirPath(io, dir); + } + + // Build the records up front: the formatter takes a slice, and the CIDR text + // has to outlive it. + var records = try self.allocator.alloc(Record, self.entries.items.len); + defer self.allocator.free(records); + var cidr_texts = try self.allocator.alloc([]u8, self.entries.items.len); + var made: usize = 0; + defer { + for (cidr_texts[0..made]) |t| self.allocator.free(t); + self.allocator.free(cidr_texts); + } + + for (self.entries.items, 0..) |e, i| { + var buf: [64]u8 = undefined; + const text = try formatCidr( + .{ .family = e.family, .network = e.network, .prefix_len = e.prefix_len }, + &buf, + ); + cidr_texts[i] = try self.allocator.dupe(u8, text); + made += 1; + records[i] = .{ + .cidr = cidr_texts[i], + .lat = e.coords.latitude, + .lon = e.coords.longitude, + .name = e.name, + }; + } + + { + const file = try std.Io.Dir.cwd().createFile(io, tmp_path, .{}); + errdefer std.Io.Dir.cwd().deleteFile(io, tmp_path) catch |err| + log.warn("could not remove partial pins file {s}: {t}", .{ tmp_path, err }); + defer file.close(io); + + var buf: [4096]u8 = undefined; + var fw = file.writer(io, &buf); + const w = &fw.interface; + + // Long format: one field per line, which is what an operator reading or + // editing this file by hand wants. + try w.print("{f}", .{srf.fmt(Record, records, .{ .long_format = true })}); + try w.flush(); + } + + try std.Io.Dir.cwd().rename(tmp_path, .cwd(), path, io); +} + +test "parseCidr: bare IPv4 becomes a host route" { + const c = try parseCidr("12.94.132.170"); + try std.testing.expectEqual(@as(u8, 4), c.family); + try std.testing.expectEqual(@as(u8, 32), c.prefix_len); +} + +test "parseCidr: masks host bits off the network" { + // .170 must be discarded by the /24. + const c = try parseCidr("12.94.132.170/24"); + const expected = try parseCidr("12.94.132.0/24"); + try std.testing.expectEqual(expected.network, c.network); + try std.testing.expectEqual(@as(u8, 24), c.prefix_len); +} + +test "parseCidr: IPv6 with prefix" { + const c = try parseCidr("2001:db8::/32"); + try std.testing.expectEqual(@as(u8, 6), c.family); + try std.testing.expectEqual(@as(u8, 32), c.prefix_len); +} + +test "parseCidr: bare IPv6 becomes a /128" { + const c = try parseCidr("2001:db8::1"); + try std.testing.expectEqual(@as(u8, 6), c.family); + try std.testing.expectEqual(@as(u8, 128), c.prefix_len); +} + +test "parseCidr: rejects an out-of-range prefix" { + try std.testing.expectError(error.InvalidPrefixLength, parseCidr("10.0.0.0/33")); + try std.testing.expectError(error.InvalidPrefixLength, parseCidr("2001:db8::/129")); +} + +test "parseCidr: rejects garbage" { + try std.testing.expectError(error.InvalidAddress, parseCidr("not-an-ip/24")); + try std.testing.expectError(error.InvalidPrefixLength, parseCidr("10.0.0.0/abc")); +} + +test "parseCidr: /0 matches everything in its family" { + const c = try parseCidr("0.0.0.0/0"); + try std.testing.expectEqual(@as(u128, 0), c.network); + try std.testing.expectEqual(@as(u8, 0), c.prefix_len); +} + +test "lookup: matches an address inside the range" { + const allocator = std.testing.allocator; + var pins: Pins = .init(allocator); + defer pins.deinit(); + + try pins.put(try parseCidr("12.94.132.0/24"), "San Francisco, California", .{ .latitude = 37.7749, .longitude = -122.4194 }); + + const hit = pins.lookup(allocator, "12.94.132.170") orelse return error.TestUnexpectedResult; + defer hit.deinit(); + try std.testing.expectEqualStrings("San Francisco, California", hit.name); + try std.testing.expectEqual(@as(f64, 37.7749), hit.coords.latitude); +} + +test "lookup: ignores an address outside the range" { + const allocator = std.testing.allocator; + var pins: Pins = .init(allocator); + defer pins.deinit(); + + try pins.put(try parseCidr("12.94.132.0/24"), "San Francisco", .{ .latitude = 37.7749, .longitude = -122.4194 }); + + try std.testing.expect(pins.lookup(allocator, "12.94.133.1") == null); +} + +test "lookup: longest prefix wins over a broader range" { + const allocator = std.testing.allocator; + var pins: Pins = .init(allocator); + defer pins.deinit(); + + // Deliberately inserted broad-first so the result cannot come from ordering. + try pins.put(try parseCidr("12.0.0.0/8"), "Broad", .{ .latitude = 1, .longitude = 1 }); + try pins.put(try parseCidr("12.94.132.0/24"), "Specific", .{ .latitude = 2, .longitude = 2 }); + + const hit = pins.lookup(allocator, "12.94.132.170") orelse return error.TestUnexpectedResult; + defer hit.deinit(); + try std.testing.expectEqualStrings("Specific", hit.name); +} + +test "lookup: an exact host pin beats a range containing it" { + const allocator = std.testing.allocator; + var pins: Pins = .init(allocator); + defer pins.deinit(); + + try pins.put(try parseCidr("12.94.132.170"), "Host", .{ .latitude = 2, .longitude = 2 }); + try pins.put(try parseCidr("12.94.132.0/24"), "Range", .{ .latitude = 1, .longitude = 1 }); + + const hit = pins.lookup(allocator, "12.94.132.170") orelse return error.TestUnexpectedResult; + defer hit.deinit(); + try std.testing.expectEqualStrings("Host", hit.name); +} + +test "lookup: families do not cross-match" { + const allocator = std.testing.allocator; + var pins: Pins = .init(allocator); + defer pins.deinit(); + + try pins.put(try parseCidr("0.0.0.0/0"), "All IPv4", .{ .latitude = 1, .longitude = 1 }); + + // An IPv4 /0 must not swallow IPv6 clients. + try std.testing.expect(pins.lookup(allocator, "2001:db8::1") == null); + const v4_hit = pins.lookup(allocator, "8.8.8.8") orelse return error.TestUnexpectedResult; + defer v4_hit.deinit(); +} + +test "put: replaces an existing pin for the same network" { + const allocator = std.testing.allocator; + var pins: Pins = .init(allocator); + defer pins.deinit(); + + const cidr = try parseCidr("12.94.132.0/24"); + try pins.put(cidr, "Old", .{ .latitude = 1, .longitude = 1 }); + try pins.put(cidr, "New", .{ .latitude = 2, .longitude = 2 }); + + try std.testing.expectEqual(@as(usize, 1), pins.entries.items.len); + const hit = pins.lookup(allocator, "12.94.132.5") orelse return error.TestUnexpectedResult; + defer hit.deinit(); + try std.testing.expectEqualStrings("New", hit.name); +} + +test "remove: reports whether a pin was present" { + const allocator = std.testing.allocator; + var pins: Pins = .init(allocator); + defer pins.deinit(); + + const cidr = try parseCidr("12.94.132.0/24"); + try pins.put(cidr, "SF", .{ .latitude = 1, .longitude = 1 }); + + try std.testing.expect(pins.remove(cidr)); + try std.testing.expect(!pins.remove(cidr)); + try std.testing.expectEqual(@as(usize, 0), pins.entries.items.len); +} + +test "formatCidr round-trips through parseCidr" { + var buf: [64]u8 = undefined; + for ([_][]const u8{ "12.94.132.0/24", "10.0.0.0/8", "0.0.0.0/0", "2001:db8::/32" }) |text| { + const parsed = try parseCidr(text); + const formatted = try formatCidr(parsed, &buf); + const reparsed = try parseCidr(formatted); + try std.testing.expectEqual(parsed.family, reparsed.family); + try std.testing.expectEqual(parsed.network, reparsed.network); + try std.testing.expectEqual(parsed.prefix_len, reparsed.prefix_len); + } +} + +test "save then load round-trips, including commas in the name" { + 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}/pins", .{path_buf[0..dir_len]}); + defer allocator.free(path); + + { + var pins: Pins = .init(allocator); + defer pins.deinit(); + try pins.put(try parseCidr("12.94.132.0/24"), "San Francisco, California, United States", .{ .latitude = 37.7749, .longitude = -122.4194 }); + try pins.put(try parseCidr("2001:db8::/32"), "Test, Place", .{ .latitude = -1.5, .longitude = 2.25 }); + try pins.save(io, path); + } + + var loaded = try load(allocator, io, path); + defer loaded.deinit(); + + try std.testing.expectEqual(@as(usize, 2), loaded.entries.items.len); + + const v4 = loaded.lookup(allocator, "12.94.132.170") orelse return error.TestUnexpectedResult; + defer v4.deinit(); + try std.testing.expectEqualStrings("San Francisco, California, United States", v4.name); + try std.testing.expectEqual(@as(f64, 37.7749), v4.coords.latitude); + try std.testing.expectEqual(@as(f64, -122.4194), v4.coords.longitude); + + const v6 = loaded.lookup(allocator, "2001:db8::1") orelse return error.TestUnexpectedResult; + defer v6.deinit(); + try std.testing.expectEqualStrings("Test, Place", v6.name); +} + +test "load: a missing file is an empty set, not an error" { + var pins = try load(std.testing.allocator, std.testing.io, "/nonexistent/pins"); + defer pins.deinit(); + try std.testing.expectEqual(@as(usize, 0), pins.entries.items.len); +} + +test "load: malformed records are skipped, valid ones kept" { + 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}/pins", .{path_buf[0..dir_len]}); + defer allocator.free(path); + + try std.Io.Dir.cwd().writeFile(io, .{ + .sub_path = path, + .data = + \\#!srfv1 + \\#!long + \\# an operator comment + \\ + \\cidr::999.1.1.1/24 + \\lat:num:1 + \\lon:num:2 + \\name::Bad Address + \\ + \\cidr::12.94.132.0/24 + \\lat:num:37.7749 + \\lon:num:-122.4194 + \\name::Good Entry, With Commas + \\ + , + }); + + var pins = try load(allocator, io, path); + defer pins.deinit(); + + // The bad address is dropped; a single unparseable record must not discard + // the operator's other overrides. + try std.testing.expectEqual(@as(usize, 1), pins.entries.items.len); + const hit = pins.lookup(allocator, "12.94.132.1") orelse return error.TestUnexpectedResult; + defer hit.deinit(); + try std.testing.expectEqualStrings("Good Entry, With Commas", hit.name); +} + +test "load: an empty or whitespace-only file yields no pins" { + 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}/pins", .{path_buf[0..dir_len]}); + defer allocator.free(path); + + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = "\n \n" }); + + var pins = try load(allocator, io, path); + defer pins.deinit(); + try std.testing.expectEqual(@as(usize, 0), pins.entries.items.len); +} + +test "save writes SRF that an operator can read" { + 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}/pins", .{path_buf[0..dir_len]}); + defer allocator.free(path); + + var pins: Pins = .init(allocator); + defer pins.deinit(); + try pins.put(try parseCidr("12.94.132.0/24"), "San Francisco, California", .{ .latitude = 37.7749, .longitude = -122.4194 }); + try pins.save(io, path); + + const content = try std.Io.Dir.cwd().readFileAlloc(io, path, allocator, .limited(64 * 1024)); + defer allocator.free(content); + + // Assert the shape, not just that it round-trips: this file is meant to be + // human-inspectable, and a silent switch to compact output would undermine + // that without failing any round-trip test. + try std.testing.expect(std.mem.indexOf(u8, content, "#!srfv1") != null); + try std.testing.expect(std.mem.indexOf(u8, content, "#!long") != null); + try std.testing.expect(std.mem.indexOf(u8, content, "cidr::12.94.132.0/24") != null); + try std.testing.expect(std.mem.indexOf(u8, content, "San Francisco, California") != null); +} diff --git a/src/location/Refresher.zig b/src/location/Refresher.zig index 3423c2c..b7233bc 100644 --- a/src/location/Refresher.zig +++ b/src/location/Refresher.zig @@ -103,3 +103,34 @@ test "ageOf spans days for the staleness comparison" { try std.testing.expectEqual(@as(u64, 31), @divFloor(age, std.time.s_per_day)); try std.testing.expect(age >= 30 * std.time.s_per_day); } + +/// Watches for operator-requested reloads (SIGHUP) and applies them. +/// +/// Separate from `run` because the two have unrelated cadences: staleness is +/// checked on the order of a day, while a reload request should be picked up +/// promptly. Folding them together would mean either polling the signal flag +/// once a day or waking the staleness check every second. +pub fn watchReloads(self: Refresher) void { + const Signals = @import("../Signals.zig"); + + while (true) { + std.Io.sleep(self.io, .fromMilliseconds(Signals.poll_interval_ms), .real) catch { + log.debug("reload watcher stopping", .{}); + return; + }; + + if (!Signals.takeReloadRequest()) continue; + + log.info("SIGHUP received, reloading on-disk state", .{}); + + // Pins are the common case (an operator just ran `wttr pin`), and a + // failure here leaves the previous overrides installed. + self.geoip.reloadPins() catch |err| + log.err("could not reload pins ({t}); keeping the previous overrides", .{err}); + + // Also reopen the database, so replacing the file by hand does not + // require waiting for the staleness check or restarting the server. + self.geoip.reload() catch |err| + log.err("could not reload the GeoLite2 database ({t}); keeping the previous one", .{err}); + } +} diff --git a/src/main.zig b/src/main.zig index e1a2155..14140cc 100644 --- a/src/main.zig +++ b/src/main.zig @@ -10,6 +10,8 @@ const Airports = @import("location/Airports.zig"); const Resolver = @import("location/resolver.zig").Resolver; const GeoLite2 = @import("location/GeoLite2.zig"); const Refresher = @import("location/Refresher.zig"); +const Signals = @import("Signals.zig"); +const pin_cmd = @import("cli/pin.zig"); const version = @import("build_options").version; /// Zig 0.16 entry point: the runtime supplies allocators, the `Io` @@ -18,9 +20,43 @@ pub fn main(init: std.process.Init) !u8 { const allocator = init.gpa; const io = init.io; + const args = try init.minimal.args.toSlice(allocator); + defer allocator.free(args); + const cfg = try Config.load(allocator, init.environ_map); defer cfg.deinit(allocator); + // No subcommand means "run the server", which keeps the deployed entrypoint + // (`ENTRYPOINT ["/wttr"]`) working unchanged. + if (args.len <= 1) return serve(allocator, io, init, cfg); + + // Everything below is a short-lived command rather than the server. SIGHUP + // defaults to terminating the process, and a concurrent invocation looking + // for servers to signal cannot tell a command in flight from a server, so + // ignore the signal instead of being killed by it. + Signals.ignore(); + + const cmd = args[1]; + if (std.mem.eql(u8, cmd, "pin")) return pin_cmd.runPin(allocator, io, cfg, args[2..]); + if (std.mem.eql(u8, cmd, "unpin")) return pin_cmd.runUnpin(allocator, io, cfg, args[2..]); + if (std.mem.eql(u8, cmd, "pins")) return pin_cmd.runList(allocator, io, cfg); + if (std.mem.eql(u8, cmd, "help") or std.mem.eql(u8, cmd, "--help") or std.mem.eql(u8, cmd, "-h")) + return pin_cmd.printUsage(io, args[0]); + + var stderr_buf: [512]u8 = undefined; + var stderr = std.Io.File.stderr().writer(io, &stderr_buf); + try stderr.interface.print("unknown command: {s}\n\n", .{cmd}); + try stderr.interface.flush(); + _ = try pin_cmd.printUsage(io, args[0]); + return 2; +} + +fn serve( + allocator: std.mem.Allocator, + io: std.Io, + init: std.process.Init, + cfg: Config, +) !u8 { std.log.info("wttr version {s} starting on {s}:{d}", .{ version, cfg.listen_host, cfg.listen_port }); std.log.info("Cache size: {d}", .{cfg.cache_size}); std.log.info("Cache dir: {s}", .{cfg.cache_dir}); @@ -86,6 +122,25 @@ pub fn main(init: std.process.Init) !u8 { // Only set up the server instance in debug mode if (@import("builtin").mode == .Debug) @import("http/handler.zig").server_instance = &server.httpz_server; + const refresher: Refresher = .{ + .allocator = allocator, + .io = io, + .geoip = &geoip, + .db_path = cfg.geolite_path, + .max_age_seconds = cfg.geolite_max_age_seconds, + .check_interval_seconds = cfg.geolite_check_interval_seconds, + }; + + // Watch for operator-driven reloads (SIGHUP, as sent by `wttr pin`). Started + // before the refresher so a reload is possible even if the refresher is + // disabled or could not start. + Signals.install(); + var reload_future = io.concurrent(Refresher.watchReloads, .{refresher}) catch |err| blk: { + std.log.warn("could not start the SIGHUP reload watcher ({t}); pins will apply on next start", .{err}); + break :blk null; + }; + defer if (reload_future) |*f| f.cancel(io); + // Refresh the GeoLite2 database in the background. `concurrent` rather than // `async` because this must get a real unit of concurrency: `async` is // allowed to run the function inline, which would block startup on a @@ -94,14 +149,6 @@ pub fn main(init: std.process.Init) !u8 { std.log.info("GeoLite2 background refresh disabled", .{}); break :blk null; } else blk: { - const refresher: Refresher = .{ - .allocator = allocator, - .io = io, - .geoip = &geoip, - .db_path = cfg.geolite_path, - .max_age_seconds = cfg.geolite_max_age_seconds, - .check_interval_seconds = cfg.geolite_check_interval_seconds, - }; std.log.info( "GeoLite2 refresh: checking every {d}h, refreshing when older than {d}d", .{ @@ -143,4 +190,7 @@ test { _ = @import("location/resolver.zig"); _ = @import("location/IpWhoIs.zig"); _ = @import("location/Refresher.zig"); + _ = @import("location/Pins.zig"); + _ = @import("Signals.zig"); + _ = @import("cli/pin.zig"); }