src/cli/mux.zig

Ref: main worktree-web-client-spec   Size: 5.2 KiB   History

//! Top-level dispatcher for the `mux` binary. The first argument selects `d`
//! (daemon), `a` (JSON agent), or `web` (browser hub); every other argument is
//! handled by the client. There is no top-level flag grammar or command alias,
//! so `mux run` treats `run` as a hostname.
//!
//! The ssh askpass mode has no command word. It is selected through the
//! `SSH_ASKPASS` environment and receives the prompt as argv[1].
const std = @import("std");
const daemon = @import("main.zig");
const agent = @import("agent");
const hub = @import("webhub_main.zig");
const client = @import("mux_main.zig");
const askpass = @import("client").askpass;

/// The dispatch decision as a value, allowing it to be tested without exiting
/// a process.
const DispatchMode = enum { daemon, agent, hub, client, askpass };

fn modeOf(args: []const [:0]const u8, ask_sock: ?[]const u8) DispatchMode {
    if (args.len < 2) return .client;
    if (std.mem.eql(u8, args[1], "d")) return .daemon;
    if (std.mem.eql(u8, args[1], "a")) return .agent;
    if (std.mem.eql(u8, args[1], "web")) return .hub;
    // An explicit mode takes precedence over the askpass environment. This
    // keeps remote `mux d endpoint` invocations and commands typed inside an
    // SSH process tree from being mistaken for password prompts. Askpass is
    // valid only for the single prompt argument supplied by SSH.
    if (ask_sock != null and args.len == 2) return .askpass;
    return .client;
}

pub fn main() !u8 {
    var gpa: std.heap.DebugAllocator(.{}) = .init;
    defer if (gpa.deinit() == .leak)
        std.debug.print("mux: LEAK: allocations outlived deinit\n", .{});
    const alloc = gpa.allocator();

    const args = try std.process.argsAlloc(alloc);
    defer std.process.argsFree(alloc, args);
    if (try @import("client").resolver.helper(alloc, args[1..])) |code| return code;

    // Named modes receive a slice beginning with their mode word, matching the
    // argv shape expected by their parsers. The client has no mode word and
    // therefore receives the complete argv.
    const ask_sock = std.posix.getenv(askpass.sock_env);
    const ask_kind = askpass.Kind.of(std.posix.getenv(askpass.prompt_env));
    return switch (modeOf(args, ask_sock)) {
        .daemon => daemon.main(args[1..]),
        .agent => agent.main(args[1..]),
        .hub => hub.main(args[1..]),
        .client => client.main(args),
        // SSH reads the response from stdout, so that stream must contain only
        // the askpass result. The prompt-kind environment controls whether the
        // wall masks the displayed response.
        .askpass => askpass.helperMain(args[1], ask_sock.?, ask_kind, std.posix.STDOUT_FILENO),
    };
}

test "modeOf: the three mode words, and nothing else" {
    try std.testing.expectEqual(DispatchMode.daemon, modeOf(&.{ "mux", "d", "start" }, null));
    try std.testing.expectEqual(DispatchMode.agent, modeOf(&.{ "mux", "a", "status" }, null));
    try std.testing.expectEqual(DispatchMode.hub, modeOf(&.{ "mux", "web" }, null));
    try std.testing.expectEqual(DispatchMode.client, modeOf(&.{"mux"}, null));
    try std.testing.expectEqual(DispatchMode.client, modeOf(&.{ "mux", "box" }, null));
}

test "modeOf: the helper is named by ssh's variable, and a mode word still wins" {
    const sock = "/run/user/1000/mux-ask-7.sock";
    // The whole of what ssh hands a helper: one argument, the prompt.
    try std.testing.expectEqual(DispatchMode.askpass, modeOf(&.{ "mux", "box's password: " }, sock));
    // Without the askpass environment, the same argv is treated as a hostname.
    try std.testing.expectEqual(DispatchMode.client, modeOf(&.{ "mux", "box's password: " }, null));
    // Inside the tree, a mode word is still a mode: the remote command the
    // wall's ssh runs is `mux d endpoint`, and a user typing one at that
    // ssh's far end would otherwise get a prompt helper.
    try std.testing.expectEqual(DispatchMode.daemon, modeOf(&.{ "mux", "d", "endpoint" }, sock));
    try std.testing.expectEqual(DispatchMode.agent, modeOf(&.{ "mux", "a", "status" }, sock));
    try std.testing.expectEqual(DispatchMode.hub, modeOf(&.{ "mux", "web" }, sock));
    // SSH askpass supplies exactly one prompt argument. Longer invocations are
    // dispatched to the client even if they inherit the environment variable.
    try std.testing.expectEqual(DispatchMode.client, modeOf(&.{ "mux", "box", "-A" }, sock));
    try std.testing.expectEqual(DispatchMode.client, modeOf(&.{"mux"}, sock));
}

test "modeOf: `run` is a host, and no daemon verb has a top-level alias" {
    // There is no compatibility alias for the old bare `run` command. Older
    // daemons reject this binary during their version check before executing
    // it, and every unrecognized mode word remains a client target.
    try std.testing.expectEqual(DispatchMode.client, modeOf(&.{ "mux", "run" }, null));
    try std.testing.expectEqual(DispatchMode.client, modeOf(&.{ "mux", "run", "--resume-fd", "5" }, null));
}

test {
    std.testing.refAllDeclsRecursive(@This());
    // These imports make the child modules' tests reachable from this root test
    // artifact. The agent and webhub modules also have separate test artifacts.
    _ = @import("main.zig");
    _ = @import("mux_main.zig");
    _ = @import("webhub_main.zig");
}