/
shell_session.go
112 lines (93 loc) · 2.27 KB
/
shell_session.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
package runner
import (
"io"
"os"
"os/exec"
"sync"
"github.com/creack/pty"
"github.com/pkg/errors"
"github.com/stateful/runme/internal/identity"
xpty "github.com/stateful/runme/internal/pty"
"golang.org/x/term"
)
type ShellSession struct {
id string
cmd *exec.Cmd
ptmx *os.File
cancelResize xpty.CancelFn
stdinOldState *term.State
done chan struct{}
mu sync.Mutex
err error
}
func NewShellSession(command string) (*ShellSession, error) {
id := identity.GenerateID()
cmd := exec.Command(command)
cmd.Env = append(os.Environ(), "RUNMESHELL="+id)
ptmx, err := pty.Start(cmd)
if err != nil {
return nil, errors.WithStack(err)
}
cancelResize := xpty.ResizeOnSig(ptmx)
// Set stdin in raw mode.
stdinOldState, err := term.MakeRaw(int(os.Stdin.Fd()))
if err != nil {
return nil, errors.Wrap(err, "failed to put stdin in raw mode")
}
s := &ShellSession{
id: id,
cmd: cmd,
ptmx: ptmx,
cancelResize: cancelResize,
stdinOldState: stdinOldState,
done: make(chan struct{}),
}
go func() {
defer close(s.done)
_, err := io.Copy(os.Stdout, ptmx)
if err != nil && errors.Is(err, io.EOF) {
err = nil
}
s.setErrorf(err, "failed to copy ptmx to stdout")
}()
go func() {
// TODO: cancel this goroutine when ptmx is closed
_, err := io.Copy(ptmx, os.Stdin)
s.setErrorf(err, "failed to copy stdin to ptmx")
}()
return s, nil
}
func (s *ShellSession) ID() string {
return s.id
}
func (s *ShellSession) Close() error {
s.cancelResize()
if err := term.Restore(int(os.Stdin.Fd()), s.stdinOldState); err != nil {
return errors.WithStack(err)
}
if err := s.cmd.Process.Kill(); err != nil {
return errors.Wrap(err, "failed to kill command")
}
<-s.done
return nil
}
func (s *ShellSession) Done() <-chan struct{} {
return s.done
}
func (s *ShellSession) Err() error {
s.mu.Lock()
defer s.mu.Unlock()
return s.err
}
func (s *ShellSession) Send(data []byte) error {
_, err := s.ptmx.Write(data)
return errors.Wrap(err, "failed to write data to ptmx")
}
func (s *ShellSession) setErrorf(err error, msg string, args ...interface{}) {
if s.err != nil || err == nil {
return
}
s.mu.Lock()
s.err = errors.Wrapf(err, msg, args...)
s.mu.Unlock()
}