terraform/vendor/github.com/masterzen/winrm/command.go

240 lines
4.7 KiB
Go

package winrm
import (
"bytes"
"errors"
"io"
"strings"
)
type commandWriter struct {
*Command
eof bool
}
type commandReader struct {
*Command
write *io.PipeWriter
read *io.PipeReader
stream string
}
// Command represents a given command running on a Shell. This structure allows to get access
// to the various stdout, stderr and stdin pipes.
type Command struct {
client *Client
shell *Shell
id string
exitCode int
finished bool
err error
Stdin *commandWriter
Stdout *commandReader
Stderr *commandReader
done chan struct{}
cancel chan struct{}
}
func newCommand(shell *Shell, ids string) *Command {
command := &Command{
shell: shell,
client: shell.client,
id: ids,
exitCode: 0,
err: nil,
done: make(chan struct{}),
cancel: make(chan struct{}),
}
command.Stdout = newCommandReader("stdout", command)
command.Stdin = &commandWriter{
Command: command,
eof: false,
}
command.Stderr = newCommandReader("stderr", command)
go fetchOutput(command)
return command
}
func newCommandReader(stream string, command *Command) *commandReader {
read, write := io.Pipe()
return &commandReader{
Command: command,
stream: stream,
write: write,
read: read,
}
}
func fetchOutput(command *Command) {
for {
select {
case <-command.cancel:
close(command.done)
return
default:
finished, err := command.slurpAllOutput()
if finished {
command.err = err
close(command.done)
return
}
}
}
}
func (c *Command) check() error {
if c.id == "" {
return errors.New("Command has already been closed")
}
if c.shell == nil {
return errors.New("Command has no associated shell")
}
if c.client == nil {
return errors.New("Command has no associated client")
}
return nil
}
// Close will terminate the running command
func (c *Command) Close() error {
if err := c.check(); err != nil {
return err
}
select { // close cancel channel if it's still open
case <-c.cancel:
default:
close(c.cancel)
}
request := NewSignalRequest(c.client.url, c.shell.id, c.id, &c.client.Parameters)
defer request.Free()
_, err := c.client.sendRequest(request)
return err
}
func (c *Command) slurpAllOutput() (bool, error) {
if err := c.check(); err != nil {
c.Stderr.write.CloseWithError(err)
c.Stdout.write.CloseWithError(err)
return true, err
}
request := NewGetOutputRequest(c.client.url, c.shell.id, c.id, "stdout stderr", &c.client.Parameters)
defer request.Free()
response, err := c.client.sendRequest(request)
if err != nil {
if strings.Contains(err.Error(), "OperationTimeout") {
// Operation timeout because there was no command output
return false, err
}
if strings.Contains(err.Error(), "EOF") {
c.exitCode = 16001
}
c.Stderr.write.CloseWithError(err)
c.Stdout.write.CloseWithError(err)
return true, err
}
var exitCode int
var stdout, stderr bytes.Buffer
finished, exitCode, err := ParseSlurpOutputErrResponse(response, &stdout, &stderr)
if err != nil {
c.Stderr.write.CloseWithError(err)
c.Stdout.write.CloseWithError(err)
return true, err
}
if stdout.Len() > 0 {
c.Stdout.write.Write(stdout.Bytes())
}
if stderr.Len() > 0 {
c.Stderr.write.Write(stderr.Bytes())
}
if finished {
c.exitCode = exitCode
c.Stderr.write.Close()
c.Stdout.write.Close()
}
return finished, nil
}
func (c *Command) sendInput(data []byte) error {
if err := c.check(); err != nil {
return err
}
request := NewSendInputRequest(c.client.url, c.shell.id, c.id, data, &c.client.Parameters)
defer request.Free()
_, err := c.client.sendRequest(request)
return err
}
// ExitCode returns command exit code when it is finished. Before that the result is always 0.
func (c *Command) ExitCode() int {
return c.exitCode
}
// Wait function will block the current goroutine until the remote command terminates.
func (c *Command) Wait() {
// block until finished
<-c.done
}
// Write data to this Pipe
// commandWriter implements io.Writer interface
func (w *commandWriter) Write(data []byte) (int, error) {
var (
written int
err error
)
for len(data) > 0 {
if w.eof {
return written, io.EOF
}
// never send more data than our EnvelopeSize.
n := min(w.client.Parameters.EnvelopeSize-1000, len(data))
if err := w.sendInput(data[:n]); err != nil {
break
}
data = data[n:]
written += n
}
return written, err
}
func min(a int, b int) int {
if a < b {
return a
}
return b
}
// Close method wrapper
// commandWriter implements io.Closer interface
func (w *commandWriter) Close() error {
w.eof = true
return w.Close()
}
// Read data from this Pipe
func (r *commandReader) Read(buf []byte) (int, error) {
n, err := r.read.Read(buf)
if err != nil && err != io.EOF {
return 0, err
}
return n, err
}