//! Linux kernel-mount launcher. Installed as an ordinary executable beside
//! Pardes; sudo authorizes each launch. Never install setuid.
const std = @import("std");
const builtin = @import("builtin");

extern "c" fn unshare(flags: c_int) c_int;
extern "c" fn mount(source: ?[*:0]const u8, target: [*:0]const u8, filesystemtype: ?[*:0]const u8, flags: c_ulong, data: ?*const anyopaque) c_int;
extern "c" fn getuid() c_uint;
extern "c" fn geteuid() c_uint;
extern "c" fn setgroups(size: usize, list: ?[*]const c_uint) c_int;
extern "c" fn initgroups(user: [*:0]const u8, group: c_uint) c_int;
extern "c" fn setresgid(real: c_uint, effective: c_uint, saved: c_uint) c_int;
extern "c" fn setresuid(real: c_uint, effective: c_uint, saved: c_uint) c_int;
extern "c" fn execv(path: [*:0]const u8, argv: [*:null]const ?[*:0]const u8) c_int;
extern "c" fn execvp(path: [*:0]const u8, argv: [*:null]const ?[*:0]const u8) c_int;
extern "c" fn readlink(path: [*:0]const u8, buf: [*]u8, size: usize) isize;
extern "c" fn mkdtemp(template: [*:0]u8) ?[*:0]u8;
extern "c" fn rmdir(path: [*:0]const u8) c_int;
extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int;
extern "c" fn unsetenv(name: [*:0]const u8) c_int;
extern "c" fn fork() c_int;
extern "c" fn waitpid(pid: c_int, status: *c_int, flags: c_int) c_int;
extern "c" fn kill(pid: c_int, sig: c_int) c_int;
extern "c" fn _exit(status: c_int) noreturn;

pub const executable = "pardes-v9fs";

/// One command for the normal interactive shell's startup queue. Quote every
/// argument so paths and shell setup strings remain data in bash, fish and sh.
pub fn writeLaunchCommand(writer: *std.Io.Writer, helper: []const u8, socket: []const u8, shell_argv: []const ?[*:0]const u8) !void {
    try quote(writer, helper);
    try writer.writeAll(" --launch ");
    try quote(writer, socket);
    try writer.writeAll(" --");
    for (shell_argv) |maybe| {
        const arg = maybe orelse break;
        try writer.writeByte(' ');
        try quote(writer, std.mem.span(arg));
    }
}

fn quote(writer: *std.Io.Writer, value: []const u8) !void {
    // A literal newline would submit the interactive command before it is
    // complete. These are paths/setup arguments, not arbitrary shell input.
    if (std.mem.indexOfAny(u8, value, "\r\n") != null) return error.MultilineLaunchArgument;
    try writer.writeByte('\'');
    for (value) |byte| {
        if (byte == '\'') try writer.writeAll("'\\''") else try writer.writeByte(byte);
    }
    try writer.writeByte('\'');
}

fn selfPath(buf: []u8) ![:0]u8 {
    const n = readlink("/proc/self/exe", buf.ptr, buf.len - 1);
    if (n < 0 or n >= buf.len - 1) return error.ExecutablePathUnavailable;
    const len: usize = @intCast(n);
    buf[len] = 0;
    return buf[0..len :0];
}

/// Resolve before fork, without relying on the shell's PATH. The override is
/// useful for build-cache binaries whose helper is in a different directory.
pub fn helperPath(buf: []u8) ![:0]u8 {
    const path = if (std.c.getenv("PARDES_V9FS_HELPER")) |env| blk: {
        const override = std.mem.span(env);
        if (!std.fs.path.isAbsolute(override)) return error.AbsoluteHelperPathRequired;
        break :blk try std.fmt.bufPrintZ(buf, "{s}", .{override});
    } else blk: {
        var own: [4096]u8 = undefined;
        const parent = std.fs.path.dirname(try selfPath(&own)) orelse return error.ExecutablePathUnavailable;
        break :blk try std.fmt.bufPrintZ(buf, "{s}/{s}", .{ parent, executable });
    };
    if (std.c.access(path.ptr, 1) != 0) return error.V9fsHelperNotFound;
    return path;
}

