testscript: add ttyin/ttyout commands

Author: FiloSottile

testscript/cmd.go

@@ -13,10 +13,12 @@ import (
 	"os/exec"
 	"path/filepath"
 	"regexp"
+	"runtime"
 	"strconv"
 	"strings"
 
 	"github.com/rogpeppe/go-internal/diff"
+	"github.com/rogpeppe/go-internal/testscript/internal/pty"
 	"github.com/rogpeppe/go-internal/txtar"
 )
 
@@ -41,6 +43,8 @@ var scriptCmds = map[string]func(*TestScript, bool, []string){
 	"stderr":   (*TestScript).cmdStderr,
 	"stdin":    (*TestScript).cmdStdin,
 	"stdout":   (*TestScript).cmdStdout,
+	"ttyin":    (*TestScript).cmdTtyin,
+	"ttyout":   (*TestScript).cmdTtyout,
 	"stop":     (*TestScript).cmdStop,
 	"symlink":  (*TestScript).cmdSymlink,
 	"unix2dos": (*TestScript).cmdUNIX2DOS,
@@ -178,6 +182,10 @@ func (ts *TestScript) cmdCp(neg bool, args []string) {
 			src = arg
 			data = []byte(ts.stderr)
 			mode = 0o666
+		case "ttyout":
+			src = arg
+			data = []byte(ts.ttyout)
+			mode = 0o666
 		default:
 			src = ts.MkAbs(arg)
 			info, err := os.Stat(src)
@@ -382,6 +390,9 @@ func (ts *TestScript) cmdStdin(neg bool, args []string) {
 	if len(args) != 1 {
 		ts.Fatalf("usage: stdin filename")
 	}
+	if ts.stdinPty {
+		ts.Fatalf("conflicting use of 'stdin' and 'ttyin -stdin'")
+	}
 	ts.stdin = ts.ReadFile(args[0])
 }
 
@@ -401,6 +412,37 @@ func (ts *TestScript) cmdGrep(neg bool, args []string) {
 	scriptMatch(ts, neg, args, "", "grep")
 }
 
+func (ts *TestScript) cmdTtyin(neg bool, args []string) {
+	if !pty.Supported {
+		ts.Fatalf("unsupported: ttyin on %s", runtime.GOOS)
+	}
+	if neg {
+		ts.Fatalf("unsupported: ! ttyin")
+	}
+	switch len(args) {
+	case 1:
+		ts.ttyin = ts.ReadFile(args[0])
+	case 2:
+		if args[0] != "-stdin" {
+			ts.Fatalf("usage: ttyin [-stdin] filename")
+		}
+		if ts.stdin != "" {
+			ts.Fatalf("conflicting use of 'stdin' and 'ttyin -stdin'")
+		}
+		ts.stdinPty = true
+		ts.ttyin = ts.ReadFile(args[1])
+	default:
+		ts.Fatalf("usage: ttyin [-stdin] filename")
+	}
+	if ts.ttyin == "" {
+		ts.Fatalf("tty input file is empty")
+	}
+}
+
+func (ts *TestScript) cmdTtyout(neg bool, args []string) {
+	scriptMatch(ts, neg, args, ts.ttyout, "ttyout")
+}
+
 // stop stops execution of the test (marking it passed).
 func (ts *TestScript) cmdStop(neg bool, args []string) {
 	if neg {

testscript/doc.go

@@ -205,6 +205,16 @@ The predefined commands are:
     Apply the grep command (see above) to the standard output
     from the most recent exec or wait command.
 
+  - ttyin [-stdin] file
+    Attach the next exec command to a controlling pseudo-terminal, and use the
+    contents of the given file as the raw terminal input. If -stdin is specified,
+    also attach the terminal to standard input.
+    Note that this does not attach the terminal to standard output/error.
+
+  - [!] ttyout [-count=N] pattern
+    Apply the grep command (see above) to the raw controlling terminal output
+    from the most recent exec command.
+
   - stop [message]
     Stop the test early (marking it as passing), including the message if given.
 

testscript/envvarname.go

@@ -0,0 +1,62 @@
+//go:build linux || darwin
+// +build linux darwin
+
+package pty
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"syscall"
+)
+
+const Supported = true
+
+func SetCtty(cmd *exec.Cmd, tty *os.File) {
+	cmd.SysProcAttr = &syscall.SysProcAttr{
+		Setctty: true,
+		Setsid:  true,
+		Ctty:    3,
+	}
+	cmd.ExtraFiles = []*os.File{tty}
+}
+
+func Open() (pty, tty *os.File, err error) {
+	p, err := os.OpenFile("/dev/ptmx", os.O_RDWR, 0)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to open pty multiplexer: %v", err)
+	}
+	defer func() {
+		if err != nil {
+			p.Close()
+		}
+	}()
+
+	name, err := ptyName(p)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to obtain tty name: %v", err)
+	}
+
+	if err := ptyGrant(p); err != nil {
+		return nil, nil, fmt.Errorf("failed to grant pty: %v", err)
+	}
+
+	if err := ptyUnlock(p); err != nil {
+		return nil, nil, fmt.Errorf("failed to unlock pty: %v", err)
+	}
+
+	t, err := os.OpenFile(name, os.O_RDWR|syscall.O_NOCTTY, 0)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to open TTY: %v", err)
+	}
+
+	return p, t, nil
+}
+
+func ioctl(f *os.File, name string, cmd, ptr uintptr) error {
+	_, _, err := syscall.Syscall(syscall.SYS_IOCTL, f.Fd(), cmd, ptr)
+	if err != 0 {
+		return fmt.Errorf("%s ioctl failed: %v", name, err)
+	}
+	return nil
+}

