a73x

src/server/upgrade.zig

Ref:   Size: 19.2 KiB   History

//! The upgrade vocabulary: the version rule, and the manifest an exec-ing
//! daemon leaves for its replacement. It crosses inside ONE process, yet it is
//! encoded as if for a stranger — length-prefixed sections, unknown tags
//! skipped — because the reader IS one: a newer binary, or an older one on
//! rollback, and neither may be held to this build's struct layout.

const std = @import("std");
const proto = @import("term").protocol;

/// Skew rule: an upgrade candidate must outrank the running daemon.
/// SemanticVersion because our releases are `0.0.1-13`-shaped: the `-13`
/// is the prerelease field and numeric identifiers order numerically,
/// which is exactly the order the release sequence needs.
pub fn strictlyNewer(candidate: []const u8, mine: []const u8) error{BadVersion}!bool {
    const c = std.SemanticVersion.parse(candidate) catch return error.BadVersion;
    const m = std.SemanticVersion.parse(mine) catch return error.BadVersion;
    return c.order(m) == .gt;
}

test "strictlyNewer: the next prerelease number is newer (0.0.1-14 over 0.0.1-13)" {
    try std.testing.expect(try strictlyNewer("0.0.1-14", "0.0.1-13"));
}
test "strictlyNewer: prerelease numbers order numerically, not lexically (0.0.1-100 over 0.0.1-99)" {
    try std.testing.expect(try strictlyNewer("0.0.1-100", "0.0.1-99"));
}
test "strictlyNewer: the same version is not newer" {
    try std.testing.expect(!try strictlyNewer("0.0.1-13", "0.0.1-13"));
}
test "strictlyNewer: a downgrade is not newer" {
    try std.testing.expect(!try strictlyNewer("0.0.1-12", "0.0.1-13"));
}
test "strictlyNewer: a release outranks its own prereleases" {
    try std.testing.expect(try strictlyNewer("0.0.1", "0.0.1-13"));
}
test "strictlyNewer: garbage is an error, not a verdict" {
    try std.testing.expectError(error.BadVersion, strictlyNewer("not-a-version", "0.0.1-13"));
}

/// Bumped when a section's LAYOUT changes, independent of the release
/// version: the release number says what the binary can do, this says
/// what these bytes mean, and rollback needs the two decoupled.
pub const manifest_version: u16 = 1;
const magic = "MUXU";

pub const SectionTag = enum(u8) { daemon = 1, session = 2, _ };

pub const EnvPair = struct { key: []const u8, value: ?[]const u8 };

pub const QuicArm = enum(u8) { none = 0, borrowed = 1, owned = 2 };

/// Runtime QUIC state, not launch flags: lazyBindQuic can own an
/// ephemeral-port listener no flag names, so the manifest records what IS
/// bound, arm and all. The key crosses as bytes, never a path — no path
/// names the carrier and the file the key came from may have moved.
pub const QuicState = struct {
    arm: QuicArm = .none,
    fd: i32 = -1,
    addr: [addr_cap]u8 = @splat(0),
    addr_len: u32 = 0,
    idle_ms: u64 = 0,
    key: [key_len]u8 = @splat(0),

    const addr_cap = 128; // sockaddr_storage
    pub const key_len = 32;
};

/// Cumulative counters cross so an upgrade is not mistaken for a restart
/// by anything sampling `mux d stats`.
pub const Counters = struct {
    snapshots: u64 = 0,
    snapshot_bytes: u64 = 0,
    deltas: u64 = 0,
    delta_bytes: u64 = 0,
    /// What the same updates would have cost a snapshot-only daemon:
    /// measured, not estimated — dumpState length at each delta send.
    snapshot_equiv_bytes: u64 = 0,
    /// Every `.attach` this daemon ACCEPTED, cumulative and monotonic.
    ///
    /// A counter where `clients=` is a gauge, because the gauge cannot answer
    /// "did anyone attach since I last looked": a client that attaches and
    /// leaves between two samples is invisible to it, and so is one that closes
    /// as another opens. `assert_attach_delta` in e2e_lib.sh is what needs the
    /// counter — it holds a wall of N tiles to N attaches however far the focus
    /// moves, which is a negative no gauge can witness.
    ///
    /// Counted where an attach SUCCEEDS — a session resolved and the client
    /// seated — not where the frame arrives: a refusal attached nobody. Both
    /// arms count, since a socket client's first attach promotes an observer
    /// and every later one, a QUIC first attach included, arrives established.
    attaches: u64 = 0,
    agent_refused_no_offer: u64 = 0,
    agent_refused_full: u64 = 0,
};

