19166: Close ssh session when exec/shell command exits.

[arvados.git] / cmd / arvados-client / container_gateway.go

diff --git a/cmd/arvados-client/container_gateway.go b/cmd/arvados-client/container_gateway.go
index e3b6b9332c23c92c0d62ab9417c3af6d56d0c9c2..55f8c33bc70c77d31f13f16bb924ee4c2a6a1613 100644 (file)
--- a/cmd/arvados-client/container_gateway.go
+++ b/cmd/arvados-client/container_gateway.go
@@ -17,6 +17,7 @@ import (
        "strings"
        "syscall"
 
+       "git.arvados.org/arvados.git/lib/cmd"
        "git.arvados.org/arvados.git/lib/controller/rpc"
        "git.arvados.org/arvados.git/sdk/go/arvados"
 )
@@ -27,26 +28,11 @@ type shellCommand struct{}
 
 func (shellCommand) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
        f := flag.NewFlagSet(prog, flag.ContinueOnError)
-       f.SetOutput(stderr)
-       f.Usage = func() {
-               _, prog := filepath.Split(prog)
-               fmt.Fprint(stderr, prog+`: open an interactive shell on a running container.
-
-Usage: `+prog+` [options] [username@]container-uuid [ssh-options] [remote-command [args...]]
-
-Options:
-`)
-               f.PrintDefaults()
-       }
        detachKeys := f.String("detach-keys", "ctrl-],ctrl-]", "set detach key sequence, as in docker-attach(1)")
-       err := f.Parse(args)
-       if err != nil {
-               fmt.Fprintln(stderr, err)
-               return 2
-       }
-
-       if f.NArg() < 1 {
-               f.Usage()
+       if ok, code := cmd.ParseFlags(f, prog, args, "[username@]container-uuid [ssh-options] [remote-command [args...]]", stderr); !ok {
+               return code
+       } else if f.NArg() < 1 {
+               fmt.Fprintf(stderr, "missing required argument: container-uuid (try -help)\n")
                return 2
        }
        target := f.Args()[0]
@@ -67,13 +53,19 @@ Options:
        // kex_exchange_identification: Connection closed by remote host
        // Connection closed by UNKNOWN port 65535
        // exit status 255
+       //
+       // In case our target is a container request, the probe also
+       // resolves it to a container, so we don't connect to two
+       // different containers in a race.
+       var probetarget bytes.Buffer
        exitcode := connectSSHCommand{}.RunCommand(
                "arvados-client connect-ssh",
                []string{"-detach-keys=" + *detachKeys, "-probe-only=true", target},
-               &bytes.Buffer{}, &bytes.Buffer{}, stderr)
+               &bytes.Buffer{}, &probetarget, stderr)
        if exitcode != 0 {
                return exitcode
        }
+       target = strings.Trim(probetarget.String(), "\n")
 
        selfbin, err := os.Readlink("/proc/self/exe")
        if err != nil {
@@ -119,13 +111,12 @@ Options:
 `)
                f.PrintDefaults()
        }
-       probeOnly := f.Bool("probe-only", false, "do not transfer IO, just exit 0 immediately if tunnel setup succeeds")
+       probeOnly := f.Bool("probe-only", false, "do not transfer IO, just setup tunnel, print target UUID, and exit")
        detachKeys := f.String("detach-keys", "", "set detach key sequence, as in docker-attach(1)")
-       if err := f.Parse(args); err != nil {
-               fmt.Fprintln(stderr, err)
-               return 2
+       if ok, code := cmd.ParseFlags(f, prog, args, "[username@]container-uuid", stderr); !ok {
+               return code
        } else if f.NArg() != 1 {
-               f.Usage()
+               fmt.Fprintf(stderr, "missing required argument: [username@]container-uuid\n")
                return 2
        }
        targetUUID := f.Args()[0]
@@ -169,7 +160,9 @@ Options:
                fmt.Fprintf(stderr, "target UUID is not a container or container request UUID: %s\n", targetUUID)
                return 1
        }
-       sshconn, err := rpcconn.ContainerSSH(context.TODO(), arvados.ContainerSSHOptions{
+       ctx, cancel := context.WithCancel(context.Background())
+       defer cancel()
+       sshconn, err := rpcconn.ContainerSSH(ctx, arvados.ContainerSSHOptions{
                UUID:          targetUUID,
                DetachKeys:    *detachKeys,
                LoginUsername: loginUsername,
@@ -181,10 +174,10 @@ Options:
        defer sshconn.Conn.Close()
 
        if *probeOnly {
+               fmt.Fprintln(stdout, targetUUID)
                return 0
        }
 
-       ctx, cancel := context.WithCancel(context.Background())
        go func() {
                defer cancel()
                _, err := io.Copy(stdout, sshconn.Conn)