a73x

src/link.zig

Ref:   Size: 22.7 KiB   History

//! The live connection to a daemon, however it was reached: a unix-socket fd,
//! the stdio of a `--via`/handoff child, or a QUIC client. Mechanics only —
//! send a frame, read a frame, wait, close. Policy stays with the owners:
//! redial and backoff are `client.Transport`'s and muxa's, attach semantics
//! are dial's and the callers'. This row exists because the fd|pipe|quic
//! union used to live twice (client.Transport, muxa.AgentConnection), each
//! with its own send, await and close.
const std = @import("std");
const proto = @import("term").protocol;
const quic = @import("quic");

/// Reap an owned helper even if it ignores SIGTERM. Never signal a child
/// already waited for: its PID may have been reused. wait also closes streams.
pub fn terminateChild(child: *std.process.Child) void {
    if (child.term == null) std.posix.kill(child.id, std.posix.SIG.KILL) catch {};
    _ = child.wait() catch {};
}

/// Three outcomes, not two: QUIC's socket goes readable for acks and half
/// frames, so `null` cannot keep the socket path's meaning of "peer gone"
/// without making every partial frame a reconnect. (Moved from client.zig.)
pub const Incoming = union(enum) {
    frame: proto.Frame,
    incomplete,
    closed,
};

/// What a non-matching frame does to an `awaitFrame` wait. `on == null` is
/// drop: the frame is freed and the wait continues — the observer-verb
/// policy dial.ask always had. A non-null `on` BORROWS the frame for the
/// duration of the call and must copy anything it keeps; awaitFrame frees
/// the frame when `on` returns. An error out of `on` ends the wait with
/// that error — some "other" frames are answers, not noise (muxa's
/// exit_status), and only the caller knows which.
pub const Sink = struct {
    ctx: ?*anyopaque = null,
    on: ?*const fn (ctx: ?*anyopaque, frame: proto.Frame) anyerror!void = null,
};