pub const Daemon = struct {
    writer_version: []const u8,
    /// The rollback target: adoption failure execs this path back.
    writer_path: []const u8,
    sock_path: []const u8,
    listener_fd: i32,
    /// Both dirs are pid-named and held open by live shells; they cross as
    /// paths because the new process (same pid) can keep them but could
    /// never re-mint them.
    shellint_dir: ?[]const u8,
    agent_dir: ?[]const u8,
    shell: []const u8,
    shell_integration: bool,
    extra_env: []const EnvPair,
    quic: QuicState,
    counters: Counters,
};

const CmdRec = struct {
    phase: u8,
    marks_seen: bool,
    start_row: u32,
    end_row: u32,
    exit_code: ?u8,
};

pub const SessionRec = struct {
    name: []const u8,
    pty_fd: i32,
    child_pid: i32,
    cols: u16,
    rows: u16,
    /// Engine.dumpState bytes — viewport only; scrollback loss is the
    /// spec's named non-goal.
    vt: []const u8,
    /// dumpState has no title field; carried apart or every title goes
    /// blank until the shell next sets one.
    title: ?[]const u8,
    cmd: CmdRec,
    /// The return watermark `mux a await --since` answers from; losing it
    /// turns a satisfiable await into a timeout.
    last_return: ?proto.CmdState,
    agent_fd: i32,
    agent_path: ?[]const u8,
};

// ---- encoding ----
// magic ++ u16 manifest_version, then sections: u8 tag ++ u32 LE len ++ bytes.
// Strings are u32 LE len ++ bytes; optionals are u8 present ++ payload. An
// unknown tag is skipped by its length — the whole forward-compat story.

fn writeInt(w: anytype, comptime T: type, v: T) !void {
    var buf: [@sizeOf(T)]u8 = undefined;
    std.mem.writeInt(T, &buf, v, .little);
    try w.writeAll(&buf);
}

fn writeStr(w: anytype, s: []const u8) !void {
    try writeInt(w, u32, @intCast(s.len));
    try w.writeAll(s);
}

fn writeOptStr(w: anytype, s: ?[]const u8) !void {
    try w.writeAll(&.{@intFromBool(s != null)});
    if (s) |x| try writeStr(w, x);
}

fn writeDaemonBody(w: anytype, d: Daemon) !void {
    try writeStr(w, d.writer_version);
    try writeStr(w, d.writer_path);
    try writeStr(w, d.sock_path);
    try writeInt(w, i32, d.listener_fd);
    try writeOptStr(w, d.shellint_dir);
    try writeOptStr(w, d.agent_dir);
    try writeStr(w, d.shell);
    try w.writeAll(&.{@intFromBool(d.shell_integration)});
    try writeInt(w, u32, @intCast(d.extra_env.len));
    for (d.extra_env) |p| {
        try writeStr(w, p.key);
        try writeOptStr(w, p.value);
    }
    try w.writeAll(&.{@intFromEnum(d.quic.arm)});
    try writeInt(w, i32, d.quic.fd);
    try writeInt(w, u32, d.quic.addr_len);
    try w.writeAll(&d.quic.addr);
    try writeInt(w, u64, d.quic.idle_ms);
    try w.writeAll(&d.quic.key);
    inline for (@typeInfo(Counters).@"struct".fields) |f|
        try writeInt(w, u64, @field(d.counters, f.name));
}

fn writeSessionBody(w: anytype, s: SessionRec) !void {
    try writeStr(w, s.name);
    try writeInt(w, i32, s.pty_fd);
    try writeInt(w, i32, s.child_pid);
    try writeInt(w, u16, s.cols);
    try writeInt(w, u16, s.rows);
    try writeStr(w, s.vt);
    try writeOptStr(w, s.title);
    try w.writeAll(&.{s.cmd.phase});
    try w.writeAll(&.{@intFromBool(s.cmd.marks_seen)});
    try writeInt(w, u32, s.cmd.start_row);
    try writeInt(w, u32, s.cmd.end_row);
    try w.writeAll(&.{@intFromBool(s.cmd.exit_code != null)});
    try w.writeAll(&.{s.cmd.exit_code orelse 0});
    try w.writeAll(&.{@intFromBool(s.last_return != null)});
    if (s.last_return) |lr| try w.writeAll(&proto.encodeCmdState(lr));
    try writeInt(w, i32, s.agent_fd);
    try writeOptStr(w, s.agent_path);
}

