a73x

src/client/discovery.zig

Ref:   Size: 14.4 KiB   History

//! Cancellable picker requests. A Job owns its target, request, result and
//! thread until stop; notifications contain immutable request identity only.
const std = @import("std");
const client = @import("client.zig");
const proto = @import("term").protocol;
const Wait = @import("open_wait.zig").Wait;
const Wire = @import("buffered_wire.zig").Wire;
pub const Ticket = struct { generation: u64, owner: u64, attachment_generation: u64 };
pub const Operation = union(enum) { list, create: struct { name: []const u8, cols: u16, rows: u16 } };
pub const Phase = enum { working, sessions, created, exists, refused, cancelled, failed };
pub const Result = struct {
    ticket: Ticket,
    phase: Phase = .working,
    /// Cancellation/timeout after dispatch cannot undo a remote creation.
    may_have_created: bool = false,
    bytes: [proto.sessions_reply_max]u8 = @splat(0),
    len: usize = 0,
    reason: [1024]u8 = @splat(0),
    reason_len: usize = 0,
    pub fn text(self: *const Result) []const u8 {
        return self.bytes[0..self.len];
    }
    pub fn reasonText(self: *const Result) []const u8 {
        return self.reason[0..self.reason_len];
    }
};
pub const Options = struct {
    ticket: Ticket,
    target: client.Target,
    operation: Operation = .list,
    /// Background catalogues use a noninteractive SSH recipe. An explicit
    /// opening on an existing target must retain that target's custom argv.
    /// Listings remain read-only in either case: inherited startup is cleared.
    use_poll_recipe: bool = true,
    timeout_ms: u32 = 15000,
    wake_ctx: ?*anyopaque = null,
    wake: ?*const fn (?*anyopaque, Ticket) void = null,
};
pub const Job = struct {
    alloc: std.mem.Allocator,
    arena: std.heap.ArenaAllocator,
    opts: Options,
    cancel_pipe: [2]std.posix.fd_t,
    mu: std.Thread.Mutex = .{},
    result: Result,
    thread: ?std.Thread = null,
    done: std.atomic.Value(bool) = .init(false),

    pub fn start(alloc: std.mem.Allocator, options: Options) !*Job {
        var arena = std.heap.ArenaAllocator.init(alloc);
        errdefer arena.deinit();
        var opts = options;
        const source = if (opts.operation == .list and opts.use_poll_recipe) try client.pollTargetFor(arena.allocator(), opts.target) else opts.target;
        opts.target = try cloneTarget(arena.allocator(), source);
        if (opts.operation == .list and opts.target == .hand) {
            opts.target.hand.asked = false;
            opts.target.hand.narrate = false;
        }
        if (opts.operation == .create) {
            const req = &opts.operation.create;
            if (!proto.validSessionName(req.name) or req.cols < 2 or req.rows == 0 or req.cols > proto.max_cols) return error.InvalidCreate;
            req.name = try arena.allocator().dupe(u8, req.name);
        }
        const pipe = try std.posix.pipe2(.{ .NONBLOCK = true, .CLOEXEC = true });
        errdefer for (pipe) |fd| std.posix.close(fd);
        const self = try alloc.create(Job);
        errdefer alloc.destroy(self);
        self.* = .{ .alloc = alloc, .arena = arena, .opts = opts, .cancel_pipe = pipe, .result = .{ .ticket = opts.ticket } };
        self.thread = try std.Thread.spawn(.{}, entry, .{self});
        return self;
    }
    pub fn cancel(self: *Job) void {
        _ = std.posix.write(self.cancel_pipe[1], &.{client.interrupt.detach_key}) catch {};
    }
    /// Join before discarding this request or its callback context. Cancel is
    /// independent of allocation and interrupts open, partial IO and reply waits.
    pub fn stop(self: *Job) void {
        self.cancel();
        _ = self.join();
        for (self.cancel_pipe) |fd| std.posix.close(fd);
        self.arena.deinit();
        self.alloc.destroy(self);
    }
    pub fn join(self: *Job) Result {
        if (self.thread) |t| {
            t.join();
            self.thread = null;
        }
        return self.snapshot();
    }
    pub fn snapshot(self: *Job) Result {
        self.mu.lock();
        defer self.mu.unlock();
        return self.result;
    }
    fn entry(self: *Job) void {
        var out: Result = .{ .ticket = self.opts.ticket };
        self.run(&out) catch |err| {
            out.phase = if (err == error.UserAbort) .cancelled else .failed;
            const reason = @errorName(err);
            if (out.reason_len == 0) {
                out.reason_len = @min(reason.len, out.reason.len);
                @memcpy(out.reason[0..out.reason_len], reason[0..out.reason_len]);
            }
        };
        self.mu.lock();
        self.result = out;
        self.mu.unlock();
        self.done.store(true, .release);
        if (self.opts.wake) |wake| wake(self.opts.wake_ctx, self.opts.ticket);
    }
    fn run(self: *Job, out: *Result) !void {
        const end = std.time.milliTimestamp() + self.opts.timeout_ms;
        var wait: Wait = .{ .alloc = self.alloc, .abort_fd = self.cancel_pipe[0], .deadline = end };
        var dial: client.handoff.Dial = .{};
        var tr = client.Transport.openUntil(self.alloc, self.opts.target, null, self.cancel_pipe[0], &dial, end) catch |err| {
            var buf: [1024]u8 = undefined;
            const failure = client.openFailure(&buf, self.opts.target, err, dial.reason.slice());
            out.reason_len = @min(failure.msg.len, out.reason.len);
            @memcpy(out.reason[0..out.reason_len], failure.msg[0..out.reason_len]);
            return err;
        };
        defer tr.close();
        var wire = try Wire.init(self.alloc, &tr);
        defer wire.deinit();
        var buf: [proto.create_req_max_len]u8 = undefined;
        const kind: proto.MsgType = if (self.opts.operation == .list) .sessions_req else .create_req;
        const payload = switch (self.opts.operation) {
            .create => |req| proto.encodeCreateReq(&buf, req.cols, req.rows, req.name),
            .list => "",
        };
        try wait.check();
        // Conservative boundary: a partial send can commit remotely even if
        // the local write later fails. Never turn cancellation into rollback.
        out.may_have_created = self.opts.operation == .create;
        try wire.send(kind, payload);
        const want: proto.MsgType = if (kind == .sessions_req) .sessions_reply else .create_reply;
        const frame = try awaitReply(&wire, &wait, want);
        defer frame.deinit(self.alloc);
        if (want == .sessions_reply) {
            if (frame.payload.len > out.bytes.len) return error.BadPayload;
            @memcpy(out.bytes[0..frame.payload.len], frame.payload);
            out.len = frame.payload.len;
            out.phase = .sessions;
        } else {
            const reply = try proto.parseCreateReply(frame.payload);
            out.phase = switch (reply.status) {
                .created => .created,
                .exists => .exists,
                .refused => .refused,
            };
            out.may_have_created = reply.status == .created;
            out.reason_len = @min(reply.reason.len, out.reason.len);
            @memcpy(out.reason[0..out.reason_len], reply.reason[0..out.reason_len]);
        }
    }
};
fn awaitReply(wire: *Wire, wait: *Wait, want: proto.MsgType) !proto.Frame {
    while (true) {
        try wait.check();
        try wire.flush();
        wire.tr.service();
        // Stream reads are incremental/nonblocking. QUIC drains buffered
        // frames eagerly, but always returns to cancellation between frames.
        switch (try wire.read()) {
            .closed => return error.Closed,
            .frame => |frame| {
                if (frame.type == want) return frame;
                frame.deinit(wire.alloc);
                continue;
            },
            .incomplete => {},
        }
        var fds = [_]std.posix.pollfd{
            .{ .fd = wire.tr.pollFd(), .events = std.posix.POLL.IN, .revents = 0 },
            .{ .fd = if (wire.tr.link != .quic and wire.pending()) wire.writeFd() else -1, .events = std.posix.POLL.OUT, .revents = 0 },
            .{ .fd = wire.tr.errFd() orelse -1, .events = std.posix.POLL.IN, .revents = 0 },
        };
        try wait.poll(&fds, @intCast(wire.tr.timeoutMs(1000)));
        if (fds[2].revents != 0) wire.tr.drainErr();
    }
}

