-
Notifications
You must be signed in to change notification settings - Fork 1
/
run.go
196 lines (177 loc) · 6.3 KB
/
run.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
package cliqmstr
import (
"fmt"
"io/ioutil"
"log"
"os"
"os/exec"
"path"
"path/filepath"
"regexp"
"strings"
"syscall"
"github.com/spf13/cobra"
)
var (
keepTmpDirectories bool // Keep intermediate files
instdir string // create instrumentation in this dir
wrappedCmds = []string{"gcc", "g++", "ar", "ld", "as", "objcopy", "strip"} // constant (which Go does not do with arrays)
)
var runCmd = &cobra.Command{
Use: "run",
Short: "Run a shell command with QMSTR instrumentation",
Long: `run executes the argument after applying temporary QMSTR instrumentation.
For example, "qmstr make" sets up a QMSTR instrumented environment, then executes "make", and finally restores the environment and deletes all temporary files.
The environment variable QMSTR_INSTRUMENTATION_HOME will be defined to point to the QMSTR instrumentation directory.`,
Args: func(cmd *cobra.Command, args []string) error {
if len(args) < 1 && !keepTmpDirectories {
return fmt.Errorf("no command specified")
}
return nil
},
Run: executeRun,
}
func init() {
rootCmd.AddCommand(runCmd)
runCmd.Flags().SetInterspersed(false) // this stops flag parsing at the first argument, allowing for e.g. "ls -la" as payload
runCmd.Flags().BoolVarP(&keepTmpDirectories, "keep", "k", false,
"Keep the created directories instead of cleaning up.")
runCmd.Flags().StringVarP(&instdir, "instdir", "i", "", "Create instrumentation in this directory (optional)")
}
func executeRun(cmd *cobra.Command, args []string) {
exitCode := Run(args)
os.Exit(exitCode)
}
// Run does everything
// It also makes sure that even though os.Exit() is called later, all defered functions are properly called.
func Run(payloadCmd []string) int {
if len(payloadCmd) == 0 && !keepTmpDirectories {
// This should have been caught by runCmd's Args validator:
log.Fatal("command validation constraint violated")
}
var tmpWorkDir string
var err error
if instdir != "" {
tmpWorkDir = instdir
err = os.MkdirAll(tmpWorkDir, os.ModePerm)
} else {
tmpWorkDir, err = ioutil.TempDir("", "qmstr-bin-")
}
if err != nil {
Log.Fatalf("error creating temporary working directory: %v", err)
}
defer func() {
if keepTmpDirectories {
Debug.Printf("keeping temporary directory at %v", tmpWorkDir)
} else {
Debug.Printf("deleting temporary instrumentation directory in %v", tmpWorkDir)
if err := os.RemoveAll(tmpWorkDir); err != nil {
// it is a warning because the program is exiting and we cannot recover anymore
Log.Printf("warning - error deleting temporary instrumentation bin directory in %v: %v", tmpWorkDir, err)
}
}
}()
SetupCompilerInstrumentation(tmpWorkDir)
if len(payloadCmd) > 0 {
exitCode, err := RunPayloadCommand(payloadCmd[0], payloadCmd[1:]...)
if err != nil {
Debug.Printf("payload command exited with non-zero exit code: %v", exitCode)
fmt.Print(err)
}
return exitCode
}
return 0
}
// SetupCompilerInstrumentation creates the QMSTR instrumentation symlinks in the given path
func SetupCompilerInstrumentation(tmpWorkDir string) {
executable, err := os.Executable()
if err != nil {
Log.Fatalf("unable to find myself: %v", err)
}
ownPath, err := filepath.Abs(filepath.Dir(executable))
if err != nil {
Log.Fatalf("unable to determine path to executable: %v", err)
}
const wrapper = "qmstr-wrapper"
wrapperPath := path.Join(ownPath, wrapper)
if _, err := os.Stat(wrapperPath); err != nil {
Debug.Printf("cannot find %s at %s: %v (optional)", wrapper, wrapperPath, err)
// optionally, search the path and use a qmstr-wrapper if found there
wrapperPath, err = exec.LookPath(wrapper)
if err != nil {
Log.Fatalf("%v not found next in %v or in the PATH", wrapper, executable)
}
}
// create a "bin" directory in the temporary directory
binDir := strings.TrimSpace(filepath.Join(tmpWorkDir, "bin"))
if err := link(wrapperPath, binDir, wrappedCmds); err != nil {
Log.Fatalf("setting up instrumentation failed: %v", err)
}
// Setup environment variables
envvars := make(map[string][]string)
envvars["gcc"] = []string{"CMAKE_LINKER", "CC"}
envvars["g++"] = []string{"CXX"}
files, err := ioutil.ReadDir(binDir)
for _, file := range files {
cmd := file.Name()
if envs, ok := envvars[cmd]; ok {
for _, envvar := range envs {
os.Setenv(envvar, filepath.Join(binDir, file.Name()))
}
}
}
// extend the PATH variable to include the created bin/ directory
paths := filepath.SplitList(os.Getenv("PATH"))
hasWhiteSpace, _ := regexp.Compile("\\s+")
for index, value := range paths {
if hasWhiteSpace.MatchString(value) {
Debug.Printf("NOTE - your PATH contains an element with whitespace in it: %v", value)
paths[index] = fmt.Sprintf("\"%s\"", value)
}
}
paths = append([]string{binDir}, paths...)
separator := string(os.PathListSeparator)
newPath := strings.Join(paths, separator)
os.Setenv("PATH", newPath)
Debug.Printf("PATH is now %v\n", os.Getenv("PATH"))
os.Setenv("QMSTR_INSTRUMENTATION_HOME", tmpWorkDir)
Debug.Printf("QMSTR_INSTRUMENTATION_HOME is now %v\n", os.Getenv("QMSTR_INSTRUMENTATION_HOME"))
if keepTmpDirectories {
fmt.Printf("export PATH=%v\n", os.Getenv("PATH"))
fmt.Printf("export QMSTR_INSTRUMENTATION_HOME=%v\n", os.Getenv("QMSTR_INSTRUMENTATION_HOME"))
}
}
func link(source string, binDir string, targets []string) error {
if err := os.Mkdir(binDir, 0700); err != nil {
return fmt.Errorf("unable to create %v: %v", binDir, err)
}
// create the symlinks to qmstr-wrapper in there
symlinks := make(map[string]string)
for _, cmd := range wrappedCmds {
symlink := path.Join(binDir, cmd)
symlinks[symlink] = source
}
for from, to := range symlinks {
if err := os.Symlink(to, from); err != nil {
return fmt.Errorf("cannot symlink %s to %s: %v", from, to, err)
}
}
return nil
}
// RunPayloadCommand performs the payload command and returns it's exit code and/or an error
func RunPayloadCommand(command string, arguments ...string) (int, error) {
cmd := exec.Command(command, arguments...)
cmd.Stdin = os.Stdin
cmd.Stderr = os.Stderr
cmd.Stdout = os.Stdout
if err := cmd.Run(); err != nil {
switch value := err.(type) {
case *exec.ExitError:
ws := value.Sys().(syscall.WaitStatus)
return ws.ExitStatus(), err
default:
return 1, err
}
}
return 0, nil
}