fn writeSection(w: anytype, alloc: std.mem.Allocator, tag: SectionTag, body: []const u8) !void {
    _ = alloc;
    try w.writeAll(&.{@intFromEnum(tag)});
    try writeInt(w, u32, @intCast(body.len));
    try w.writeAll(body);
}

pub fn writeManifest(w: anytype, alloc: std.mem.Allocator, d: Daemon, sessions: []const SessionRec) !void {
    try w.writeAll(magic);
    try writeInt(w, u16, manifest_version);

    var body: std.ArrayList(u8) = .empty;
    defer body.deinit(alloc);

    try writeDaemonBody(body.writer(alloc), d);
    try writeSection(w, alloc, .daemon, body.items);

    for (sessions) |s| {
        body.clearRetainingCapacity();
        try writeSessionBody(body.writer(alloc), s);
        try writeSection(w, alloc, .session, body.items);
    }
}

// ---- parsing ----

const Cursor = struct {
    bytes: []const u8,
    i: usize = 0,

    fn take(self: *Cursor, n: usize) error{BadManifest}![]const u8 {
        if (self.bytes.len - self.i < n) return error.BadManifest;
        defer self.i += n;
        return self.bytes[self.i..][0..n];
    }
    fn int(self: *Cursor, comptime T: type) error{BadManifest}!T {
        return std.mem.readInt(T, (try self.take(@sizeOf(T)))[0..@sizeOf(T)], .little);
    }
    fn str(self: *Cursor, alloc: std.mem.Allocator) ![]const u8 {
        const n = try self.int(u32);
        return alloc.dupe(u8, try self.take(n));
    }
    fn optStr(self: *Cursor, alloc: std.mem.Allocator) !?[]const u8 {
        if ((try self.int(u8)) == 0) return null;
        return try self.str(alloc);
    }
    fn done(self: *const Cursor) bool {
        return self.i == self.bytes.len;
    }
};

pub const Parsed = struct {
    daemon: Daemon,
    sessions: []SessionRec,
    arena: std.heap.ArenaAllocator,

    pub fn deinit(self: *Parsed) void {
        self.arena.deinit();
    }
};

fn parseDaemon(c: *Cursor, alloc: std.mem.Allocator) !Daemon {
    var d: Daemon = undefined;
    d.writer_version = try c.str(alloc);
    d.writer_path = try c.str(alloc);
    d.sock_path = try c.str(alloc);
    d.listener_fd = try c.int(i32);
    d.shellint_dir = try c.optStr(alloc);
    d.agent_dir = try c.optStr(alloc);
    d.shell = try c.str(alloc);
    d.shell_integration = (try c.int(u8)) != 0;
    const env_n = try c.int(u32);
    const env = try alloc.alloc(EnvPair, env_n);
    for (env) |*p| {
        p.key = try c.str(alloc);
        p.value = try c.optStr(alloc);
    }
    d.extra_env = env;
    d.quic = .{};
    d.quic.arm = std.meta.intToEnum(QuicArm, try c.int(u8)) catch return error.BadManifest;
    d.quic.fd = try c.int(i32);
    d.quic.addr_len = try c.int(u32);
    @memcpy(&d.quic.addr, try c.take(QuicState.addr_cap));
    d.quic.idle_ms = try c.int(u64);
    @memcpy(&d.quic.key, try c.take(QuicState.key_len));
    inline for (@typeInfo(Counters).@"struct".fields) |f|
        @field(d.counters, f.name) = try c.int(u64);
    if (!c.done()) return error.BadManifest;
    return d;
}

fn parseSession(c: *Cursor, alloc: std.mem.Allocator) !SessionRec {
    var s: SessionRec = undefined;
    s.name = try c.str(alloc);
    s.pty_fd = try c.int(i32);
    s.child_pid = try c.int(i32);
    s.cols = try c.int(u16);
    s.rows = try c.int(u16);
    s.vt = try c.str(alloc);
    s.title = try c.optStr(alloc);
    s.cmd.phase = try c.int(u8);
    s.cmd.marks_seen = (try c.int(u8)) != 0;
    s.cmd.start_row = try c.int(u32);
    s.cmd.end_row = try c.int(u32);
    const has_code = (try c.int(u8)) != 0;
    const code = try c.int(u8);
    s.cmd.exit_code = if (has_code) code else null;
    if ((try c.int(u8)) != 0) {
        s.last_return = proto.decodeCmdState(try c.take(proto.cmd_state_len)) catch
            return error.BadManifest;
    } else s.last_return = null;
    s.agent_fd = try c.int(i32);
    s.agent_path = try c.optStr(alloc);
    if (!c.done()) return error.BadManifest;
    return s;
}