pub const Link = union(enum) {
    /// A unix socket: one fd, read and write. -1 once closed — close() is
    /// idempotent because a re-dial releases the dead link on entry and an
    /// abort then closes the same value again through a defer; a second
    /// close(2) on a stale fd is EBADF, which std.posix maps to unreachable.
    fd: std.posix.fd_t,
    /// `--via`, and the ssh half of a handoff: the child whose stdio IS the
    /// transport. `r` is its stdout, `w` its stdin.
    pipe: Pipe,
    /// The connection that IS the transport; bytes go through the stream
    /// layer, so there is nothing to write(2) to.
    quic: Quic,

    pub const Pipe = struct {
        child: std.process.Child,
        r: std.posix.fd_t,
        w: std.posix.fd_t,
    };

    pub const Quic = struct {
        cl: *quic.Client,
        /// Bytes the QUIC ring would not take yet. Frames are appended whole
        /// and handed over a prefix at a time, so a short accept can never
        /// split one on the wire — the remainder is offered again next pass.
        /// Lives here and not in the wrappers because the staging is a
        /// property of the quic link, and it used to exist twice
        /// (Transport.qout, muxa.sendFrameQuic's stack buffer).
        qout: std.ArrayList(u8) = .empty,
        alloc: std.mem.Allocator,
    };

    pub fn pollFd(self: *const Link) std.posix.fd_t {
        return switch (self.*) {
            .fd => |fd| fd,
            .pipe => |p| p.r,
            .quic => |q| q.cl.pollFd(),
        };
    }

    /// Folds ngtcp2's next deadline in, so retransmits and idle timeouts
    /// happen on time without a second timer.
    pub fn timeoutMs(self: *Link, cap_ms: i32) i32 {
        return switch (self.*) {
            .quic => |*q| q.cl.timeoutMs(cap_ms),
            .fd, .pipe => cap_ms,
        };
    }

    /// Unconditional: a QUIC connection's timers are the only thing that
    /// notices a peer which stopped answering.
    pub fn service(self: *Link) void {
        switch (self.*) {
            .quic => |*q| {
                q.cl.pump();
                self.flushQuic();
            },
            .fd, .pipe => {},
        }
    }

    /// Set an opaque inbound-byte budget for a role that cannot retain the
    /// terminal wire's normal framing allowance. The transport enforces the
    /// budget before it extends QUIC flow control; it still knows no frames.
    pub fn setInboundCap(self: *Link, cap: ?usize) void {
        switch (self.*) {
            .quic => |*q| q.cl.inbound_cap = cap,
            .fd, .pipe => {},
        }
    }

    /// Queue-and-offer, never blocking: fd and pipe write through; quic
    /// appends whole and flushes what the ring takes. A caller that must
    /// KNOW the bytes left (muxa's verbs) follows with flushWithin.
    pub fn sendFrame(self: *Link, t: proto.MsgType, payload: []const u8) !void {
        switch (self.*) {
            .quic => |*q| {
                try proto.appendFrame(&q.qout, q.alloc, t, payload);
                self.flushQuic();
            },
            .fd => |fd| return proto.writeFrame(fd, t, payload),
            .pipe => |p| return proto.writeFrame(p.w, t, payload),
        }
    }

    /// Offer the outbound queue to the ring again. After every write and on
    /// every service pass, because the room to accept comes from
    /// acknowledgements, which arrive on their own schedule.
    pub fn flushQuic(self: *Link) void {
        const q = switch (self.*) {
            .quic => |*q| q,
            .fd, .pipe => return,
        };
        if (q.qout.items.len == 0) return;
        const n = q.cl.send(q.qout.items);
        if (n == 0) return;
        q.qout.replaceRangeAssumeCapacity(0, n, &.{});
    }

    /// Drive the staged bytes out or say why not, within the deadline. A
    /// no-op for fd/pipe (their sendFrame already either took the bytes or
    /// failed). The QUIC arm is muxa's old sendFrameQuic loop: the ring is
    /// full, only the peer's acks empty it, and they arrive through pump —
    /// polling first keeps this from spinning.
    pub fn flushWithin(self: *Link, deadline_ms: u32) !void {
        const q = switch (self.*) {
            .quic => |*q| q,
            .fd, .pipe => return,
        };
        const end = std.time.milliTimestamp() + deadline_ms;
        while (q.qout.items.len != 0) {
            if (q.cl.dead) return error.ConnectionLost;
            self.flushQuic();
            if (q.qout.items.len == 0) return;
            if (std.time.milliTimestamp() >= end) return error.SendStalled;
            var fds = [_]std.posix.pollfd{
                .{ .fd = q.cl.pollFd(), .events = std.posix.POLL.IN, .revents = 0 },
            };
            _ = std.posix.poll(&fds, q.cl.timeoutMs(50)) catch return error.ConnectionLost;
            q.cl.pump();
        }
    }

    /// The next whole frame, if there is one. (Moved from
    /// client.Transport.readFrame; see Incoming for why a missing frame is
    /// not automatically a dead transport.)
    ///
    /// Two errors are NOT transport events and are returned as themselves.
    /// `error.OutOfMemory` is this side's failure, not the peer's. And a
    /// header claiming more than `max_payload` is a peer that spoke and got
    /// the framing wrong, which is a different fact from a peer that went
    /// away: `dial.ask` reports a closed peer as "no answer" (null) and every
    /// other error as an error, precisely so a corrupt reply to `mux d dump`
    /// or `stats` is never reported as an absent daemon. Everything else —
    /// a read errno, a frame cut short mid-payload — is the peer gone.
    pub fn readFrame(self: *Link, alloc: std.mem.Allocator) !Incoming {
        switch (self.*) {
            .quic => |*q| {
                // Death is checked after the pump, so bytes that arrived in
                // the same pass as the close are still delivered before the
                // tear.
                // takeFrame reads a buffer rather than a socket, so
                // OutOfMemory and FrameTooLarge are its ONLY failures: there
                // is no errno arm left to fold into .closed here.
                const got = try proto.takeFrame(alloc, &q.cl.in);
                if (got) |frame| return .{ .frame = frame };
                return if (q.cl.dead) .closed else .incomplete;
            },
            .fd => |fd| {
                const frame = (proto.readFrame(alloc, fd) catch |err| switch (err) {
                    error.OutOfMemory, error.FrameTooLarge => return err,
                    else => return .closed,
                }) orelse return .closed;
                return .{ .frame = frame };
            },
            .pipe => |p| {
                const frame = (proto.readFrame(alloc, p.r) catch |err| switch (err) {
                    error.OutOfMemory, error.FrameTooLarge => return err,
                    else => return .closed,
                }) orelse return .closed;
                return .{ .frame = frame };
            },
        }
    }

    /// Wait for one frame of type `want`. THE primitive this row exists
    /// for; every hand-rolled poll+readFrame loop in the tree is a copy of
    /// this. Returns null when the deadline runs out (a null deadline waits
    /// forever), error.Closed when the peer is gone — callers own the
    /// wording for both (dial.ask maps Closed to its "no answer" null; muxa
    /// maps it to DaemonGone/ConnectionLost). Non-matching frames go to the
    /// sink (see Sink for ownership). A `deadline_ms` of 0, or a deadline
    /// already spent by the time the loop is re-entered, returns null off an
    /// fd or pipe arm without polling once — a caller that means "hand me
    /// whatever has already arrived" must pass a nonzero deadline. QUIC is
    /// the exception: its read is eager, so a frame already buffered from an
    /// earlier pump still comes back.
    pub fn awaitFrame(
        self: *Link,
        alloc: std.mem.Allocator,
        want: proto.MsgType,
        deadline_ms: ?u32,
        sink: Sink,
    ) !?proto.Frame {
        const end: ?i64 = if (deadline_ms) |ms| std.time.milliTimestamp() + ms else null;
        // Whether this link may be asked for a frame without poll saying so
        // first. QUIC may: a whole frame can already sit in cl.in from an
        // earlier pump, and poll would never fire for bytes that are already
        // in userspace. A BOUNDED fd or pipe may not: readFrame there is a
        // read(2) on a stream, which on a quiet peer returns only when that
        // peer speaks or goes away — the wait has to be spent in poll, where
        // the deadline can end it, which is why dial.ask has always polled
        // before reading too.
        // With no deadline the blocking read IS the wait, so they read
        // straight away and never reach the poll below.
        const eager = switch (self.*) {
            .quic => true,
            .fd, .pipe => end == null,
        };
        var ready = eager;
        while (true) {
            if (ready) {
                // A non-match falls back to `eager`: on a bounded fd the next
                // frame has to be announced by poll again, or draining noise
                // would walk straight into a blocking read.
                ready = eager;
                switch (try self.readFrame(alloc)) {
                    .frame => |frame| {
                        if (frame.type == want) return frame;
                        if (sink.on) |on| {
                            defer frame.deinit(alloc);
                            try on(sink.ctx, frame);
                        } else frame.deinit(alloc);
                        continue;
                    },
                    .closed => return error.Closed,
                    .incomplete => {},
                }
            }
            // Nothing whole in hand: wait. The wait is capped at 250ms even
            // with no caller deadline, because a QUIC link's timers (loss
            // detection, keepalive) need servicing on schedule rather than
            // whenever the daemon happens to say something; a plain fd with
            // no deadline may block indefinitely, which is what a caller
            // wants when a stopped daemon should be a visible hang.
            var cap: i32 = undefined;
            if (end) |e| {
                const left = e - std.time.milliTimestamp();
                if (left <= 0) return null;
                cap = @intCast(@min(left, 250));
            } else cap = switch (self.*) {
                .quic => 250,
                .fd, .pipe => -1,
            };
            var fds = [_]std.posix.pollfd{
                .{ .fd = self.pollFd(), .events = std.posix.POLL.IN, .revents = 0 },
            };
            const n = std.posix.poll(&fds, self.timeoutMs(cap)) catch return error.Closed;
            self.service();
            // A poll that timed out leaves a bounded fd with nothing to read,
            // and going back to read(2) on it would block past the deadline;
            // the next pass re-checks the clock instead. QUIC always looks
            // again — readFrame never blocks there, and the pump just now may
            // have completed a frame out of bytes that arrived earlier.
            ready = eager or n != 0;
            // Once poll HAS announced bytes, the fd arms still block in
            // readFrame until that frame is whole, even past the deadline —
            // the price of frames over a stream, same as dial.ask always
            // paid; a malformed frame is an error rather than an absent
            // reply.
        }
    }

    /// Idempotent, and it has to be: a re-dial releases the dead link on
    /// entry, and an abort then closes the same value again through the
    /// pump's defer. (The pipe-kill ordering is client.Transport.close's,
    /// moved: stdin first so the command sees EOF and can wind down its
    /// remote end, then forced termination and wait — no zombie or a
    /// SIGTERM-ignoring child that can hold cancellation indefinitely.)
    pub fn close(self: *Link) void {
        switch (self.*) {
            .fd => |fd| {
                if (fd == -1) return; // already released
                std.posix.close(fd);
            },
            .pipe => |*p| {
                if (p.child.stdin) |*in| {
                    in.close();
                    p.child.stdin = null;
                }
                terminateChild(&p.child);
            },
            .quic => |*q| {
                q.qout.deinit(q.alloc);
                // Closes the UDP socket with it, so pollFd's value must not
                // be closed again.
                q.cl.deinit();
            },
        }
        self.* = .{ .fd = -1 };
    }
};

