github.com/superfly/nomad@v0.10.5-fly/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 inferredTty := isTty() 109 flags.BoolVar(&ttyOpt, "t", inferredTty, "") 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 return 1 198 } 199 } 200 201 if err := validateTaskExistsInAllocation(task, alloc); err != nil { 202 l.Ui.Error(err.Error()) 203 return 1 204 } 205 206 if l.Stdin == nil { 207 l.Stdin = os.Stdin 208 } 209 if l.Stdout == nil { 210 l.Stdout = os.Stdout 211 } 212 if l.Stderr == nil { 213 l.Stderr = os.Stderr 214 } 215 216 var stdin io.Reader = l.Stdin 217 if !stdinOpt { 218 stdin = bytes.NewReader(nil) 219 } 220 221 code, err := l.execImpl(client, alloc, task, ttyOpt, command, escapeChar, stdin, l.Stdout, l.Stderr) 222 if err != nil { 223 l.Ui.Error(fmt.Sprintf("failed to exec into task: %v", err)) 224 return 1 225 } 226 227 return code 228 } 229 230 // execImpl invokes the Alloc Exec api call, it also prepares and restores terminal states as necessary. 231 func (l *AllocExecCommand) execImpl(client *api.Client, alloc *api.Allocation, task string, tty bool, 232 command []string, escapeChar string, stdin io.Reader, stdout, stderr io.WriteCloser) (int, error) { 233 234 sizeCh := make(chan api.TerminalSize, 1) 235 236 ctx, cancelFn := context.WithCancel(context.Background()) 237 defer cancelFn() 238 239 // When tty, ensures we capture all user input and monitor terminal resizes. 240 if tty { 241 if stdin == nil { 242 return -1, fmt.Errorf("stdin is null") 243 } 244 245 inCleanup, err := setRawTerminal(stdin) 246 if err != nil { 247 return -1, err 248 } 249 defer inCleanup() 250 251 outCleanup, err := setRawTerminalOutput(stdout) 252 if err != nil { 253 return -1, err 254 } 255 defer outCleanup() 256 257 sizeCleanup, err := watchTerminalSize(stdout, sizeCh) 258 if err != nil { 259 return -1, err 260 } 261 defer sizeCleanup() 262 263 if escapeChar != "" { 264 stdin = escapingio.NewReader(stdin, escapeChar[0], func(c byte) bool { 265 switch c { 266 case '.': 267 // need to restore tty state so error reporting here 268 // gets emitted at beginning of line 269 outCleanup() 270 inCleanup() 271 272 stderr.Write([]byte("\nConnection closed\n")) 273 cancelFn() 274 return true 275 default: 276 return false 277 } 278 }) 279 } 280 } 281 282 signalCh := make(chan os.Signal, 1) 283 signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM) 284 go func() { 285 for range signalCh { 286 cancelFn() 287 } 288 }() 289 290 return client.Allocations().Exec(ctx, 291 alloc, task, tty, command, stdin, stdout, stderr, sizeCh, nil) 292 } 293 294 // isTty returns true if both stdin and stdout are a TTY 295 func isTty() bool { 296 _, isStdinTerminal := term.GetFdInfo(os.Stdin) 297 _, isStdoutTerminal := term.GetFdInfo(os.Stdout) 298 return isStdinTerminal && isStdoutTerminal 299 } 300 301 // setRawTerminal sets the stream terminal in raw mode, so process captures 302 // Ctrl+C and other commands to forward to remote process. 303 // It returns a cleanup function that restores terminal to original mode. 304 func setRawTerminal(stream interface{}) (cleanup func(), err error) { 305 fd, isTerminal := term.GetFdInfo(stream) 306 if !isTerminal { 307 return nil, errors.New("not a terminal") 308 } 309 310 state, err := term.SetRawTerminal(fd) 311 if err != nil { 312 return nil, err 313 } 314 315 return func() { term.RestoreTerminal(fd, state) }, nil 316 } 317 318 // setRawTerminalOutput sets the output stream in Windows to raw mode, 319 // so it disables LF -> CRLF translation. 320 // It's basically a no-op on unix. 321 func setRawTerminalOutput(stream interface{}) (cleanup func(), err error) { 322 fd, isTerminal := term.GetFdInfo(stream) 323 if !isTerminal { 324 return nil, errors.New("not a terminal") 325 } 326 327 state, err := term.SetRawTerminalOutput(fd) 328 if err != nil { 329 return nil, err 330 } 331 332 return func() { term.RestoreTerminal(fd, state) }, nil 333 } 334 335 // watchTerminalSize watches terminal size changes to propagate to remote tty. 336 func watchTerminalSize(out io.Writer, resize chan<- api.TerminalSize) (func(), error) { 337 fd, isTerminal := term.GetFdInfo(out) 338 if !isTerminal { 339 return nil, errors.New("not a terminal") 340 } 341 342 ctx, cancel := context.WithCancel(context.Background()) 343 344 signalCh := make(chan os.Signal, 1) 345 setupWindowNotification(signalCh) 346 347 sendTerminalSize := func() { 348 s, err := term.GetWinsize(fd) 349 if err != nil { 350 return 351 } 352 353 resize <- api.TerminalSize{ 354 Height: int(s.Height), 355 Width: int(s.Width), 356 } 357 } 358 go func() { 359 for { 360 select { 361 case <-ctx.Done(): 362 return 363 case <-signalCh: 364 sendTerminalSize() 365 } 366 } 367 }() 368 369 go func() { 370 // send initial size 371 sendTerminalSize() 372 }() 373 374 return cancel, nil 375 }