pub fn parseManifest(gpa: std.mem.Allocator, bytes: []const u8) error{ BadManifest, OutOfMemory }!Parsed {
    var arena = std.heap.ArenaAllocator.init(gpa);
    errdefer arena.deinit();
    const alloc = arena.allocator();

    var c = Cursor{ .bytes = bytes };
    if (!std.mem.eql(u8, try c.take(magic.len), magic)) return error.BadManifest;
    if ((try c.int(u16)) != manifest_version) return error.BadManifest;

    var daemon: ?Daemon = null;
    var sessions: std.ArrayList(SessionRec) = .empty;

    while (!c.done()) {
        const tag = try c.int(u8);
        const len = try c.int(u32);
        var body = Cursor{ .bytes = try c.take(len) };
        switch (@as(SectionTag, @enumFromInt(tag))) {
            .daemon => daemon = try parseDaemon(&body, alloc),
            .session => try sessions.append(alloc, try parseSession(&body, alloc)),
            // An unknown tag was written by a newer binary; its length
            // already walked the cursor past it, so skipping is one arm.
            _ => {},
        }
    }

    return .{
        .daemon = daemon orelse return error.BadManifest,
        .sessions = sessions.items,
        .arena = arena,
    };
}

// ---- tests ----

fn sampleDaemon() Daemon {
    var q = QuicState{ .arm = .owned, .fd = 7, .addr_len = 16, .idle_ms = 30_000 };
    q.addr[0] = 2;
    q.key[0] = 0xAB;
    q.key[31] = 0xCD;
    return .{
        .writer_version = "0.0.1-13",
        .writer_path = "/usr/bin/mux.old",
        .sock_path = "/run/user/1000/muxd.sock",
        .listener_fd = 3,
        .shellint_dir = "/tmp/mux-shim-1234",
        .agent_dir = null,
        .shell = "/bin/zsh",
        .shell_integration = true,
        .extra_env = &.{ .{ .key = "FOO", .value = "bar" }, .{ .key = "UNSET_ME", .value = null } },
        .quic = q,
        .counters = .{
            .snapshots = 5,
            .snapshot_bytes = 6,
            .deltas = 7,
            .delta_bytes = 8,
            .snapshot_equiv_bytes = 10,
            .attaches = 9,
            .agent_refused_no_offer = 11,
            .agent_refused_full = 12,
        },
    };
}

fn sampleSessions() [2]SessionRec {
    return .{
        .{
            .name = "build",
            .pty_fd = 10,
            .child_pid = 4242,
            .cols = 120,
            .rows = 40,
            .vt = "\x1b[2J\x1b[Hhello",
            .title = "make: all",
            .cmd = .{ .phase = 1, .marks_seen = true, .start_row = 3, .end_row = 9, .exit_code = 0 },
            .last_return = .{
                .phase = .at_prompt,
                .mechanism = .marks,
                .exit_code = 0,
                .start_row = 3,
                .end_row = 9,
                .seq = 77,
            },
            .agent_fd = 11,
            .agent_path = "/tmp/mux-agent/agent-build.sock",
        },
        .{
            .name = "0",
            .pty_fd = 12,
            .child_pid = 4243,
            .cols = 80,
            .rows = 24,
            .vt = "",
            .title = null,
            .cmd = .{ .phase = 0, .marks_seen = false, .start_row = 0, .end_row = 0, .exit_code = null },
            .last_return = null,
            .agent_fd = -1,
            .agent_path = null,
        },
    };
}

fn encodeSample(alloc: std.mem.Allocator) ![]u8 {
    var out: std.ArrayList(u8) = .empty;
    errdefer out.deinit(alloc);
    const ss = sampleSessions();
    try writeManifest(out.writer(alloc), alloc, sampleDaemon(), &ss);
    return out.toOwnedSlice(alloc);
}

