github.com/ferranbt/nomad@v0.9.3-0.20190607002617-85c449b7667c/command/alloc_exec.go (about) 1 package command 2 3 import ( 4 "bytes" 5 "context" 6 "errors" 7 "fmt" 8 "io" 9 "os" 10 "os/signal" 11 "strings" 12 "syscall" 13 14 "github.com/docker/docker/pkg/term" 15 "github.com/hashicorp/nomad/api" 16 "github.com/hashicorp/nomad/api/contexts" 17 "github.com/hashicorp/nomad/helper/escapingio" 18 "github.com/posener/complete" 19 ) 20 21 type AllocExecCommand struct { 22 Meta 23 24 Stdin io.Reader 25 Stdout io.WriteCloser 26 Stderr io.WriteCloser 27 } 28 29 func (l *AllocExecCommand) Help() string { 30 helpText := ` 31 Usage: nomad alloc exec [options] <allocation> <command> 32 33 Run command inside the environment of the given allocation and task. 34 35 General Options: 36 37 ` + generalOptionsUsage() + ` 38 39 Exec Specific Options: 40 41 -task <task-name> 42 Sets the task to exec command in 43 44 -job 45 Use a random allocation from the specified job ID. 46 47 -i 48 Pass stdin to the container, defaults to true. Pass -i=false to disable. 49 50 -t 51 Allocate a pseudo-tty, defaults to true if stdin is detected to be a tty session. 52 Pass -t=false to disable explicitly. 53 54 -e <escape_char> 55 Sets the escape character for sessions with a pty (default: '~'). The escape 56 character is only recognized at the beginning of a line. The escape character 57 followed by a dot ('.') closes the connection. Setting the character to 58 'none' disables any escapes and makes the session fully transparent. 59 ` 60 return strings.TrimSpace(helpText) 61 } 62 63 func (l *AllocExecCommand) Synopsis() string { 64 return "Execute commands in task" 65 } 66 67 func (c *AllocExecCommand) AutocompleteFlags() complete.Flags { 68 return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient), 69 complete.Flags{ 70 "--task": complete.PredictAnything, 71 "-job": complete.PredictAnything, 72 "-i": complete.PredictNothing, 73 "-t": complete.PredictNothing, 74 "-e": complete.PredictSet("none", "~"), 75 }) 76 } 77 78 func (l *AllocExecCommand) AutocompleteArgs() complete.Predictor { 79 return complete.PredictFunc(func(a complete.Args) []string { 80 client, err := l.Meta.Client() 81 if err != nil { 82 return nil 83 } 84 85 resp, _, err := client.Search().PrefixSearch(a.Last, contexts.Allocs, nil) 86 if err != nil { 87 return []string{} 88 } 89 return resp.Matches[contexts.Allocs] 90 }) 91 } 92 93 func (l *AllocExecCommand) Name() string { return "alloc exec" } 94 95 func (l *AllocExecCommand) Run(args []string) int { 96 var job, stdinOpt, ttyOpt bool 97 var task string 98 var escapeChar string 99 100 flags := l.Meta.FlagSet(l.Name(), FlagSetClient) 101 flags.Usage = func() { l.Ui.Output(l.Help()) } 102 flags.BoolVar(&job, "job", false, "") 103 flags.StringVar(&task, "task", "", "") 104 flags.StringVar(&escapeChar, "e", "~", "") 105 106 flags.BoolVar(&stdinOpt, "i", true, "") 107 108 stdinTty := isStdinTty() 109 flags.BoolVar(&ttyOpt, "t", stdinTty, "") 110 111 if err := flags.Parse(args); err != nil { 112 return 1 113 } 114 args = flags.Args() 115 116 if ttyOpt && !stdinOpt { 117 l.Ui.Error("-i must be enabled if running with tty") 118 return 1 119 } 120 121 if escapeChar == "none" { 122 escapeChar = "" 123 } else if len(escapeChar) > 1 { 124 l.Ui.Error("-e requires 'none' or a single character") 125 return 1 126 } 127 128 if numArgs := len(args); numArgs < 1 { 129 if job { 130 l.Ui.Error("A job ID is required") 131 } else { 132 l.Ui.Error("An allocation ID is required") 133 } 134 135 l.Ui.Error(commandErrorText(l)) 136 return 1 137 } else if numArgs < 2 { 138 l.Ui.Error("A command is required") 139 l.Ui.Error(commandErrorText(l)) 140 return 1 141 } 142 143 command := args[1:] 144 145 client, err := l.Meta.Client() 146 if err != nil { 147 l.Ui.Error(fmt.Sprintf("Error initializing client: %v", err)) 148 return 1 149 } 150 151 // If -job is specified, use random allocation, otherwise use provided allocation 152 allocID := args[0] 153 if job { 154 allocID, err = getRandomJobAlloc(client, args[0]) 155 if err != nil { 156 l.Ui.Error(fmt.Sprintf("Error fetching allocations: %v", err)) 157 return 1 158 } 159 } 160 161 length := shortId 162 163 // Query the allocation info 164 if len(allocID) == 1 { 165 l.Ui.Error(fmt.Sprintf("Alloc ID must contain at least two characters.")) 166 return 1 167 } 168 169 allocID = sanitizeUUIDPrefix(allocID) 170 allocs, _, err := client.Allocations().PrefixList(allocID) 171 if err != nil { 172 l.Ui.Error(fmt.Sprintf("Error querying allocation: %v", err)) 173 return 1 174 } 175 if len(allocs) == 0 { 176 l.Ui.Error(fmt.Sprintf("No allocation(s) with prefix or id %q found", allocID)) 177 return 1 178 } 179 if len(allocs) > 1 { 180 // Format the allocs 181 out := formatAllocListStubs(allocs, false, length) 182 l.Ui.Error(fmt.Sprintf("Prefix matched multiple allocations\n\n%s", out)) 183 return 1 184 } 185 // Prefix lookup matched a single allocation 186 alloc, _, err := client.Allocations().Info(allocs[0].ID, nil) 187 if err != nil { 188 l.Ui.Error(fmt.Sprintf("Error querying allocation: %s", err)) 189 return 1 190 } 191 192 if task == "" { 193 task, err = lookupAllocTask(alloc) 194 195 if err != nil { 196 l.Ui.Error(err.Error()) 197 l.Ui.Error("\nPlease specify the task.") 198 return 1 199 } 200 } 201 202 if err := validateTaskExistsInAllocation(task, alloc); err != nil { 203 l.Ui.Error(err.Error()) 204 return 1 205 } 206 207 if l.Stdin == nil { 208 l.Stdin = os.Stdin 209 } 210 if l.Stdout == nil { 211 l.Stdout = os.Stdout 212 } 213 if l.Stderr == nil { 214 l.Stderr = os.Stderr 215 } 216 217 var stdin io.Reader = l.Stdin 218 if !stdinOpt { 219 stdin = bytes.NewReader(nil) 220 } 221 222 code, err := l.execImpl(client, alloc, task, ttyOpt, command, escapeChar, stdin, l.Stdout, l.Stderr) 223 if err != nil { 224 l.Ui.Error(fmt.Sprintf("failed to exec into task: %v", err)) 225 return 1 226 } 227 228 return code 229 } 230 231 // execImpl invokes the Alloc Exec api call, it also prepares and restores terminal states as necessary. 232 func (l *AllocExecCommand) execImpl(client *api.Client, alloc *api.Allocation, task string, tty bool, 233 command []string, escapeChar string, stdin io.Reader, stdout, stderr io.WriteCloser) (int, error) { 234 235 sizeCh := make(chan api.TerminalSize, 1) 236 237 ctx, cancelFn := context.WithCancel(context.Background()) 238 defer cancelFn() 239 240 // When tty, ensures we capture all user input and monitor terminal resizes. 241 if tty { 242 if stdin == nil { 243 return -1, fmt.Errorf("stdin is null") 244 } 245 246 inCleanup, err := setRawTerminal(stdin) 247 if err != nil { 248 return -1, err 249 } 250 defer inCleanup() 251 252 outCleanup, err := setRawTerminalOutput(stdout) 253 if err != nil { 254 return -1, err 255 } 256 defer outCleanup() 257 258 sizeCleanup, err := watchTerminalSize(stdout, sizeCh) 259 if err != nil { 260 return -1, err 261 } 262 defer sizeCleanup() 263 264 if escapeChar != "" { 265 stdin = escapingio.NewReader(stdin, escapeChar[0], func(c byte) bool { 266 switch c { 267 case '.': 268 // need to restore tty state so error reporting here 269 // gets emitted at beginning of line 270 outCleanup() 271 inCleanup() 272 273 stderr.Write([]byte("\nConnection closed\n")) 274 cancelFn() 275 return true 276 default: 277 return false 278 } 279 }) 280 } 281 } 282 283 signalCh := make(chan os.Signal, 1) 284 signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM) 285 go func() { 286 for range signalCh { 287 cancelFn() 288 } 289 }() 290 291 return client.Allocations().Exec(ctx, 292 alloc, task, tty, command, stdin, stdout, stderr, sizeCh, nil) 293 } 294 295 func isStdinTty() bool { 296 _, isTerminal := term.GetFdInfo(os.Stdin) 297 return isTerminal 298 } 299 300 // setRawTerminal sets the stream terminal in raw mode, so process captures 301 // Ctrl+C and other commands to forward to remote process. 302 // It returns a cleanup function that restores terminal to original mode. 303 func setRawTerminal(stream interface{}) (cleanup func(), err error) { 304 fd, isTerminal := term.GetFdInfo(stream) 305 if !isTerminal { 306 return nil, errors.New("not a terminal") 307 } 308 309 state, err := term.SetRawTerminal(fd) 310 if err != nil { 311 return nil, err 312 } 313 314 return func() { term.RestoreTerminal(fd, state) }, nil 315 } 316 317 // setRawTerminalOutput sets the output stream in Windows to raw mode, 318 // so it disables LF -> CRLF translation. 319 // It's basically a no-op on unix. 320 func setRawTerminalOutput(stream interface{}) (cleanup func(), err error) { 321 fd, isTerminal := term.GetFdInfo(stream) 322 if !isTerminal { 323 return nil, errors.New("not a terminal") 324 } 325 326 state, err := term.SetRawTerminalOutput(fd) 327 if err != nil { 328 return nil, err 329 } 330 331 return func() { term.RestoreTerminal(fd, state) }, nil 332 } 333 334 // watchTerminalSize watches terminal size changes to propagate to remote tty. 335 func watchTerminalSize(out io.Writer, resize chan<- api.TerminalSize) (func(), error) { 336 fd, isTerminal := term.GetFdInfo(out) 337 if !isTerminal { 338 return nil, errors.New("not a terminal") 339 } 340 341 ctx, cancel := context.WithCancel(context.Background()) 342 343 signalCh := make(chan os.Signal, 1) 344 setupWindowNotification(signalCh) 345 346 sendTerminalSize := func() { 347 s, err := term.GetWinsize(fd) 348 if err != nil { 349 return 350 } 351 352 resize <- api.TerminalSize{ 353 Height: int(s.Height), 354 Width: int(s.Width), 355 } 356 } 357 go func() { 358 for { 359 select { 360 case <-ctx.Done(): 361 return 362 case <-signalCh: 363 sendTerminalSize() 364 } 365 } 366 }() 367 368 go func() { 369 // send initial size 370 sendTerminalSize() 371 }() 372 373 return cancel, nil 374 }