rename: script to screen
This commit is contained in:
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user