|
| 1 | +// Copyright 2018 Netflix, Inc. |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +package expect |
| 16 | + |
| 17 | +import ( |
| 18 | + "bufio" |
| 19 | + "fmt" |
| 20 | + "io" |
| 21 | + "io/ioutil" |
| 22 | + "log" |
| 23 | + "os" |
| 24 | + "time" |
| 25 | + "unicode/utf8" |
| 26 | + |
| 27 | + "github.com/coder/coder/expect/pty" |
| 28 | +) |
| 29 | + |
| 30 | +// Console is an interface to automate input and output for interactive |
| 31 | +// applications. Console can block until a specified output is received and send |
| 32 | +// input back on it's tty. Console can also multiplex other sources of input |
| 33 | +// and multiplex its output to other writers. |
| 34 | +type Console struct { |
| 35 | + opts ConsoleOpts |
| 36 | + pty pty.Pty |
| 37 | + passthroughPipe *PassthroughPipe |
| 38 | + runeReader *bufio.Reader |
| 39 | + closers []io.Closer |
| 40 | +} |
| 41 | + |
| 42 | +// ConsoleOpt allows setting Console options. |
| 43 | +type ConsoleOpt func(*ConsoleOpts) error |
| 44 | + |
| 45 | +// ConsoleOpts provides additional options on creating a Console. |
| 46 | +type ConsoleOpts struct { |
| 47 | + Logger *log.Logger |
| 48 | + Stdouts []io.Writer |
| 49 | + Closers []io.Closer |
| 50 | + ExpectObservers []ExpectObserver |
| 51 | + SendObservers []SendObserver |
| 52 | + ReadTimeout *time.Duration |
| 53 | +} |
| 54 | + |
| 55 | +// ExpectObserver provides an interface for a function callback that will |
| 56 | +// be called after each Expect operation. |
| 57 | +// matchers will be the list of active matchers when an error occurred, |
| 58 | +// or a list of matchers that matched `buf` when err is nil. |
| 59 | +// buf is the captured output that was matched against. |
| 60 | +// err is error that might have occurred. May be nil. |
| 61 | +type ExpectObserver func(matchers []Matcher, buf string, err error) |
| 62 | + |
| 63 | +// SendObserver provides an interface for a function callback that will |
| 64 | +// be called after each Send operation. |
| 65 | +// msg is the string that was sent. |
| 66 | +// num is the number of bytes actually sent. |
| 67 | +// err is the error that might have occured. May be nil. |
| 68 | +type SendObserver func(msg string, num int, err error) |
| 69 | + |
| 70 | +// WithStdout adds writers that Console duplicates writes to, similar to the |
| 71 | +// Unix tee(1) command. |
| 72 | +// |
| 73 | +// Each write is written to each listed writer, one at a time. Console is the |
| 74 | +// last writer, writing to it's internal buffer for matching expects. |
| 75 | +// If a listed writer returns an error, that overall write operation stops and |
| 76 | +// returns the error; it does not continue down the list. |
| 77 | +func WithStdout(writers ...io.Writer) ConsoleOpt { |
| 78 | + return func(opts *ConsoleOpts) error { |
| 79 | + opts.Stdouts = append(opts.Stdouts, writers...) |
| 80 | + return nil |
| 81 | + } |
| 82 | +} |
| 83 | + |
| 84 | +// WithCloser adds closers that are closed in order when Console is closed. |
| 85 | +func WithCloser(closer ...io.Closer) ConsoleOpt { |
| 86 | + return func(opts *ConsoleOpts) error { |
| 87 | + opts.Closers = append(opts.Closers, closer...) |
| 88 | + return nil |
| 89 | + } |
| 90 | +} |
| 91 | + |
| 92 | +// WithLogger adds a logger for Console to log debugging information to. By |
| 93 | +// default Console will discard logs. |
| 94 | +func WithLogger(logger *log.Logger) ConsoleOpt { |
| 95 | + return func(opts *ConsoleOpts) error { |
| 96 | + opts.Logger = logger |
| 97 | + return nil |
| 98 | + } |
| 99 | +} |
| 100 | + |
| 101 | +// WithExpectObserver adds an ExpectObserver to allow monitoring Expect operations. |
| 102 | +func WithExpectObserver(observers ...ExpectObserver) ConsoleOpt { |
| 103 | + return func(opts *ConsoleOpts) error { |
| 104 | + opts.ExpectObservers = append(opts.ExpectObservers, observers...) |
| 105 | + return nil |
| 106 | + } |
| 107 | +} |
| 108 | + |
| 109 | +// WithSendObserver adds a SendObserver to allow monitoring Send operations. |
| 110 | +func WithSendObserver(observers ...SendObserver) ConsoleOpt { |
| 111 | + return func(opts *ConsoleOpts) error { |
| 112 | + opts.SendObservers = append(opts.SendObservers, observers...) |
| 113 | + return nil |
| 114 | + } |
| 115 | +} |
| 116 | + |
| 117 | +// WithDefaultTimeout sets a default read timeout during Expect statements. |
| 118 | +func WithDefaultTimeout(timeout time.Duration) ConsoleOpt { |
| 119 | + return func(opts *ConsoleOpts) error { |
| 120 | + opts.ReadTimeout = &timeout |
| 121 | + return nil |
| 122 | + } |
| 123 | +} |
| 124 | + |
| 125 | +// NewConsole returns a new Console with the given options. |
| 126 | +func NewConsole(opts ...ConsoleOpt) (*Console, error) { |
| 127 | + options := ConsoleOpts{ |
| 128 | + Logger: log.New(ioutil.Discard, "", 0), |
| 129 | + } |
| 130 | + |
| 131 | + for _, opt := range opts { |
| 132 | + if err := opt(&options); err != nil { |
| 133 | + return nil, err |
| 134 | + } |
| 135 | + } |
| 136 | + |
| 137 | + pty, err := pty.New() |
| 138 | + if err != nil { |
| 139 | + return nil, err |
| 140 | + } |
| 141 | + closers := append(options.Closers, pty) |
| 142 | + reader := pty.Reader() |
| 143 | + |
| 144 | + passthroughPipe, err := NewPassthroughPipe(reader) |
| 145 | + if err != nil { |
| 146 | + return nil, err |
| 147 | + } |
| 148 | + closers = append(closers, passthroughPipe) |
| 149 | + |
| 150 | + c := &Console{ |
| 151 | + opts: options, |
| 152 | + pty: pty, |
| 153 | + passthroughPipe: passthroughPipe, |
| 154 | + runeReader: bufio.NewReaderSize(passthroughPipe, utf8.UTFMax), |
| 155 | + closers: closers, |
| 156 | + } |
| 157 | + |
| 158 | + /*for _, stdin := range options.Stdins { |
| 159 | + go func(stdin io.Reader) { |
| 160 | + _, err := io.Copy(c, stdin) |
| 161 | + if err != nil { |
| 162 | + c.Logf("failed to copy stdin: %s", err) |
| 163 | + } |
| 164 | + }(stdin) |
| 165 | + }*/ |
| 166 | + |
| 167 | + return c, nil |
| 168 | +} |
| 169 | + |
| 170 | +// Tty returns an input Tty for accepting input |
| 171 | +func (c *Console) InTty() *os.File { |
| 172 | + return c.pty.InPipe() |
| 173 | +} |
| 174 | + |
| 175 | +// OutTty returns an output tty for writing |
| 176 | +func (c *Console) OutTty() *os.File { |
| 177 | + return c.pty.OutPipe() |
| 178 | +} |
| 179 | + |
| 180 | +// Read reads bytes b from Console's tty. |
| 181 | +/*func (c *Console) Read(b []byte) (int, error) { |
| 182 | + return c.ptm.Read(b) |
| 183 | +}*/ |
| 184 | + |
| 185 | +// Write writes bytes b to Console's tty. |
| 186 | +/*func (c *Console) Write(b []byte) (int, error) { |
| 187 | + c.Logf("console write: %q", b) |
| 188 | + return c.ptm.Write(b) |
| 189 | +}*/ |
| 190 | + |
| 191 | +// Fd returns Console's file descripting referencing the master part of its |
| 192 | +// pty. |
| 193 | +/*func (c *Console) Fd() uintptr { |
| 194 | + return c.ptm.Fd() |
| 195 | +}*/ |
| 196 | + |
| 197 | +// Close closes Console's tty. Calling Close will unblock Expect and ExpectEOF. |
| 198 | +func (c *Console) Close() error { |
| 199 | + for _, fd := range c.closers { |
| 200 | + err := fd.Close() |
| 201 | + if err != nil { |
| 202 | + c.Logf("failed to close: %s", err) |
| 203 | + } |
| 204 | + } |
| 205 | + return nil |
| 206 | +} |
| 207 | + |
| 208 | +// Send writes string s to Console's tty. |
| 209 | +func (c *Console) Send(s string) (int, error) { |
| 210 | + c.Logf("console send: %q", s) |
| 211 | + n, err := c.pty.WriteString(s) |
| 212 | + for _, observer := range c.opts.SendObservers { |
| 213 | + observer(s, n, err) |
| 214 | + } |
| 215 | + return n, err |
| 216 | +} |
| 217 | + |
| 218 | +// SendLine writes string s to Console's tty with a trailing newline. |
| 219 | +func (c *Console) SendLine(s string) (int, error) { |
| 220 | + return c.Send(fmt.Sprintf("%s\n", s)) |
| 221 | +} |
| 222 | + |
| 223 | +// Log prints to Console's logger. |
| 224 | +// Arguments are handled in the manner of fmt.Print. |
| 225 | +func (c *Console) Log(v ...interface{}) { |
| 226 | + c.opts.Logger.Print(v...) |
| 227 | +} |
| 228 | + |
| 229 | +// Logf prints to Console's logger. |
| 230 | +// Arguments are handled in the manner of fmt.Printf. |
| 231 | +func (c *Console) Logf(format string, v ...interface{}) { |
| 232 | + c.opts.Logger.Printf(format, v...) |
| 233 | +} |
0 commit comments