const testing = std.testing;

fn mkPair() ![2]std.posix.fd_t {
    var pair: [2]std.posix.fd_t = undefined;
    try testing.expectEqual(@as(c_int, 0), std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair));
    return .{ pair[0], pair[1] };
}

test "fd link: sendFrame puts one frame on the wire, readFrame takes it back" {
    const alloc = testing.allocator;
    const pair = try mkPair();
    var a: Link = .{ .fd = pair[0] };
    var b: Link = .{ .fd = pair[1] };
    defer a.close();
    defer b.close();
    try a.sendFrame(.input, "hi");
    const inc = try b.readFrame(alloc);
    const frame = inc.frame;
    defer frame.deinit(alloc);
    try testing.expectEqual(proto.MsgType.input, frame.type);
    try testing.expectEqualStrings("hi", frame.payload);
}

test "pipe link: the child's stdio IS the transport, and r and w are not interchangeable" {
    const alloc = testing.allocator;
    // A real child rather than a socketpair standing in for one, because the
    // thing under test is which of the child's two fds each direction uses.
    // `cat` echoes, so a frame that comes back proves sendFrame wrote the
    // child's STDIN and readFrame read its STDOUT; swap the pair and both
    // directions address the wrong end of a pipe and fail. A socketpair is
    // symmetric and would pass either way.
    var child = std.process.Child.init(&.{"cat"}, alloc);
    child.stdin_behavior = .Pipe;
    child.stdout_behavior = .Pipe;
    try child.spawn();
    var l: Link = .{ .pipe = .{
        .child = child,
        .r = child.stdout.?.handle,
        .w = child.stdin.?.handle,
    } };
    defer l.close();
    try l.sendFrame(.input, "echo me");
    const got = (try l.awaitFrame(alloc, .input, 5000, .{})).?;
    defer got.deinit(alloc);
    try testing.expectEqualStrings("echo me", got.payload);
}