test "manifest: a round-trip loses nothing a session needs" {
    const alloc = std.testing.allocator;
    const bytes = try encodeSample(alloc);
    defer alloc.free(bytes);

    var p = try parseManifest(alloc, bytes);
    defer p.deinit();

    const d = sampleDaemon();
    try std.testing.expectEqualStrings(d.writer_version, p.daemon.writer_version);
    try std.testing.expectEqualStrings(d.sock_path, p.daemon.sock_path);
    try std.testing.expectEqual(d.listener_fd, p.daemon.listener_fd);
    try std.testing.expectEqualStrings(d.shellint_dir.?, p.daemon.shellint_dir.?);
    try std.testing.expectEqual(@as(?[]const u8, null), p.daemon.agent_dir);
    try std.testing.expect(p.daemon.shell_integration);
    try std.testing.expectEqual(@as(usize, 2), p.daemon.extra_env.len);
    try std.testing.expectEqualStrings("bar", p.daemon.extra_env[0].value.?);
    try std.testing.expectEqual(@as(?[]const u8, null), p.daemon.extra_env[1].value);
    try std.testing.expectEqual(QuicArm.owned, p.daemon.quic.arm);
    try std.testing.expectEqual(@as(u8, 0xCD), p.daemon.quic.key[31]);
    try std.testing.expectEqual(@as(u64, 30_000), p.daemon.quic.idle_ms);
    // Whole-struct, and every field distinct in the sample: a per-field
    // assertion is blind to the counter it forgot, and the `inline for` on
    // both sides means the one it forgets is the one that was added.
    try std.testing.expectEqual(d.counters, p.daemon.counters);

    const ss = sampleSessions();
    try std.testing.expectEqual(@as(usize, 2), p.sessions.len);
    try std.testing.expectEqualStrings(ss[0].name, p.sessions[0].name);
    try std.testing.expectEqual(ss[0].child_pid, p.sessions[0].child_pid);
    try std.testing.expectEqualStrings(ss[0].vt, p.sessions[0].vt);
    try std.testing.expectEqualStrings(ss[0].title.?, p.sessions[0].title.?);
    try std.testing.expect(p.sessions[0].cmd.marks_seen);
    try std.testing.expectEqual(@as(?u8, 0), p.sessions[0].cmd.exit_code);
    try std.testing.expectEqual(@as(u64, 77), p.sessions[0].last_return.?.seq);
    try std.testing.expectEqual(ss[0].agent_fd, p.sessions[0].agent_fd);
}

test "manifest: an unknown section tag is skipped by its length, not fatal" {
    const alloc = std.testing.allocator;
    const bytes = try encodeSample(alloc);
    defer alloc.free(bytes);

    // Splice a tag this build has never heard of between header and body.
    var spliced: std.ArrayList(u8) = .empty;
    defer spliced.deinit(alloc);
    try spliced.appendSlice(alloc, bytes[0 .. magic.len + 2]);
    try spliced.appendSlice(alloc, &.{ 0x7f, 3, 0, 0, 0, 0xAA, 0xBB, 0xCC });
    try spliced.appendSlice(alloc, bytes[magic.len + 2 ..]);

    var p = try parseManifest(alloc, spliced.items);
    defer p.deinit();
    try std.testing.expectEqual(@as(usize, 2), p.sessions.len);
}

test "manifest: a truncated section is BadManifest, not a partial adopt" {
    const alloc = std.testing.allocator;
    const bytes = try encodeSample(alloc);
    defer alloc.free(bytes);
    try std.testing.expectError(
        error.BadManifest,
        parseManifest(alloc, bytes[0 .. bytes.len - 5]),
    );
}

test "manifest: a wrong magic is BadManifest" {
    const alloc = std.testing.allocator;
    const bytes = try encodeSample(alloc);
    defer alloc.free(bytes);
    var bad = try alloc.dupe(u8, bytes);
    defer alloc.free(bad);
    bad[0] = 'X';
    try std.testing.expectError(error.BadManifest, parseManifest(alloc, bad));
}

test "manifest: a session with no agent and no title round-trips its absences" {
    const alloc = std.testing.allocator;
    const bytes = try encodeSample(alloc);
    defer alloc.free(bytes);
    var p = try parseManifest(alloc, bytes);
    defer p.deinit();
    try std.testing.expectEqual(@as(?[]const u8, null), p.sessions[1].title);
    try std.testing.expectEqual(@as(i32, -1), p.sessions[1].agent_fd);
    try std.testing.expectEqual(@as(?[]const u8, null), p.sessions[1].agent_path);
    try std.testing.expectEqual(@as(?proto.CmdState, null), p.sessions[1].last_return);
}

test "manifest: a manifest with no daemon section is BadManifest" {
    const alloc = std.testing.allocator;
    var out: std.ArrayList(u8) = .empty;
    defer out.deinit(alloc);
    try out.appendSlice(alloc, magic);
    try out.appendSlice(alloc, &.{ 1, 0 });
    try std.testing.expectError(error.BadManifest, parseManifest(alloc, out.items));
}