fn usage() void {
    std.debug.print(
        \\usage: pardes-v9fs --launch SOCKET -- /absolute/shell [args...]
        \\  Ask sudo in this terminal, mount privately, and start an unprivileged shell.
        \\  PARDES_MOUNT names the mount. The caller's environment is preserved with sudo -E.
        \\internal: pardes-v9fs SOCKET MOUNTPOINT UID GID -- /absolute/command [args...]
        \\  Requires explicit root execution. Never install setuid or grant blanket NOPASSWD.
        \\
    , .{});
}

var sudo_pid: std.atomic.Value(c_int) = .init(-1);

fn forwardSignal(sig: std.posix.SIG) callconv(.c) void {
    const pid = sudo_pid.load(.monotonic);
    if (pid > 0) _ = kill(pid, @intCast(@intFromEnum(sig)));
}

fn launch(arena: std.mem.Allocator, args: []const [:0]const u8) !u8 {
    if (args.len < 5 or !std.mem.eql(u8, args[3], "--")) return error.InvalidArguments;
    if (geteuid() == 0 or getuid() != geteuid()) return error.UnprivilegedLaunchRequired;
    if (!std.fs.path.isAbsolute(args[2]) or !std.fs.path.isAbsolute(args[4])) return error.AbsolutePathRequired;
    var own: [4096]u8 = undefined;
    const helper = try selfPath(&own);
    var target = "/tmp/pardes-v9fs-XXXXXX".*;
    if (mkdtemp(&target) == null) return error.MountDirectoryFailed;
    defer _ = rmdir(&target);
    try check(setenv("PARDES_MOUNT", &target, 1), "export mount path");
    inline for (.{ "PATH", "HOME", "USER", "LOGNAME", "SHELL" }) |name| {
        const saved = "PARDES_V9FS_" ++ name;
        if (std.c.getenv(name)) |value| {
            try check(setenv(saved, value, 1), "preserve shell environment");
        } else _ = unsetenv(saved);
    }
    const uid = try std.fmt.allocPrintSentinel(arena, "{d}", .{getuid()}, 0);
    const gid = try std.fmt.allocPrintSentinel(arena, "{d}", .{std.c.getgid()}, 0);
    const prefix = [_]?[*:0]const u8{ "sudo", "-E", "--", helper.ptr, args[2].ptr, &target, uid.ptr, gid.ptr, "--" };
    const argv = try arena.allocSentinel(?[*:0]const u8, prefix.len + args.len - 4, null);
    @memcpy(argv[0..prefix.len], &prefix);
    for (args[4..], prefix.len..) |arg, i| argv[i] = arg.ptr;
    std.debug.print("Mounting Pardes at {s}\n", .{target});
    const pid = fork();
    if (pid < 0) return error.ForkFailed;
    if (pid == 0) {
        // host_io's resetChildSignals, here in its own module: every signal
        // back to its default and none blocked before the exec.
        const default: std.posix.Sigaction = .{ .handler = .{ .handler = std.posix.SIG.DFL }, .mask = std.posix.sigemptyset(), .flags = 0 };
        var sig: u8 = 1;
        while (sig < 65) : (sig += 1) {
            if (sig == @intFromEnum(std.posix.SIG.KILL) or sig == @intFromEnum(std.posix.SIG.STOP)) continue;
            _ = std.c.sigaction(@enumFromInt(sig), &default, null);
        }
        const none = std.posix.sigemptyset();
        std.posix.sigprocmask(std.posix.SIG.SETMASK, &none, null);
        _ = execvp("sudo", argv.ptr);
        _exit(127);
    }
    sudo_pid.store(pid, .monotonic);
    const action: std.posix.Sigaction = .{ .handler = .{ .handler = forwardSignal }, .mask = std.posix.sigemptyset(), .flags = 0 };
    for ([_]std.posix.SIG{ .HUP, .INT, .TERM }) |sig| std.posix.sigaction(sig, &action, null);
    var status: c_int = 0;
    while (waitpid(pid, &status, 0) < 0) {
        if (std.posix.errno(-1) != .INTR) return error.WaitFailed;
    }
    sudo_pid.store(-1, .monotonic);
    return if (status & 0x7f == 0) @intCast((status >> 8) & 0xff) else @intCast(128 + (status & 0x7f));
}