test "close is idempotent on every arm that can be closed twice" {
    const pair = try mkPair();
    std.posix.close(pair[1]);
    var l: Link = .{ .fd = pair[0] };
    l.close();
    l.close(); // second close must be a no-op, not EBADF-unreachable
}

test "awaitFrame: null sink drops noise frames and returns the match" {
    const alloc = testing.allocator;
    const pair = try mkPair();
    var l: Link = .{ .fd = pair[0] };
    defer l.close();
    // Two noise frames BEFORE the answer: the wait must survive a stream,
    // not a single-frame fixture.
    try proto.writeFrame(pair[1], .pty_mode, &.{0});
    try proto.writeFrame(pair[1], .delta, "x");
    try proto.writeFrame(pair[1], .stats_reply, "ok");
    std.posix.close(pair[1]);
    const got = (try l.awaitFrame(alloc, .stats_reply, 1000, .{})).?;
    defer got.deinit(alloc);
    try testing.expectEqualStrings("ok", got.payload);
}

test "awaitFrame: the sink sees every non-match, in order, and may end the wait" {
    const alloc = testing.allocator;
    const Seen = struct {
        types: [8]proto.MsgType = undefined,
        n: usize = 0,
        fn on(ctx: ?*anyopaque, frame: proto.Frame) anyerror!void {
            const self: *@This() = @ptrCast(@alignCast(ctx.?));
            self.types[self.n] = frame.type;
            self.n += 1;
            if (frame.type == .exit_status) return error.SessionExited;
        }
    };
    // Leg one: the sink observes and the wait completes.
    {
        const pair = try mkPair();
        var l: Link = .{ .fd = pair[0] };
        defer l.close();
        try proto.writeFrame(pair[1], .snapshot, "");
        try proto.writeFrame(pair[1], .stats_reply, "ok");
        std.posix.close(pair[1]);
        var seen: Seen = .{};
        const got = (try l.awaitFrame(alloc, .stats_reply, 1000, .{ .ctx = &seen, .on = Seen.on })).?;
        defer got.deinit(alloc);
        try testing.expectEqual(@as(usize, 1), seen.n);
        try testing.expectEqual(proto.MsgType.snapshot, seen.types[0]);
    }
    // Leg two: the sink's error IS the outcome — an "other" frame that is
    // an answer, muxa's whole reason for a second copy.
    {
        const pair = try mkPair();
        var l: Link = .{ .fd = pair[0] };
        defer l.close();
        try proto.writeFrame(pair[1], .exit_status, &.{7});
        std.posix.close(pair[1]);
        var seen: Seen = .{};
        try testing.expectError(
            error.SessionExited,
            l.awaitFrame(alloc, .stats_reply, 1000, .{ .ctx = &seen, .on = Seen.on }),
        );
    }
}

