github.com/u-root/u-root@v7.0.1-0.20200915234505-ad7babab0a8e+incompatible/pkg/pty/pty.go (about)

     1  // Copyright 2015-2020 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package pty provides basic pty support.
     6  // It implments much of exec.Command
     7  // but the Start() function starts two goroutines that relay the
     8  // data for Stdin, Stdout, and Stdout such that proper kernel pty
     9  // processing is done. We did not simply embed an exec.Command
    10  // as we can no guarantee that we can implement all aspects of it
    11  // for all time to come.
    12  package pty
    13  
    14  import (
    15  	"fmt"
    16  	"io"
    17  	"os"
    18  	"os/exec"
    19  
    20  	"github.com/u-root/u-root/pkg/termios"
    21  )
    22  
    23  // Pty contains all the bits and pieces needed to start and control
    24  // a process at the other end of a pty, as well as whatever is needed
    25  // to control the pty, and restore modes on exit.
    26  type Pty struct {
    27  	C        *exec.Cmd
    28  	Ptm      *os.File
    29  	Pts      *os.File
    30  	Sname    string
    31  	Kid      int
    32  	TTY      *termios.TTYIO
    33  	WS       *termios.Winsize
    34  	Restorer *termios.Termios
    35  }
    36  
    37  var sys = func(p *Pty) {}
    38  
    39  // Command sets up an exec.Command at the remote end of a Pty.
    40  func (p *Pty) Command(cmd string, args ...string) {
    41  	p.C = exec.Command(cmd, args...)
    42  	p.C.Stdin, p.C.Stdout, p.C.Stderr = p.Pts, p.Pts, p.Pts
    43  	sys(p)
    44  }
    45  
    46  // Start starts Command attached to a Pty. It sets
    47  // window size and other variables as needed. It does not
    48  // block.
    49  func (p *Pty) Start() error {
    50  	tty, err := termios.New()
    51  	if err != nil {
    52  		return err
    53  	}
    54  
    55  	if p.WS, err = tty.GetWinSize(); err != nil {
    56  		return err
    57  	}
    58  
    59  	if p.Restorer, err = tty.Raw(); err != nil {
    60  		return err
    61  	}
    62  
    63  	if err := p.C.Start(); err != nil {
    64  		tty.Set(p.Restorer)
    65  		return err
    66  	}
    67  	p.Kid = p.C.Process.Pid
    68  
    69  	// We make a good faith effort to set the
    70  	// WinSize of the Pts, but it's not a deal breaker
    71  	// if we can't do it.
    72  	if err := termios.SetWinSize(p.Pts.Fd(), p.WS); err != nil {
    73  		fmt.Fprintf(p.C.Stderr, "SetWinSize of Pts: %v", err)
    74  	}
    75  
    76  	return nil
    77  }
    78  
    79  // Run runs a Command attached to a Pty, waiting for completion and
    80  // managing stdio. It uses Wait to restore tty modes when it is done.
    81  func (p *Pty) Run() error {
    82  	if err := p.Start(); err != nil {
    83  		return err
    84  	}
    85  
    86  	go io.Copy(p.TTY, p.Ptm)
    87  
    88  	// The 1 byte for IO may seem weird, but ptys are for human interaction
    89  	// and, let's face it, we don't all type fast.
    90  	go func() {
    91  		var data [1]byte
    92  		for {
    93  			if _, err := p.TTY.Read(data[:]); err != nil {
    94  				return
    95  			}
    96  			// Log the error but it may be transient.
    97  			if _, err := p.Ptm.Write(data[:]); err != nil {
    98  				fmt.Fprintf(p.C.Stderr, "Error writing input to ptm: %v: give up\n", err)
    99  			}
   100  		}
   101  	}()
   102  	return p.Wait()
   103  }
   104  
   105  // Wait waits for a previously started command to finish, and restores the
   106  // tty mode when it is done.
   107  func (p *Pty) Wait() error {
   108  	defer p.TTY.Set(p.Restorer)
   109  	return p.C.Wait()
   110  }