testscript/envvarname_windows.go

@@ -0,0 +1,32 @@
+package pty
+
+import (
+	"bytes"
+	"os"
+	"syscall"
+	"unsafe"
+)
+
+func ptyName(f *os.File) (string, error) {
+	// Parameter length is encoded in the low 13 bits of the top word.
+	// See https://github.com/apple/darwin-xnu/blob/2ff845c2e0/bsd/sys/ioccom.h#L69-L77
+	const IOCPARM_MASK = 0x1fff
+	const TIOCPTYGNAME_PARM_LEN = (syscall.TIOCPTYGNAME >> 16) & IOCPARM_MASK
+	out := make([]byte, TIOCPTYGNAME_PARM_LEN)
+
+	err := ioctl(f, "TIOCPTYGNAME", syscall.TIOCPTYGNAME, uintptr(unsafe.Pointer(&out[0])))
+	if err != nil {
+		return "", err
+	}
+
+	i := bytes.IndexByte(out, 0x00)
+	return string(out[:i]), nil
+}
+
+func ptyGrant(f *os.File) error {
+	return ioctl(f, "TIOCPTYGRANT", syscall.TIOCPTYGRANT, 0)
+}
+
+func ptyUnlock(f *os.File) error {
+	return ioctl(f, "TIOCPTYUNLK", syscall.TIOCPTYUNLK, 0)
+}

testscript/internal/pty/pty.go

@@ -0,0 +1,26 @@
+package pty
+
+import (
+	"os"
+	"strconv"
+	"syscall"
+	"unsafe"
+)
+
+func ptyName(f *os.File) (string, error) {
+	var out uint
+	err := ioctl(f, "TIOCGPTN", syscall.TIOCGPTN, uintptr(unsafe.Pointer(&out)))
+	if err != nil {
+		return "", err
+	}
+	return "/dev/pts/" + strconv.Itoa(int(out)), nil
+}
+
+func ptyGrant(f *os.File) error {
+	return nil
+}
+
+func ptyUnlock(f *os.File) error {
+	var zero int
+	return ioctl(f, "TIOCSPTLCK", syscall.TIOCSPTLCK, uintptr(unsafe.Pointer(&zero)))
+}

testscript/internal/pty/pty_darwin.go

@@ -0,0 +1,21 @@
+//go:build !linux && !darwin
+// +build !linux,!darwin
+
+package pty
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"runtime"
+)
+
+const Supported = false
+
+func SetCtty(cmd *exec.Cmd, tty *os.File) error {
+	panic("SetCtty called on unsupported platform")
+}
+
+func Open() (pty, tty *os.File, err error) {
+	return nil, nil, fmt.Errorf("pty unsupported on %s", runtime.GOOS)
+}