/// Deep target copy used by off-thread requests. The arena is owned by Job;
/// no nested argv/path/name slice borrows the caller's picker editor buffer.
pub fn cloneTarget(a: std.mem.Allocator, target: client.Target) !client.Target {
    return switch (target) {
        .sock => |s| .{ .sock = try a.dupe(u8, s) },
        .via => |s| .{ .via = try a.dupe(u8, s) },
        .quic => |q| blk: {
            var copy = q;
            copy.host_port = try a.dupe(u8, q.host_port);
            copy.key_path = try a.dupe(u8, q.key_path);
            break :blk .{ .quic = copy };
        },
        .hand => |h| blk: {
            var copy = h;
            copy.host = try a.dupe(u8, h.host);
            copy.ssh_argv = try cloneArgv(a, h.ssh_argv);
            copy.asked_argv = try cloneArgv(a, h.asked_argv);
            if (h.cache_path) |v| copy.cache_path = try a.dupe(u8, v);
            if (h.ask_sock) |v| copy.ask_sock = try a.dupe(u8, v);
            copy.ask_exe = try a.dupe(u8, h.ask_exe);
            break :blk .{ .hand = copy };
        },
    };
}
fn cloneArgv(a: std.mem.Allocator, argv: []const []const u8) ![]const []const u8 {
    const copy = try a.alloc([]const u8, argv.len);
    for (argv, copy) |s, *to| to.* = try a.dupe(u8, s);
    return copy;
}

