src/pty.zig

Ref: main   Size: 5.6 KiB

const std = @import("std");

const c = @cImport({
    @cInclude("pty.h");
    @cInclude("termios.h");
    @cInclude("unistd.h");
    @cInclude("sys/ioctl.h");
    @cInclude("fcntl.h");
    @cInclude("sys/wait.h");
    @cInclude("stdlib.h");
});

pub const Pty = struct {
    master_fd: std.posix.fd_t,
    child_pid: std.posix.pid_t,
    child_reaped: bool = false,

    pub const SpawnOptions = struct {
        cols: u16,
        rows: u16,
        shell: [:0]const u8,
        shell_args: ?[]const [:0]const u8 = null,
    };

    pub fn spawn(opts: SpawnOptions) !Pty {
        var master: c_int = undefined;
        var winsize = c.struct_winsize{
            .ws_row = opts.rows,
            .ws_col = opts.cols,
            .ws_xpixel = 0,
            .ws_ypixel = 0,
        };

        const pid = c.forkpty(&master, null, null, &winsize);
        if (pid < 0) return error.ForkptyFailed;

        if (pid == 0) {
            // Child process
            _ = c.setenv("TERM", "xterm-256color", 1);

            if (opts.shell_args) |args| {
                std.debug.assert(args.len < 15); // argv[0] = shell, must fit in 16-slot buffer
                var argv_buf: [16:null]?[*:0]const u8 = .{null} ** 16;
                argv_buf[0] = opts.shell.ptr;
                for (args, 1..) |arg, i| {
                    argv_buf[i] = arg.ptr;
                }
                std.posix.execveZ(opts.shell.ptr, &argv_buf, std.c.environ) catch {};
            } else {
                var argv = [_:null]?[*:0]const u8{ opts.shell.ptr, null };
                std.posix.execveZ(opts.shell.ptr, &argv, std.c.environ) catch {};
            }
            std.process.exit(1);
        }

        // Parent: set master fd non-blocking
        const flags = try std.posix.fcntl(master, std.posix.F.GETFL, 0);
        const nonblock_bit: usize = @as(u32, @bitCast(std.posix.O{ .NONBLOCK = true }));
        _ = try std.posix.fcntl(master, std.posix.F.SETFL, flags | nonblock_bit);

        return .{
            .master_fd = master,
            .child_pid = pid,
        };
    }

    pub fn read(self: *Pty, buf: []u8) !usize {
        return std.posix.read(self.master_fd, buf);
    }

    pub fn write(self: *Pty, data: []const u8) !usize {
        return std.posix.write(self.master_fd, data);
    }

    pub fn resize(self: *Pty, cols: u16, rows: u16) !void {
        var ws = c.struct_winsize{
            .ws_row = rows,
            .ws_col = cols,
            .ws_xpixel = 0,
            .ws_ypixel = 0,
        };
        if (c.ioctl(self.master_fd, c.TIOCSWINSZ, &ws) < 0) {
            return error.IoctlFailed;
        }
    }

    pub fn isChildAlive(self: *Pty) bool {
        if (self.child_reaped) return false;
        var status: c_int = 0;
        const rc = c.waitpid(self.child_pid, &status, c.WNOHANG);
        if (rc == 0) return true;
        if (rc == self.child_pid) {
            self.child_reaped = true;
            return false;
        }
        return true;
    }

    pub fn deinit(self: *Pty) void {
        std.posix.close(self.master_fd);
        if (!self.child_reaped) {
            _ = std.c.kill(self.child_pid, std.c.SIG.TERM);
            _ = std.c.waitpid(self.child_pid, null, 0);
        }
    }
};

test "Pty.write and read echoes through shell" {
    var pty = try Pty.spawn(.{
        .cols = 80,
        .rows = 24,
        .shell = "/bin/sh",
    });
    defer pty.deinit();

    // Give shell a moment to start
    std.Thread.sleep(100 * std.time.ns_per_ms);

    // Write a command
    _ = try pty.write("echo hello\n");

    // Drain output for up to 1 second
    var buf: [4096]u8 = undefined;
    var seen_hello = false;
    const deadline = std.time.nanoTimestamp() + 1 * std.time.ns_per_s;
    while (std.time.nanoTimestamp() < deadline) {
        const n = pty.read(&buf) catch |err| switch (err) {
            error.WouldBlock => {
                std.Thread.sleep(10 * std.time.ns_per_ms);
                continue;
            },
            else => return err,
        };
        if (n == 0) break;
        if (std.mem.indexOf(u8, buf[0..n], "hello") != null) {
            seen_hello = true;
            break;
        }
    }
    try std.testing.expect(seen_hello);
}

test "Pty.spawn launches /bin/sh and returns valid fd" {
    var pty = try Pty.spawn(.{
        .cols = 80,
        .rows = 24,
        .shell = "/bin/sh",
    });
    defer pty.deinit();
    try std.testing.expect(pty.master_fd >= 0);
    try std.testing.expect(pty.child_pid > 0);
}

test "Pty.resize sets winsize via ioctl" {
    var pty = try Pty.spawn(.{
        .cols = 80,
        .rows = 24,
        .shell = "/bin/sh",
    });
    defer pty.deinit();

    try pty.resize(120, 40);

    var ws: c.struct_winsize = undefined;
    const rc = c.ioctl(pty.master_fd, c.TIOCGWINSZ, &ws);
    try std.testing.expectEqual(@as(c_int, 0), rc);
    try std.testing.expectEqual(@as(c_ushort, 120), ws.ws_col);
    try std.testing.expectEqual(@as(c_ushort, 40), ws.ws_row);
}

test "Pty.isChildAlive returns true while shell runs, false after exit" {
    var pty = try Pty.spawn(.{
        .cols = 80,
        .rows = 24,
        .shell = "/bin/sh",
    });
    defer pty.deinit();

    try std.testing.expect(pty.isChildAlive());

    // Tell the shell to exit
    _ = try pty.write("exit\n");

    // Wait up to 1 second for it to actually exit
    const deadline = std.time.nanoTimestamp() + 1 * std.time.ns_per_s;
    var saw_dead = false;
    while (std.time.nanoTimestamp() < deadline) {
        if (!pty.isChildAlive()) {
            saw_dead = true;
            break;
        }
        std.Thread.sleep(20 * std.time.ns_per_ms);
    }
    try std.testing.expect(saw_dead);
}