testscript/internal/pty/pty_linux.go

@@ -0,0 +1,10 @@
+[!linux] [!darwin] skip
+
+ttyin secretwords.txt
+terminalprompt
+ttyout 'magic words'
+! stderr .
+! stdout .
+
+-- secretwords.txt --
+SQUEAMISHOSSIFRAGE

testscript/internal/pty/pty_unsupported.go

@@ -34,6 +34,7 @@ import (
 	"github.com/rogpeppe/go-internal/internal/os/execpath"
 	"github.com/rogpeppe/go-internal/par"
 	"github.com/rogpeppe/go-internal/testenv"
+	"github.com/rogpeppe/go-internal/testscript/internal/pty"
 	"github.com/rogpeppe/go-internal/txtar"
 )
 
@@ -100,6 +101,13 @@ func (e *Env) Getenv(key string) string {
 	return ""
 }
 
+func envvarname(k string) string {
+	if runtime.GOOS == "windows" {
+		return strings.ToLower(k)
+	}
+	return k
+}
+
 // Setenv sets the value of the environment variable named by the key. It
 // panics if key is invalid.
 func (e *Env) Setenv(key, value string) {
@@ -357,6 +365,9 @@ type TestScript struct {
 	stdin         string                      // standard input to next 'go' command; set by 'stdin' command.
 	stdout        string                      // standard output from last 'go' command; for 'stdout' command
 	stderr        string                      // standard error from last 'go' command; for 'stderr' command
+	ttyin         string                      // terminal input; set by 'ttyin' command
+	stdinPty      bool                        // connect pty to standard input; set by 'ttyin -stdin' command
+	ttyout        string                      // terminal output; for 'ttyout' command
 	stopped       bool                        // test wants to stop early
 	start         time.Time                   // time phase started
 	background    []backgroundCmd             // backgrounded 'exec' and 'go' commands
@@ -940,16 +951,49 @@ func (ts *TestScript) exec(command string, args ...string) (stdout, stderr strin
 	var stdoutBuf, stderrBuf strings.Builder
 	cmd.Stdout = &stdoutBuf
 	cmd.Stderr = &stderrBuf
+	if ts.ttyin != "" {
+		ctrl, tty, err := pty.Open()
+		if err != nil {
+			return "", "", err
+		}
+		doneR, doneW := make(chan struct{}), make(chan struct{})
+		var ptyBuf strings.Builder
+		go func() {
+			io.Copy(ctrl, strings.NewReader(ts.ttyin))
+			ctrl.Write([]byte{4 /* EOT */})
+			close(doneW)
+		}()
+		go func() {
+			io.Copy(&ptyBuf, ctrl)
+			close(doneR)
+		}()
+		defer func() {
+			tty.Close()
+			ctrl.Close()
+			<-doneR
+			<-doneW
+			ts.ttyin = ""
+			ts.ttyout = ptyBuf.String()
+		}()
+		pty.SetCtty(cmd, tty)
+		if ts.stdinPty {
+			cmd.Stdin = tty
+		}
+	}
 	if err = cmd.Start(); err == nil {
 		err = waitOrStop(ts.ctxt, cmd, ts.gracePeriod)
 	}
 	ts.stdin = ""
+	ts.stdinPty = false
 	return stdoutBuf.String(), stderrBuf.String(), err
 }
 
 // execBackground starts the given command line (an actual subprocess, not simulated)
 // in ts.cd with environment ts.env.
 func (ts *TestScript) execBackground(command string, args ...string) (*exec.Cmd, error) {
+	if ts.ttyin != "" {
+		return nil, errors.New("ttyin is not supported by background commands")
+	}
 	cmd, err := ts.buildExecCmd(command, args...)
 	if err != nil {
 		return nil, err
@@ -1126,6 +1170,8 @@ func (ts *TestScript) ReadFile(file string) string {
 		return ts.stdout
 	case "stderr":
 		return ts.stderr
+	case "ttyout":
+		return ts.ttyout
 	default:
 		file = ts.MkAbs(file)
 		data, err := ioutil.ReadFile(file)