const TestPeer = struct {
    tmp: @import("testtmp").TmpDir,
    listener: std.net.Server,
    path: []u8,
    fn init() !TestPeer {
        const a = std.testing.allocator;
        var tmp = try @import("testtmp").TmpDir.make();
        errdefer tmp.cleanup();
        const path = try std.fmt.allocPrint(a, "{s}/picker.sock", .{tmp.path()});
        errdefer a.free(path);
        const addr = try std.net.Address.initUnix(path);
        return .{ .tmp = tmp, .path = path, .listener = try addr.listen(.{}) };
    }
    fn deinit(self: *TestPeer) void {
        self.listener.deinit();
        std.testing.allocator.free(self.path);
        self.tmp.cleanup();
    }
    fn accept(self: *TestPeer) !std.net.Stream {
        var fds = [_]std.posix.pollfd{.{ .fd = self.listener.stream.handle, .events = std.posix.POLL.IN, .revents = 0 }};
        if (try std.posix.poll(&fds, 2000) == 0) return error.AcceptTimeout;
        return (try self.listener.accept()).stream;
    }
};
fn resultOf(job: *Job) !Result {
    const end = std.time.milliTimestamp() + 2000;
    while (!job.done.load(.acquire)) {
        if (std.time.milliTimestamp() >= end) return error.JobTimeout;
        std.Thread.sleep(std.time.ns_per_ms);
    }
    return job.snapshot();
}
fn requestFrom(stream: std.net.Stream, expected: proto.MsgType) !void {
    var link = client.Link{ .fd = stream.handle };
    const frame = (try link.awaitFrame(std.testing.allocator, expected, 1000, .{})) orelse return error.RequestTimeout;
    defer frame.deinit(std.testing.allocator);
}

test "picker retarget cancels partial reply and independent generation completes" {
    const a = std.testing.allocator;
    var peer = try TestPeer.init();
    defer peer.deinit();
    const source = try a.dupe(u8, peer.path);
    defer a.free(source);
    const old = try Job.start(a, .{ .ticket = .{ .generation = 1, .owner = 7, .attachment_generation = 2 }, .target = .{ .sock = source } });
    defer old.stop();
    @memset(source, 'x');
    const first = try peer.accept();
    defer first.close();
    try requestFrom(first, .sessions_req);
    try first.writeAll(&.{ @intFromEnum(proto.MsgType.sessions_reply), 20, 0 });
    old.cancel();
    const next = try Job.start(a, .{ .ticket = .{ .generation = 2, .owner = 9, .attachment_generation = 1 }, .target = .{ .sock = peer.path } });
    defer next.stop();
    const second = try peer.accept();
    defer second.close();
    try requestFrom(second, .sessions_req);
    try proto.writeFrame(second.handle, .sessions_reply, "one\ntwo\n");
    const fresh = try resultOf(next);
    const stale = try resultOf(old);
    try std.testing.expectEqual(Phase.sessions, fresh.phase);
    try std.testing.expectEqualStrings("one\ntwo\n", fresh.text());
    try std.testing.expectEqual(Phase.cancelled, stale.phase);
    try std.testing.expect(!std.meta.eql(fresh.ticket, stale.ticket));
    try std.testing.expect(!stale.may_have_created);
}

test "picker create uses explicit operation and reports ambiguous cancellation after dispatch" {
    const a = std.testing.allocator;
    var peer = try TestPeer.init();
    defer peer.deinit();
    for ([_]bool{ false, true }) |cancelled| {
        const job = try Job.start(a, .{ .ticket = .{ .generation = 1, .owner = 1, .attachment_generation = 1 }, .target = .{ .sock = peer.path }, .operation = .{ .create = .{ .name = "0", .cols = 80, .rows = 24 } } });
        defer job.stop();
        const stream = try peer.accept();
        defer stream.close();
        var link = client.Link{ .fd = stream.handle };
        const frame = (try link.awaitFrame(a, .create_req, 1000, .{})) orelse return error.RequestTimeout;
        defer frame.deinit(a);
        try std.testing.expectEqual(proto.MsgType.create_req, frame.type);
        try std.testing.expectEqualStrings("0", (try proto.parseCreateReq(frame.payload)).name);
        if (cancelled) job.cancel() else {
            var buf: [proto.create_reply_max_len]u8 = undefined;
            try proto.writeFrame(stream.handle, .create_reply, proto.encodeCreateReply(&buf, .exists, "name already exists"));
        }
        const result = try resultOf(job);
        try std.testing.expectEqual(if (cancelled) Phase.cancelled else Phase.exists, result.phase);
        try std.testing.expectEqual(cancelled, result.may_have_created);
    }
}

test "same-target listing retains custom SSH argv without inheriting daemon startup" {
    const job = try Job.start(std.testing.allocator, .{
        .ticket = .{ .generation = 1, .owner = 1, .attachment_generation = 1 },
        .use_poll_recipe = false,
        .target = .{ .hand = .{
            .host = "unused.invalid",
            .ssh_argv = &.{ "/bin/sh", "-c", "printf 'endpoint none\\n'; dd bs=5 count=1 of=/dev/null 2>/dev/null; printf '\\221\\002\\000\\000\\0000\\n'; cat >/dev/null" },
            .asked_argv = &.{"false"},
            .cache_path = null,
            .asked = true,
        } },
    });
    defer job.stop();
    const result = try resultOf(job);
    try std.testing.expectEqual(Phase.sessions, result.phase);
    try std.testing.expectEqualStrings("0\n", result.text());
    try std.testing.expect(!result.may_have_created);
}