fn check(rc: c_int, operation: []const u8) !void {
    if (rc == 0) return;
    const err = std.posix.errno(rc);
    std.debug.print("v9fs: {s}: {s}\n", .{ operation, @tagName(err) });
    return error.SystemCallFailed;
}

fn userId(text: []const u8) !c_uint {
    const id = std.fmt.parseInt(c_uint, text, 10) catch return error.InvalidUserId;
    if (id == 0 or id == std.math.maxInt(c_uint)) return error.InvalidUserId;
    return id;
}

pub fn main(init: std.process.Init) !void {
    if (builtin.os.tag != .linux) return error.LinuxRequired;
    const arena = init.arena.allocator();
    const args = try init.minimal.args.toSlice(arena);
    if (args.len == 2 and std.mem.eql(u8, args[1], "--help")) {
        usage();
        return;
    }
    if (args.len > 1 and std.mem.eql(u8, args[1], "--launch")) {
        const status = try launch(arena, args);
        std.process.exit(status);
    }
    if (args.len < 7 or !std.mem.eql(u8, args[5], "--")) {
        usage();
        return error.InvalidArguments;
    }
    for ([_][]const u8{ args[1], args[2], args[6] }) |path| {
        if (!std.fs.path.isAbsolute(path)) return error.AbsolutePathRequired;
    }
    const uid = try userId(args[3]);
    const gid = try userId(args[4]);
    if (getuid() != geteuid()) return error.SetuidInstallationUnsupported;
    if (geteuid() != 0) {
        std.debug.print("v9fs: native 9P mounts need CAP_SYS_ADMIN in the initial user namespace; use --launch to ask sudo in this terminal.\n", .{});
        return error.MountPrivilegeRequired;
    }

    // Prepare everything before changing namespace or credentials.
    const options = try std.fmt.allocPrintSentinel(arena, "trans=unix,version=9p2000,cache=none,access=any,uname={d},dfltuid={d},dfltgid={d}", .{ uid, uid, gid }, 0);
    const argv = try arena.allocSentinel(?[*:0]const u8, args.len - 6, null);
    for (args[6..], 0..) |arg, i| argv[i] = arg.ptr;

    try check(unshare(0x00020000), "create private mount namespace (CAP_SYS_ADMIN required)"); // CLONE_NEWNS
    try check(mount(null, "/", null, (1 << 14) | (1 << 18), null), "make mount propagation recursively private"); // MS_REC | MS_PRIVATE
    try check(mount(args[1].ptr, args[2].ptr, "9p", 2 | 4 | 8, options.ptr), "mount 9P2000 over Unix socket (kernel 9p and 9pnet_fd support required)"); // NOSUID | NODEV | NOEXEC

    try check(setgroups(0, null), "clear supplementary groups");
    const account = std.c.getpwuid(uid) orelse return error.UserAccountNotFound;
    try check(initgroups(account.name orelse return error.UserAccountNotFound, gid), "restore user groups");
    try check(setresgid(gid, gid, gid), "drop group privileges");
    try check(setresuid(uid, uid, uid), "drop user privileges");
    // sudo can replace identity variables and PATH even with -E. Restore
    // those values only after dropping privileges.
    inline for (.{ "PATH", "HOME", "USER", "LOGNAME", "SHELL" }) |name| {
        const saved = "PARDES_V9FS_" ++ name;
        if (std.c.getenv(saved)) |value| {
            try check(setenv(name, value, 1), "restore shell environment");
            _ = unsetenv(saved);
        }
    }
    _ = execv(args[6].ptr, argv.ptr);
    // Namespace destruction releases the mount when its last process exits.
    try check(-1, "execute unprivileged command");
}