github.com/lmorg/murex@v0.0.0-20240217211045-e081c89cd4ef/builtins/pipes/psuedotty/pty_unix.go (about)

     1  //go:build !windows && !js && !plan9 && !no_pty
     2  // +build !windows,!js,!plan9,!no_pty
     3  
     4  package psuedotty
     5  
     6  import (
     7  	"context"
     8  	"fmt"
     9  	"io"
    10  	"os"
    11  	"sync/atomic"
    12  	"time"
    13  
    14  	"github.com/creack/pty"
    15  	"github.com/lmorg/murex/builtins/pipes/streams"
    16  	"github.com/lmorg/murex/config"
    17  	"github.com/lmorg/murex/debug"
    18  	"github.com/lmorg/murex/lang/stdio"
    19  	"github.com/lmorg/murex/utils"
    20  	"github.com/lmorg/murex/utils/readline"
    21  )
    22  
    23  func init() {
    24  	stdio.RegisterPipe("pty", registerPipe)
    25  }
    26  
    27  func registerPipe(_ string) (stdio.Io, error) {
    28  	return NewPTY(80, 25)
    29  }
    30  
    31  type PTY struct {
    32  	in         *os.File
    33  	replica    *os.File
    34  	out        *streams.Stdin
    35  	dependents int32
    36  }
    37  
    38  func NewPTY(width, height int) (*PTY, error) {
    39  	primary, replica, err := pty.Open()
    40  	if err != nil {
    41  		return nil, fmt.Errorf("unable to open pty: %s", err.Error())
    42  	}
    43  
    44  	size := pty.Winsize{
    45  		Cols: uint16(width),
    46  		Rows: uint16(height),
    47  	}
    48  
    49  	err = pty.Setsize(primary, &size)
    50  	if err != nil {
    51  		return nil, fmt.Errorf("unable to set pty size: %s", err.Error())
    52  	}
    53  
    54  	_, err = readline.MakeRaw(int(primary.Fd()))
    55  	if err != nil {
    56  		return nil, fmt.Errorf("unable to set pty state: %s", err.Error())
    57  	}
    58  
    59  	p := new(PTY)
    60  	p.in = primary
    61  	p.replica = replica
    62  	p.out = streams.NewStdin()
    63  	p.out.Open()
    64  
    65  	go func() {
    66  		i, err := io.Copy(p.out, p.replica)
    67  		if err != nil && debug.Enabled {
    68  			os.Stderr.WriteString(fmt.Sprintf("!!! read failed from PTY after %d bytes: %v !!!", i, err))
    69  		}
    70  	}()
    71  
    72  	return p, nil
    73  }
    74  
    75  // GetDataType returns the murex data type for the stream.Io interface
    76  func (p *PTY) GetDataType() (dt string) { return p.out.GetDataType() }
    77  
    78  // SetDataType defines the murex data type for the stream.Io interface
    79  func (p *PTY) SetDataType(dt string) { p.out.SetDataType(dt) }
    80  
    81  // Stats provides real time stream stats. Useful for progress bars etc.
    82  func (p *PTY) Stats() (uint64, uint64) { return p.out.Stats() }
    83  
    84  // IsTTY returns true because the PTY stream is a pseudo-TTY
    85  func (p *PTY) IsTTY() bool { return true }
    86  
    87  // File returns the os.File struct for the stream.Io interface if a TTY
    88  func (p *PTY) File() *os.File { return p.in }
    89  
    90  // Open the stream.Io interface for another dependant
    91  func (p *PTY) Open() {
    92  	atomic.AddInt32(&p.dependents, 1)
    93  	p.out.Open()
    94  }
    95  
    96  // Close the stream.Io interface
    97  func (p *PTY) Close() {
    98  	i := atomic.AddInt32(&p.dependents, -1)
    99  	if i < 0 {
   100  		panic("More closed dependents than open")
   101  	}
   102  	p.out.Close()
   103  	if i == 0 {
   104  		go p.close()
   105  	}
   106  }
   107  
   108  func (p *PTY) close() {
   109  	defer p.out.ForceClose()
   110  
   111  	for {
   112  		time.Sleep(1 * time.Second)
   113  		w, r := p.out.Stats()
   114  		if r >= w {
   115  			err := p.in.Close()
   116  			if err != nil {
   117  				panic(err)
   118  			}
   119  			err = p.replica.Close()
   120  			if err != nil {
   121  				panic(err)
   122  			}
   123  			return
   124  		}
   125  	}
   126  }
   127  
   128  // ForceClose forces the stream.Io interface to close. This should only be called by a STDIN reader
   129  func (p *PTY) ForceClose() {
   130  	p.in.Close()
   131  	p.replica.Close()
   132  	p.out.ForceClose()
   133  }
   134  
   135  func (p *PTY) Read(b []byte) (int, error) {
   136  	return p.out.Read(b)
   137  }
   138  
   139  func (p *PTY) ReadLine(callback func([]byte)) error {
   140  	return p.out.ReadLine(callback)
   141  }
   142  
   143  func (p *PTY) ReadArray(ctx context.Context, callback func([]byte)) error {
   144  	return p.out.ReadArray(ctx, callback)
   145  }
   146  
   147  func (p *PTY) ReadArrayWithType(ctx context.Context, callback func(interface{}, string)) error {
   148  	return p.out.ReadArrayWithType(ctx, callback)
   149  }
   150  
   151  func (p *PTY) ReadMap(conf *config.Config, callback func(*stdio.Map)) error {
   152  	return p.out.ReadMap(conf, callback)
   153  }
   154  
   155  func (p *PTY) ReadAll() ([]byte, error) {
   156  	b, err := p.out.ReadAll()
   157  	_ = p.in.Close()
   158  	p.out.Close()
   159  	return b, err
   160  }
   161  
   162  func (p *PTY) Write(b []byte) (int, error) {
   163  	i, err := p.in.Write(b)
   164  	return i, err
   165  }
   166  
   167  func (p *PTY) Writeln(b []byte) (int, error) {
   168  	slice := append(b, utils.NewLineByte...)
   169  	return p.in.Write(slice)
   170  }
   171  
   172  func (p *PTY) WriteArray(dataType string) (stdio.ArrayWriter, error) {
   173  	return stdio.WriteArray(p, dataType)
   174  }
   175  
   176  func (p *PTY) WriteTo(w io.Writer) (int64, error) {
   177  	return stdio.WriteTo(p, w)
   178  }