rename: script to screen

2026-03-21 23:30:47 +00:00
parent 3dfe7a8b67
commit 5586051bb1
4 changed files with 47 additions and 7 deletions
@@ -0,0 +1,399 @@
+/*
+Intially mire dependend on screen(1)
+But that had some quite a few issues
+- kernel tty line discipline: script controlled a slave bash terminal so backspaces if any
+were truncated due to output piping as the bash propmt wasn't ready. Initially we moves to
+do disable that as a pty option entirely. That broke echo so now we have a marker based
+approach which I guess would've worked in screen to.
+- Second issue, screen runs took ~200 ms while with this custom one it's 8ms for the same
+fixtures so this was kept.
+*/
+
+package script
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"os"
+	"os/exec"
+	"os/signal"
+	"syscall"
+	"time"
+
+	"github.com/creack/pty"
+	"golang.org/x/sys/unix"
+	"golang.org/x/term"
+)
+
+const (
+	defaultRows            = 24
+	defaultCols            = 80
+	terminalEOF            = byte(0x04)
+	replayReadySettleDelay = 5 * time.Millisecond
+)
+
+type RecordRequest struct {
+	Cmd       *exec.Cmd
+	Input     io.Reader
+	Output    io.Writer
+	TTY       *os.File
+	InputLog  io.Writer
+	OutputLog io.Writer
+}
+
+type ReplayRequest struct {
+	Cmd        *exec.Cmd
+	Input      []byte
+	InputReady <-chan struct{}
+	OutputLog  io.Writer
+}
+
+// Record runs a live PTY session so we can mirror interactive output while capturing stable input and output logs.
+func Record(req RecordRequest) error {
+	if req.Cmd == nil {
+		return errors.New("record session command is required")
+	}
+
+	ptmx, err := pty.StartWithSize(req.Cmd, sessionSize(req.TTY))
+	if err != nil {
+		return err
+	}
+	defer ptmx.Close()
+
+	restoreTTY, err := makeRaw(req.TTY)
+	if err != nil {
+		return err
+	}
+	defer restoreTTY()
+
+	// Keep the child PTY aligned with the real terminal so full-screen apps
+	// redraw against the size the user is actually seeing.
+	stopResize := watchResize(req.TTY, ptmx)
+	defer stopResize()
+
+	// Duplicate file-backed input so the stop path can close our copy without
+	// accidentally tearing down the caller's stdin handle.
+	input, closeInput, err := duplicateInput(req.Input)
+	if err != nil {
+		return err
+	}
+	defer closeInput()
+
+	outputDone := copyAsync(combineWriters(req.Output, req.OutputLog), ptmx)
+	inputDone, stopInput, err := copyInputAsync(combineWriters(ptmx, newInputLogWriter(req.InputLog)), input)
+	if err != nil {
+		return err
+	}
+	defer stopInput()
+
+	waitErr := req.Cmd.Wait()
+	stopInput()
+	ptmx.Close()
+
+	outputErr := <-outputDone
+	inputErr := <-inputDone
+
+	return firstErr(waitErr, outputErr, inputErr)
+}
+
+// Replay feeds recorded keystrokes back into a fresh PTY session to verify behavior against captured output.
+func Replay(req ReplayRequest) error {
+	if req.Cmd == nil {
+		return errors.New("replay session command is required")
+	}
+
+	ptmx, err := pty.StartWithSize(req.Cmd, sessionSize(nil))
+	if err != nil {
+		return err
+	}
+	defer ptmx.Close()
+
+	outputDone := copyAsync(combineWriters(req.OutputLog), ptmx)
+	processDone := make(chan struct{})
+	inputDone := copyAsyncWhenReady(ptmx, bytes.NewReader(replayInput(req.Input)), req.InputReady, processDone)
+
+	waitErr := req.Cmd.Wait()
+	close(processDone)
+	ptmx.Close()
+
+	outputErr := <-outputDone
+	inputErr := <-inputDone
+
+	return firstErr(waitErr, outputErr, inputErr)
+}
+
+// combineWriters skips nil destinations so callers can fan out conditionally without repeated nil checks.
+func combineWriters(writers ...io.Writer) io.Writer {
+	active := make([]io.Writer, 0, len(writers))
+	for _, writer := range writers {
+		if writer != nil {
+			active = append(active, writer)
+		}
+	}
+
+	switch len(active) {
+	case 0:
+		return io.Discard
+	case 1:
+		return active[0]
+	default:
+		return io.MultiWriter(active...)
+	}
+}
+
+// copyAsync moves a stream in the background so PTY input and output can progress concurrently.
+func copyAsync(dst io.Writer, src io.Reader) <-chan error {
+	done := make(chan error, 1)
+	go func() {
+		_, err := io.Copy(dst, src)
+		done <- normalizeCopyError(err)
+	}()
+	return done
+}
+
+// copyAsyncWhenReady delays replay input until the child shell is ready enough to receive it reliably.
+func copyAsyncWhenReady(dst io.Writer, src io.Reader, ready <-chan struct{}, stop <-chan struct{}) <-chan error {
+	done := make(chan error, 1)
+	go func() {
+		if ready != nil {
+			select {
+			case <-ready:
+				// Some programs signal readiness before they have finished their
+				// first prompt, so we wait briefly to avoid replaying input too early.
+				time.Sleep(replayReadySettleDelay)
+			case <-stop:
+				done <- nil
+				return
+			}
+		}
+		_, err := io.Copy(dst, src)
+		done <- normalizeCopyError(err)
+	}()
+	return done
+}
+
+// replayInput appends terminal EOF so non-interactive replays still tell the shell when scripted input is finished.
+func replayInput(input []byte) []byte {
+	if len(input) > 0 && input[len(input)-1] == terminalEOF {
+		return input
+	}
+
+	data := make([]byte, 0, len(input)+1)
+	data = append(data, input...)
+	data = append(data, terminalEOF)
+	return data
+}
+
+// copyInputAsync gives file-backed input an interruptible read loop so shutdown does not hang on blocked stdin reads.
+func copyInputAsync(dst io.Writer, src io.Reader) (<-chan error, func(), error) {
+	file, ok := src.(*os.File)
+	if !ok {
+		done := copyAsync(dst, src)
+		return done, func() {}, nil
+	}
+
+	stopReader, stopWriter, err := os.Pipe()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	done := make(chan error, 1)
+	go func() {
+		defer stopReader.Close()
+
+		buf := make([]byte, 4096)
+		// Poll lets us break out of a blocked terminal read immediately when the
+		// session ends instead of waiting for the next keystroke.
+		fds := []unix.PollFd{
+			{Fd: int32(file.Fd()), Events: unix.POLLIN | unix.POLLHUP},
+			{Fd: int32(stopReader.Fd()), Events: unix.POLLIN | unix.POLLHUP},
+		}
+
+		for {
+			_, err := unix.Poll(fds, -1)
+			if err != nil {
+				if err == unix.EINTR {
+					continue
+				}
+				done <- err
+				return
+			}
+			if fds[1].Revents != 0 {
+				done <- nil
+				return
+			}
+			if fds[0].Revents&(unix.POLLIN|unix.POLLHUP) == 0 {
+				continue
+			}
+
+			n, readErr := file.Read(buf)
+			if n > 0 {
+				if _, writeErr := dst.Write(buf[:n]); writeErr != nil {
+					done <- normalizeCopyError(writeErr)
+					return
+				}
+			}
+			if readErr != nil {
+				done <- normalizeCopyError(readErr)
+				return
+			}
+		}
+	}()
+
+	stop := func() {
+		_ = stopWriter.Close()
+		_ = file.Close()
+	}
+
+	return done, stop, nil
+}
+
+// normalizeCopyError treats expected PTY shutdown conditions as success so callers only see actionable failures.
+func normalizeCopyError(err error) error {
+	switch {
+	case err == nil:
+		return nil
+	case errors.Is(err, io.EOF):
+		return nil
+	case errors.Is(err, os.ErrClosed):
+		return nil
+	case errors.Is(err, syscall.EIO):
+		return nil
+	default:
+		return err
+	}
+}
+
+// firstErr preserves the earliest real failure because later shutdown errors are often just fallout.
+func firstErr(errs ...error) error {
+	for _, err := range errs {
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// sessionSize prefers the caller's terminal geometry so interactive programs render as if they were attached directly.
+func sessionSize(tty *os.File) *pty.Winsize {
+	if tty != nil && term.IsTerminal(int(tty.Fd())) {
+		if cols, rows, err := term.GetSize(int(tty.Fd())); err == nil {
+			return &pty.Winsize{
+				Rows: uint16(rows),
+				Cols: uint16(cols),
+			}
+		}
+	}
+
+	return &pty.Winsize{
+		Rows: defaultRows,
+		Cols: defaultCols,
+	}
+}
+
+// makeRaw disables local terminal processing so the child PTY sees the user's exact keystrokes and control bytes.
+func makeRaw(tty *os.File) (func(), error) {
+	if tty == nil || !term.IsTerminal(int(tty.Fd())) {
+		return func() {}, nil
+	}
+
+	// Raw mode lets the child process receive keystrokes and control sequences
+	// directly instead of having the local terminal preprocess them first.
+	state, err := term.MakeRaw(int(tty.Fd()))
+	if err != nil {
+		return nil, err
+	}
+
+	return func() {
+		_ = term.Restore(int(tty.Fd()), state)
+	}, nil
+}
+
+// watchResize forwards host terminal resizes so curses-style apps redraw against the current viewport.
+func watchResize(tty *os.File, ptmx *os.File) func() {
+	if tty == nil || !term.IsTerminal(int(tty.Fd())) {
+		return func() {}
+	}
+
+	applySize := func() {
+		_ = pty.Setsize(ptmx, sessionSize(tty))
+	}
+	applySize()
+
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, syscall.SIGWINCH)
+
+	done := make(chan struct{})
+	go func() {
+		for {
+			select {
+			case <-signals:
+				applySize()
+			case <-done:
+				return
+			}
+		}
+	}()
+
+	return func() {
+		signal.Stop(signals)
+		close(done)
+	}
+}
+
+// duplicateInput gives the recorder ownership of file-backed input without mutating the caller's descriptor lifecycle.
+func duplicateInput(input io.Reader) (io.Reader, func(), error) {
+	if input == nil {
+		return bytes.NewReader(nil), func() {}, nil
+	}
+
+	file, ok := input.(*os.File)
+	if !ok {
+		return input, func() {}, nil
+	}
+
+	fd, err := syscall.Dup(int(file.Fd()))
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// The recorder may need to close its input to stop cleanly, but that should
+	// only affect the duplicated descriptor it owns.
+	dup := os.NewFile(uintptr(fd), file.Name())
+	return dup, func() {
+		_ = dup.Close()
+	}, nil
+}
+
+type inputLogWriter struct {
+	dst io.Writer
+}
+
+// newInputLogWriter hides optional logging behind a writer so the input path stays linear.
+func newInputLogWriter(dst io.Writer) io.Writer {
+	if dst == nil {
+		return io.Discard
+	}
+	return inputLogWriter{dst: dst}
+}
+
+// Write normalizes terminal line endings so recorded input fixtures are easy to diff and reuse.
+func (w inputLogWriter) Write(p []byte) (int, error) {
+	normalized := make([]byte, len(p))
+	for i, b := range p {
+		if b == '\r' {
+			// Logs are easier to diff and replay reasoning against when Enter is
+			// stored as a regular newline instead of terminal carriage returns.
+			normalized[i] = '\n'
+			continue
+		}
+		normalized[i] = b
+	}
+
+	if _, err := w.dst.Write(normalized); err != nil {
+		return 0, err
+	}
+
+	return len(p), nil
+}
@@ -0,0 +1,95 @@
+package script
+
+import (
+	"bytes"
+	"os"
+	"os/exec"
+	"testing"
+	"time"
+)
+
+func TestRecordCopiesInputAndOutput(t *testing.T) {
+	reader, writer, err := os.Pipe()
+	if err != nil {
+		t.Fatalf("os.Pipe() error = %v", err)
+	}
+	t.Cleanup(func() {
+		_ = reader.Close()
+		_ = writer.Close()
+	})
+
+	go func() {
+		_, _ = writer.Write([]byte("hello\n"))
+		_ = writer.Close()
+	}()
+
+	cmd := exec.Command("sh", "-c", `printf 'ready\n'; read line; printf 'line:%s\n' "$line"`)
+
+	var live bytes.Buffer
+	var inputLog bytes.Buffer
+	var outputLog bytes.Buffer
+	if err := Record(RecordRequest{
+		Cmd:       cmd,
+		Input:     reader,
+		Output:    &live,
+		InputLog:  &inputLog,
+		OutputLog: &outputLog,
+	}); err != nil {
+		t.Fatalf("Record() error = %v", err)
+	}
+
+	if got := inputLog.String(); got != "hello\n" {
+		t.Fatalf("input log = %q, want %q", got, "hello\n")
+	}
+	for _, got := range []string{live.String(), outputLog.String()} {
+		if !bytes.Contains([]byte(got), []byte("ready")) || !bytes.Contains([]byte(got), []byte("line:hello")) {
+			t.Fatalf("output = %q, want ready + line:hello", got)
+		}
+	}
+}
+
+func TestReplayCapturesOutput(t *testing.T) {
+	cmd := exec.Command("sh", "-c", `read line; printf '__MIRE_E2E_BEGIN__\nline:%s\n' "$line"`)
+
+	var outputLog bytes.Buffer
+	if err := Replay(ReplayRequest{
+		Cmd:       cmd,
+		Input:     []byte("hello\n"),
+		OutputLog: &outputLog,
+	}); err != nil {
+		t.Fatalf("Replay() error = %v", err)
+	}
+
+	got := outputLog.String()
+	if !bytes.Contains([]byte(got), []byte("__MIRE_E2E_BEGIN__")) || !bytes.Contains([]byte(got), []byte("line:hello")) {
+		t.Fatalf("output log = %q, want marker + line:hello", got)
+	}
+}
+
+func TestReplayWaitsForInputReadySignal(t *testing.T) {
+	cmd := exec.Command("sh", "-c", `IFS= read -r line; printf 'line:%s\n' "$line"`)
+	ready := make(chan struct{})
+
+	var outputLog bytes.Buffer
+	start := time.Now()
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		close(ready)
+	}()
+
+	if err := Replay(ReplayRequest{
+		Cmd:        cmd,
+		Input:      []byte("hello\n"),
+		InputReady: ready,
+		OutputLog:  &outputLog,
+	}); err != nil {
+		t.Fatalf("Replay() error = %v", err)
+	}
+
+	if elapsed := time.Since(start); elapsed < 100*time.Millisecond {
+		t.Fatalf("Replay() elapsed = %v, want >= %v", elapsed, 100*time.Millisecond)
+	}
+	if !bytes.Contains(outputLog.Bytes(), []byte("line:hello")) {
+		t.Fatalf("output log = %q, want line:hello", outputLog.String())
+	}
+}