/
command_executor.go
201 lines (180 loc) · 5.87 KB
/
command_executor.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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
// Copyright © 2016 Asteris, LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package shell
import (
"io"
"io/ioutil"
"os"
"os/exec"
"syscall"
"time"
log "github.com/Sirupsen/logrus"
"github.com/pkg/errors"
)
// NB: Known Bug with timed script execution:
// Currently when a script executes beyond it's alloted time a timeout will
// occur and nil is returned by timeoutExec. The goroutine running the script
// will continue to drain stdout and stderr from the process sockets and they
// will be GCed when the script finally finishes. This means that there is no
// mechanism for getting the output of a script when it has timed out. A
// proposed solution to this would be to implement a ReadUntilWouldBlock-type
// function that would allow us to read into a buffer from a ReadCloser until
// the read operation would block, then return the contents of the buffer (along
// with some marker if we recevied an error or EOF). Then exec function would
// then take in pointers to buffers for stdout and stderr and populate them
// directly, so that if the script execution timed out we would still have a
// reference to those buffers.
var (
ErrTimedOut = errors.New("execution timed out")
)
// A CommandExecutor supports running a script and returning the results wrapped
// in a *CommandResults structure.
type CommandExecutor interface {
Run(string) (*CommandResults, error)
}
// CommandGenerator provides a container to wrap generating a system command
type CommandGenerator struct {
Interpreter string
Flags []string
Dir string
Env []string
Timeout *time.Duration
}
// Run will generate a new command and run it with optional timeout parameters
func (cmd *CommandGenerator) Run(script string) (*CommandResults, error) {
ctx, err := cmd.start()
if err != nil {
return nil, err
}
return ctx.Run(script, cmd.Timeout)
}
func (cmd *CommandGenerator) start() (*commandIOContext, error) {
command := newCommand(cmd)
stdin, stdout, stderr, err := cmdGetPipes(command)
return &commandIOContext{
Command: command,
Stdin: stdin,
Stdout: stdout,
Stderr: stderr,
}, err
}
// commandIOContext provides the context for a command that includes it's stdin,
// stdout, and stderr pipes along with the underlying command.
type commandIOContext struct {
Command *exec.Cmd
Stdin io.WriteCloser
Stdout io.ReadCloser
Stderr io.ReadCloser
}
// Run wraps exec and timeoutExec, executing the script with or without a
// timeout depending whether or not timeout is nil.
func (c *commandIOContext) Run(script string, timeout *time.Duration) (results *CommandResults, err error) {
if timeout == nil {
results, err = c.exec(script)
} else {
results, err = c.timeoutExec(script, *timeout)
}
return
}
// timeoutExec will run the given script with a timelimit specified by
// timeout. If the script does not return with that time duration a
// ScriptTimeoutError is returned.
func (c *commandIOContext) timeoutExec(script string, timeout time.Duration) (*CommandResults, error) {
timeoutChannel := make(chan []interface{}, 1)
go func() {
cmdResults, err := c.exec(script)
timeoutChannel <- []interface{}{cmdResults, err}
}()
select {
case result := <-timeoutChannel:
var errResult error
if result[1] != nil {
errResult = result[1].(error)
}
return result[0].(*CommandResults), errResult
case <-time.After(timeout):
return nil, ErrTimedOut
}
}
func (c *commandIOContext) exec(script string) (results *CommandResults, err error) {
results = &CommandResults{
Stdin: script,
}
// if working dir does not exist, we want the check to return a non-zero
// result. otherwise, Command.Start will return an error and short-circuit
// plan/apply
if c.Command.Dir != "" {
_, err = os.Stat(c.Command.Dir)
if os.IsNotExist(err) {
results.ExitStatus = 1
results.Stdout = err.Error()
return results, nil
}
}
if err = c.Command.Start(); err != nil {
return
}
if _, err = c.Stdin.Write([]byte(script)); err != nil {
return
}
if err = c.Stdin.Close(); err != nil {
return
}
if data, readErr := ioutil.ReadAll(c.Stdout); readErr == nil {
results.Stdout = string(data)
} else {
log.WithField("module", "shell").Warn("cannot read stdout from script")
}
if data, readErr := ioutil.ReadAll(c.Stderr); readErr == nil {
results.Stderr = string(data)
} else {
log.WithField("module", "shell").Warn("cannot read stdout from script")
}
if waitErr := c.Command.Wait(); waitErr == nil {
results.ExitStatus = 0
} else {
exitErr, ok := waitErr.(*exec.ExitError)
if !ok {
err = errors.Wrap(waitErr, "failed to wait on process")
return
}
status, ok := exitErr.Sys().(syscall.WaitStatus)
if !ok {
err = errors.New("unexpected error getting exit status")
}
results.ExitStatus = uint32(status.ExitStatus())
}
results.State = c.Command.ProcessState
return
}
func newCommand(cmd *CommandGenerator) *exec.Cmd {
var command *exec.Cmd
if cmd.Interpreter == "" {
if len(cmd.Flags) > 0 {
log.WithField("module", "shell").WithField("interpreter", "/bin/sh").Debug("passing flags to default interpreter")
command = exec.Command(defaultInterpreter, cmd.Flags...)
} else {
command = exec.Command(defaultInterpreter, defaultExecFlags...)
}
} else {
command = exec.Command(cmd.Interpreter, cmd.Flags...)
}
command.Dir = cmd.Dir
if len(cmd.Env) > 0 {
env := os.Environ()
env = append(env, cmd.Env...)
command.Env = env
}
return command
}