src/server/server_upgrade.zig
Ref: Size: 24.2 KiB History
//! Live daemon upgrade orchestration: adopt and write manifests, validate a
//! candidate, and preserve descriptors across exec and rollback. The wire
//! format lives in upgrade.zig; Server retains the public entry points.
const std = @import("std");
const srv_mod = @import("server.zig");
const Server = srv_mod.Server;
const Options = Server.Options;
const max_sessions = srv_mod.max_sessions;
const SessionTable = @import("server_sessions.zig").SessionTable;
const Engine = @import("engine").Engine;
const Pty = @import("pty").Pty;
const proto = @import("term").protocol;
const shellint = @import("shellint.zig");
const serve = @import("serve");
const server_os = @import("server_os");
const quic_server = @import("quic_server.zig");
const upgrade = @import("upgrade.zig");
/// Adopt the manifest an exec-ing daemon left in its carrier.
pub fn initFromManifest(
alloc: std.mem.Allocator,
parsed: *const upgrade.Parsed,
version: []const u8,
) !Server {
// Before any request can ask: the comparison is against the image
// that BOOTED, not the first one asked about. `mux d upgrade` keeps
// the pid but execs a new image, which zeroes every global — so this
// records the CANDIDATE, and the daemon is graded against the file
// it is now running rather than the one it started life as.
server_os.noteBootImage();
// Same pid, same children, same descriptors. No `sockpath.claim`: the
// inherited listener fd IS the claim, and claim's probe would find our
// own socket answering. `version` is THIS binary's, never the
// manifest's — the skew check measures against what is running now.
const d = parsed.daemon;
// Everything the Server keeps a slice of is copied out of the
// manifest arena, which dies with the caller's `Parsed` while the
// Server outlives it.
var shellint_arena = std.heap.ArenaAllocator.init(alloc);
errdefer shellint_arena.deinit();
const a = shellint_arena.allocator();
const sock_path = try a.dupe(u8, d.sock_path);
const shell_z = try a.dupeZ(u8, d.shell);
const extra_env = try a.alloc(Pty.EnvPair, d.extra_env.len);
for (d.extra_env, extra_env) |src, *dst| dst.* = .{
.key = try a.dupeZ(u8, src.key),
.value = if (src.value) |v| try a.dupeZ(u8, v) else null,
};
// The shim directory crossed as a path: live shells hold it in
// ZDOTDIR and it is pid-named, so a fresh one would orphan the old.
// Rewritten with THIS binary's scripts under the name shells hold —
// `shellint.prepare` creates exclusively, since adopting is an attack.
const injection: shellint.Injection = if (d.shellint_dir) |src| blk: {
const dir = try a.dupe(u8, src);
std.fs.cwd().deleteTree(dir) catch {};
break :blk shellint.prepare(a, dir, shell_z) catch |err| {
std.debug.print(
"mux d: resumed without shell integration ({s}: {t}); " ++
"sessions spawned from here run without command marks\n",
.{ dir, err },
);
break :blk shellint.no_injection;
};
} else shellint.no_injection;
const opts: Options = .{
.sock_path = sock_path,
.shell = shell_z,
.shell_integration = d.shell_integration,
.extra_env = extra_env,
.version = version,
};
const plan = try Server.planFrom(a, opts, injection);
const agent_dir: ?[]const u8 = if (d.agent_dir) |dir| try alloc.dupe(u8, dir) else null;
errdefer if (agent_dir) |dir| alloc.free(dir);
var srv: Server = .{
.alloc = alloc,
.spawn_plan = plan,
.spawn_shell = shell_z,
.spawn_shell_integration = d.shell_integration,
.spawn_extra_env = extra_env,
.version = version,
// The fd is the listener: no bind, no listen, no claim. A claim
// here would probe the path, find our own inherited listener
// answering, and refuse the daemon its own socket. `adopt` only
// re-stamps the id from the file as found.
.bound = try serve.adopt(d.listener_fd, sock_path),
.sock_path = sock_path,
.shellint_arena = shellint_arena,
.shellint_dir = injection.dir,
.agents = .{ .dir = agent_dir },
.forwards = .{ .alloc = alloc },
};
srv_mod.logSocket(sock_path, "adopted across upgrade dev={d} ino={d}", .{ srv.bound.path_id.dev, srv.bound.path_id.ino });
// Cumulative, so an upgrade is not mistaken for a restart by
// anything sampling `mux d stats`. Saturating rather than @intCast
// below: a manifest is bytes, and a corrupt counter must not panic a
// daemon that is otherwise able to serve.
srv.stats = d.counters;
srv.agents.refused_no_offer =
std.math.cast(u32, d.counters.agent_refused_no_offer) orelse std.math.maxInt(u32);
srv.agents.refused_full =
std.math.cast(u32, d.counters.agent_refused_full) orelse std.math.maxInt(u32);
// Frees what this constructor allocated and NOTHING the manifest
// handed over: no descriptor is closed and no child is signalled on
// a failed adoption, because the rollback exec is about to hand all
// of them to the old binary.
errdefer for (&srv.sessions.table) |*slot| {
if (slot.*) |*s| {
s.eng.deinit();
if (s.agentPath()) |p| alloc.free(p);
slot.* = null;
}
};
for (parsed.sessions, 0..) |rec, i| {
// A manifest from a binary with a bigger table would otherwise
// index past ours; the sessions that fit are still served.
if (i >= max_sessions) break;
const eng = try Engine.init(alloc, .{
.cols = rec.cols,
.rows = rec.rows,
.clipboard_max = proto.clipboard_base64_max,
});
// Into the slot before anything else can fail, so the errdefer
// above owns the engine from here on.
srv.sessions.table[i] = .{
.eng = eng,
.pty = Pty.adopt(rec.pty_fd, rec.child_pid),
.epoch = SessionTable.freshEpoch(),
};
const s = &srv.sessions.table[i].?;
eng.feed(rec.vt);
// The title goes back in as an OSC 0 so the ENGINE owns it
// again: `sampleTermTitle` then announces it to the first client
// that attaches, with no resume-shaped exception anywhere
// downstream. title_sent stays null — nobody has been told.
if (rec.title) |t| {
const osc = try std.fmt.allocPrint(alloc, "\x1b]0;{s}\x07", .{t});
defer alloc.free(osc);
eng.feed(osc);
}
s.cmd = .{
.phase = std.meta.intToEnum(proto.CmdPhase, rec.cmd.phase) catch .at_prompt,
.marks_seen = rec.cmd.marks_seen,
.start_row = rec.cmd.start_row,
.end_row = rec.cmd.end_row,
.exit_code = rec.cmd.exit_code,
};
// The verdict crosses; its watermark cannot. The delta tracker is
// rebuilt from zero here, so an old-space seq is a watermark from
// the future that no later return can exceed.
s.last_return = if (rec.last_return) |lr| blk: {
var restamped = lr;
restamped.seq = s.tracker.seq;
break :blk restamped;
} else null;
if (rec.agent_path) |p| {
const dup = try alloc.dupeZ(u8, p);
// `adopt`, not `bind`: the fd crossed the exec already. The
// id is re-stamped from the file as found, which is the
// manifest rule — a watermark belongs to the space that
// minted it, and the pre-exec id was stamped in another.
//
// Re-stamping needs the FILE, so it fails when something has
// deleted the socket out of the runtime dir. That degrades
// per session — the same answer `AgentRelay.bindSock` gives a
// socket it cannot bind — and never fails the adoption: a
// whole daemon refusing to come up would take every shell in
// the table with it, and the rollback exec would then hit the
// identical missing file with MUX_UPGRADE_ROLLBACK already
// set. One session loses agent forwarding instead.
if (serve.adopt(rec.agent_fd, dup)) |b| {
s.agent_sock = .{ .bound = b, .path = dup };
} else |err| {
std.debug.print(
"mux d: no agent socket for session {s} ({t})\n",
.{ SessionTable.safeName(rec.name), err },
);
alloc.free(dup);
// The descriptor is NOT closed here, for the reason the
// errdefer above states: nothing the manifest handed over
// is released on this path, because a later failure hands
// every one of them back to the old binary. It is
// recorded instead — no session names it now, so nothing
// else could find it again — and `sealAdoptedFds` closes
// it once that hand-back window has passed.
srv.orphaned_fds[srv.orphaned_n] = rec.agent_fd;
srv.orphaned_n += 1;
}
}
const n = @min(rec.name.len, proto.session_name_max);
@memcpy(s.name_buf[0..n], rec.name[0..n]);
s.name_len = @intCast(n);
}
// The UDP socket crossed bound and the key as bytes, so TLS stands
// back up on the same port with the same PSK. Both arms become
// `.owned`: whatever held the previous reference went with the image.
switch (d.quic.arm) {
.none => {},
.borrowed, .owned => {
const l = try quic_server.Listener.initFromFd(
alloc,
d.quic.fd,
.{ .bytes = d.quic.key },
undefined,
d.quic.idle_ms,
);
srv.quic = .{ .owned = l };
},
}
return srv;
}
// The manifest an exec-ing daemon leaves for its replacement. The
// writer's version and path are the rollback target's identity.
pub fn writeManifestTo(
self: *Server,
fd: std.posix.fd_t,
writer_version: []const u8,
) !void {
const a = self.alloc;
// The rollback target is THIS binary, resolved here rather than
// taken from a caller: the first draft took it as a parameter and
// both call sites handed it the CANDIDATE's path, which would have
// made a failed adoption exec the broken binary again in a loop.
var exe_buf: [std.fs.max_path_bytes]u8 = undefined;
const writer_path = try std.fs.selfExePath(&exe_buf);
// QUIC runtime state: arm, fd, bound address, idle, key. The key
// crosses as bytes — no path names the carrier and the file may have
// moved. None means no QUIC at all.
var quic_state: upgrade.QuicState = .{};
switch (self.quic) {
.none => {},
.borrowed, .owned => |l| {
quic_state.arm = if (self.quic == .owned) .owned else .borrowed;
quic_state.fd = l.fd;
quic_state.idle_ms = l.idle_ms;
const addr = l.boundAddr();
const addr_bytes = std.mem.asBytes(&addr.any);
quic_state.addr_len = @intCast(addr_bytes.len);
@memcpy(quic_state.addr[0..addr_bytes.len], addr_bytes);
if (quic_server.Listener.currentKey()) |k| @memcpy(&quic_state.key, &k.bytes);
},
}
// extra_env: upgrade.EnvPair is { []const u8, ?[]const u8 } while
// Pty.EnvPair is { [:0]const u8, ?[:0]const u8 }; the manifest's
// shape is what crosses, so the mapping lives here.
const extra_env = try a.alloc(upgrade.EnvPair, self.spawn_extra_env.len);
defer a.free(extra_env);
for (self.spawn_extra_env, extra_env) |src, *dst| {
dst.* = .{ .key = src.key, .value = if (src.value) |v| v else null };
}
// The relay hands over its counters as one value; nothing here reads
// its fields, so the manifest cannot drift when the table's shape does.
const ac = self.agents.counters();
// AgentRelay owns the two refusal counters; `self.stats` only ever
// carried them, so they are stamped from the relay on the way out
// and whatever adoption left in the carrier is overwritten here.
var counters = self.stats;
counters.agent_refused_no_offer = ac.refused_no_offer;
counters.agent_refused_full = ac.refused_full;
const daemon: upgrade.Daemon = .{
.writer_version = writer_version,
.writer_path = writer_path,
.sock_path = self.sock_path,
.listener_fd = self.bound.fd,
.shellint_dir = self.shellint_dir,
.agent_dir = self.agents.dir,
.shell = self.spawn_shell,
.shell_integration = self.spawn_shell_integration,
.extra_env = extra_env,
.quic = quic_state,
.counters = counters,
};
// One SessionRec per live session. dumpState is viewport-only by
// construction; the title is carried apart because dumpState has
// no title field.
var recs: std.ArrayList(upgrade.SessionRec) = .empty;
defer recs.deinit(a);
for (&self.sessions.table) |*slot| {
// BY POINTER: `slot.* orelse continue` copies the Session, and
// `name()` slices that copy's `name_buf`, which dies with the
// iteration — every record would dangle into one reused stack slot.
if (slot.* == null) continue;
const s = &slot.*.?;
const vt = try s.eng.dumpState(a);
const title_bytes = s.eng.title();
const title: ?[]const u8 = if (title_bytes.len > 0) title_bytes else null;
try recs.append(a, .{
.name = s.name(),
.pty_fd = s.pty.master,
.child_pid = s.pty.child,
.cols = @intCast(s.eng.term.cols),
.rows = @intCast(s.eng.term.rows),
.vt = vt,
.title = title,
.cmd = .{
.phase = @intFromEnum(s.cmd.phase),
.marks_seen = s.cmd.marks_seen,
.start_row = s.cmd.start_row,
.end_row = s.cmd.end_row,
.exit_code = s.cmd.exit_code,
},
.last_return = s.last_return,
.agent_fd = s.agentFd(),
.agent_path = s.agentPath(),
});
}
// Write to a buffer first, then to the fd: the manifest writer's
// anytype contract is an ArrayList writer (as the upgrade module's
// own tests use), and serializing to memory avoids coupling the
// manifest format to a particular fd writer's interface.
var manifest_buf: std.ArrayList(u8) = .empty;
defer manifest_buf.deinit(a);
try upgrade.writeManifest(manifest_buf.writer(a), a, daemon, recs.items);
var file = std.fs.File{ .handle = fd };
try file.seekTo(0);
try file.writeAll(manifest_buf.items);
// dumpState allocated each iteration; freed now that the manifest
// carries a copy.
for (recs.items) |r| a.free(r.vt);
}
// A refusal reason, or null to accept. Ordered cheapest refusal first,
// and each reason distinct so the requester can tell what failed.
pub fn validateUpgrade(self: *Server, req: proto.UpgradeReq, my_version: []const u8) ?[]const u8 {
const a = self.alloc;
// 0. No session mid-hangup: `endSession` already closed that master,
// and the exec's `clearCloexec` is `unreachable` on -1, not an error.
// Before the child-run checks, which cost two spawns.
for (self.sessions.table) |slot| {
if (slot) |session| {
if (session.pty.master < 0) return a.dupe(u8, "session ending, retry") catch null;
}
}
// 1. Version: strictly newer, or equal under allow_same_version.
// The reason names BOTH versions so the operator can see the skew.
const newer = upgrade.strictlyNewer(req.version, my_version) catch
return a.dupe(u8, "version: unparseable version string") catch null;
if (!newer) {
if (req.allow_same_version and std.mem.eql(u8, req.version, my_version)) {
// Equal under the flag: allowed.
} else {
return std.fmt.allocPrint(a, "version: {s} is not newer than {s}", .{ req.version, my_version }) catch null;
}
}
// 2. Path: absolute and executable. A relative path would be
// resolved against the daemon's cwd, not the requester's.
if (req.path.len == 0 or req.path[0] != '/')
return a.dupe(u8, "path: not absolute") catch null;
std.posix.access(req.path, std.posix.X_OK) catch
return a.dupe(u8, "path: not executable") catch null;
// 3–4 run the candidate directly: first its advertised version, then
// its ability to resume the manifest this daemon would hand over.
if (checkVersionOutput(self, req.path, req.version)) |reason| return reason;
if (checkManifestResume(self, req.path, my_version)) |reason| return reason;
return null;
}
// Child-run `path --version` must print exactly `mux <version>\n`.
// A daemon of v0.0.1-15 or older expects `muxd <version>` here and so
// refuses this binary as a candidate; `mux d upgrade` says what to do
// about it when the refusal comes back.
fn checkVersionOutput(self: *Server, path: []const u8, version: []const u8) ?[]const u8 {
const a = self.alloc;
const result = std.process.Child.run(.{
.allocator = a,
.argv = &.{ path, "--version" },
.max_output_bytes = 256,
}) catch return a.dupe(u8, "version: cannot spawn candidate") catch null;
defer a.free(result.stdout);
defer a.free(result.stderr);
if (result.term != .Exited or result.term.Exited != 0)
return a.dupe(u8, "version: nonzero exit") catch null;
const expected = std.fmt.allocPrint(a, "mux {s}\n", .{version}) catch
return a.dupe(u8, "version: oom") catch null;
defer a.free(expected);
if (!std.mem.eql(u8, result.stdout, expected))
return a.dupe(u8, "version: output mismatch") catch null;
return null;
}
// Child-run `path d start --resume-fd N --check` must exit 0. The carrier
// is written fresh (not CLOEXEC — children must inherit it).
fn checkManifestResume(self: *Server, path: []const u8, my_version: []const u8) ?[]const u8 {
const a = self.alloc;
const carrier = server_os.anonFd("mux-upgrade") catch
return a.dupe(u8, "check: cannot create the manifest carrier") catch null;
defer std.posix.close(carrier);
self.writeManifestTo(carrier, my_version) catch
return a.dupe(u8, "check: cannot write manifest") catch null;
var fdbuf: [12]u8 = undefined;
const fd_str = std.fmt.bufPrint(&fdbuf, "{d}", .{carrier}) catch
return a.dupe(u8, "check: oom") catch null;
const result = std.process.Child.run(.{
.allocator = a,
.argv = &.{ path, "d", "start", "--resume-fd", fd_str, "--check" },
.max_output_bytes = 4096,
}) catch return a.dupe(u8, "check: cannot spawn candidate") catch null;
defer a.free(result.stdout);
defer a.free(result.stderr);
if (result.term != .Exited or result.term.Exited != 0)
return a.dupe(u8, "check: nonzero exit") catch null;
return null;
}
// Clear FD_CLOEXEC on a descriptor so it survives execve. The upgrade
// exec keeps the listener, QUIC UDP, pty masters, agent listeners and
// the manifest carrier. Most of those are CLOEXEC by default; the
// carrier is not (`server_os.anonFd` never sets the flag) and is on the
// list because `restoreCloexec` re-seals every fd here after a failed
// exec, so a carrier missing from it would stay inheritable by the
// shells this daemon spawns next.
pub fn clearCloexec(fd: std.posix.fd_t) !void {
const flags = try std.posix.fcntl(fd, std.posix.F.GETFD, 0);
_ = try std.posix.fcntl(fd, std.posix.F.SETFD, flags & ~@as(usize, std.posix.FD_CLOEXEC));
}
/// The adopting side of clearCloexec: a flag cleared for the exec that
/// stays cleared is inherited by every shell this daemon spawns next.
pub fn setCloexec(fd: std.posix.fd_t) !void {
const flags = try std.posix.fcntl(fd, std.posix.F.GETFD, 0);
_ = try std.posix.fcntl(fd, std.posix.F.SETFD, flags | @as(usize, std.posix.FD_CLOEXEC));
}
fn restoreCloexec(fds: []const std.posix.fd_t) void {
for (fds) |fd| setCloexec(fd) catch {};
}
fn sealFd(fd: std.posix.fd_t) void {
setCloexec(fd) catch |e|
std.debug.print("mux d: seal fd {d}: {t}\n", .{ fd, e });
}
/// The flag back on every fd `execUpgrade` cleared; the manifest carrier
/// it also cleared is the caller's to close. The caller runs this only
/// once the last rollback point is behind it, which is why the fds no
/// session adopted are CLOSED here rather than sealed: before this point
/// they still belong to the binary a rollback would exec, and after it
/// nothing will ever name them again. Clearing the count keeps a second
/// call from double-closing.
pub fn sealAdoptedFds(self: *Server) void {
sealFd(self.bound.fd);
if (self.quicListener()) |q| sealFd(q.fd);
for (&self.sessions.table) |*slot| {
if (slot.*) |*s| {
sealFd(s.pty.master);
if (s.agentFd() != -1) sealFd(s.agentFd());
}
}
for (self.orphaned_fds[0..self.orphaned_n]) |fd| std.posix.close(fd);
self.orphaned_n = 0;
}
// The exec itself. NOTHING from `Server.deinit` runs here: the process
// that owns the socket path, the shim dir and the agent dir never exits,
// it becomes the new binary. Every sink is bare-closed, NEVER sent
// `exit_status` — that is a dying shell's word and makes clients exit
// instead of redial. A failed exec must leave a working daemon.
pub fn execUpgrade(self: *Server, path: []const u8, carrier: std.posix.fd_t) void {
const a = self.alloc;
// Every return below is a FAILED exec, so the arming is spent either
// way. A `pending_upgrade` left set would put the run loop into an
// exec on a loop.
defer {
a.free(path);
std.posix.close(carrier);
self.pending_upgrade = null;
}
// Collect the fds to clear CLOEXEC on, so they can be restored on
// failure. The listener, QUIC UDP (if any), every pty master, every
// agent listener, and the manifest carrier.
var cleared: std.ArrayList(std.posix.fd_t) = .empty;
defer cleared.deinit(a);
// After the list's own defer, so it runs before it: one restore
// covering every return keeps a half-cleared table from surviving a
// failure that took an early exit.
defer restoreCloexec(cleared.items);
const listener_fd = self.bound.fd;
clearCloexec(listener_fd) catch return;
cleared.append(a, listener_fd) catch return;
if (self.quicListener()) |q| {
clearCloexec(q.fd) catch return;
q.closeAll();
cleared.append(a, q.fd) catch return;
}
for (&self.sessions.table) |*slot| {
const s = slot.* orelse continue;
clearCloexec(s.pty.master) catch return;
cleared.append(a, s.pty.master) catch return;
if (s.agentFd() != -1) {
clearCloexec(s.agentFd()) catch return;
cleared.append(a, s.agentFd()) catch return;
}
}
clearCloexec(carrier) catch return;
cleared.append(a, carrier) catch return;
// Bare-close every client sink and observer. NOT exit_status: that
// tells a client its shell died, and it exits instead of redialing.
for (&self.clients) |*slot| {
if (slot.*) |*c| {
c.pending.deinit(a);
c.inbound.deinit(a);
c.sink.close();
slot.* = null;
}
}
for (&self.observers) |*slot| {
if (slot.*) |*o| {
o.inbound.deinit(a);
std.posix.close(o.fd);
}
}
@memset(&self.observers, null);
// Build argv: {"mux", "d", "start", "--resume-fd", "<n>"}. The
// candidate is by definition NEWER than this binary, so it reads the
// mode word.
var fd_buf: [12]u8 = undefined;
const fd_str = std.fmt.bufPrintZ(&fd_buf, "{d}", .{carrier}) catch return;
const argv = [_:null]?[*:0]const u8{
"mux",
"d",
"start",
"--resume-fd",
fd_str.ptr,
};
// path is []const u8; execveZ needs [*:0]const u8.
const path_z = a.dupeZ(u8, path) catch return;
defer a.free(path_z);
// execveZ returns a plain error set (not an error union): on success
// it never returns, so any return is a failure.
const exec_err = std.posix.execveZ(path_z.ptr, &argv, std.c.environ);
std.debug.print("mux d: upgrade exec failed: {s}\n", .{@errorName(exec_err)});
// execveZ only returns on failure (caught above); on success we
// never reach here.
}