test "awaitFrame: silence consumes the deadline and returns null; a closed peer is error.Closed" {
    const alloc = testing.allocator;
    {
        const pair = try mkPair();
        defer std.posix.close(pair[1]);
        var l: Link = .{ .fd = pair[0] };
        defer l.close();
        const t0 = std.time.milliTimestamp();
        try testing.expect((try l.awaitFrame(alloc, .stats_reply, 100, .{})) == null);
        try testing.expect(std.time.milliTimestamp() - t0 >= 100);
    }
    {
        const pair = try mkPair();
        var l: Link = .{ .fd = pair[0] };
        defer l.close();
        std.posix.close(pair[1]);
        try testing.expectError(error.Closed, l.awaitFrame(alloc, .stats_reply, 1000, .{}));
    }
}

test "awaitFrame: no deadline waits out a late reply" {
    const alloc = testing.allocator;
    const pair = try mkPair();
    var l: Link = .{ .fd = pair[0] };
    defer l.close();
    const Late = struct {
        fn run(fd: std.posix.fd_t) void {
            std.Thread.sleep(150 * std.time.ns_per_ms);
            proto.writeFrame(fd, .stats_reply, "late") catch {};
            std.posix.close(fd);
        }
    };
    const t = try std.Thread.spawn(.{}, Late.run, .{pair[1]});
    defer t.join();
    const got = (try l.awaitFrame(alloc, .stats_reply, null, .{})).?;
    defer got.deinit(alloc);
    try testing.expectEqualStrings("late", got.payload);
}

// The two pins below moved here from dial.zig when `dial.ask` stopped
// carrying its own await loop: both are properties of THIS loop, and the
// contract sentence they protect ("a corrupt frame is never reported as
// silence") is still dial.ask's doc comment.

test "awaitFrame: an empty payload of the wanted type is the answer, not a frame to skip" {
    // The verb's decoder decides what an empty reply means. This loop used to
    // carry a per-caller policy so that `upgrade_reply` could keep waiting
    // through one; `protocol.parseUpgradeReply` now answers that question
    // where the rest of the upgrade wire is read.
    const alloc = testing.allocator;
    const pair = try mkPair();
    defer std.posix.close(pair[1]);
    var l: Link = .{ .fd = pair[0] };
    defer l.close();
    try proto.writeFrame(pair[1], .upgrade_reply, "");
    try proto.writeFrame(pair[1], .upgrade_reply, &.{0});

    const frame = (try l.awaitFrame(alloc, .upgrade_reply, 500, .{})).?;
    defer frame.deinit(alloc);
    try testing.expectEqual(@as(usize, 0), frame.payload.len);
    try testing.expect(proto.parseUpgradeReply(frame.payload) == null);
}

test "awaitFrame: a frame this side cannot read is an error, never silence" {
    // Preserve corrupt-frame errors so dump and stats distinguish a broken
    // daemon response from an absent daemon.
    const pair = try mkPair();
    defer std.posix.close(pair[1]);
    var l: Link = .{ .fd = pair[0] };
    defer l.close();
    var hdr: [5]u8 = undefined;
    hdr[0] = @intFromEnum(proto.MsgType.stats_reply);
    std.mem.writeInt(u32, hdr[1..5], proto.max_payload + 1, .little);
    try proto.writeAllFd(pair[1], &hdr);
    try testing.expectError(
        error.FrameTooLarge,
        l.awaitFrame(testing.allocator, .stats_reply, null, .{}),
    );
}

test "forced child cleanup reaps a peer ignoring SIGTERM" {
    const a = std.testing.allocator;
    var child = std.process.Child.init(&.{ "sh", "-c", "trap '' TERM; printf R; while :; do :; done" }, a);
    child.stdin_behavior = .Pipe;
    child.stdout_behavior = .Pipe;
    try child.spawn();
    defer terminateChild(&child);
    const pid = child.id;
    var fds = [_]std.posix.pollfd{.{ .fd = child.stdout.?.handle, .events = std.posix.POLL.IN, .revents = 0 }};
    if (try std.posix.poll(&fds, 1000) == 0) return error.ChildTimeout;
    var byte: [1]u8 = undefined;
    try std.testing.expectEqual(@as(usize, 1), try std.posix.read(child.stdout.?.handle, &byte));
    const start = std.time.milliTimestamp();
    terminateChild(&child);
    try std.testing.expect(std.time.milliTimestamp() - start < 500);
    try std.testing.expectError(error.ProcessNotFound, std.posix.kill(pid, 0));
}