cmd/arvados-client/container_gateway_test.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package main
   6
   7 import (
   8         "bytes"
   9         "context"
  10         "crypto/hmac"
  11         "crypto/sha256"
  12         "fmt"
  13         "io/ioutil"
  14         "net"
  15         "net/http"
  16         "net/url"
  17         "os"
  18         "os/exec"
  19         "strings"
  20         "sync"
  21         "time"
  22
  23         "git.arvados.org/arvados.git/lib/controller/rpc"
  24         "git.arvados.org/arvados.git/lib/crunchrun"
  25         "git.arvados.org/arvados.git/sdk/go/arvados"
  26         "git.arvados.org/arvados.git/sdk/go/arvadostest"
  27         "git.arvados.org/arvados.git/sdk/go/httpserver"
  28         check "gopkg.in/check.v1"
  29 )
  30
  31 func (s *ClientSuite) TestShellGatewayNotAvailable(c *check.C) {
  32         var stdout, stderr bytes.Buffer
  33         cmd := exec.Command("go", "run", ".", "shell", arvadostest.QueuedContainerUUID, "-o", "controlpath=none", "echo", "ok")
  34         cmd.Env = append(cmd.Env, os.Environ()...)
  35         cmd.Env = append(cmd.Env, "ARVADOS_API_TOKEN="+arvadostest.ActiveTokenV2)
  36         cmd.Stdout = &stdout
  37         cmd.Stderr = &stderr
  38         c.Check(cmd.Run(), check.NotNil)
  39         c.Log(stderr.String())
  40         c.Check(stderr.String(), check.Matches, `(?ms).*container is not running yet \(state is "Queued"\).*`)
  41 }
  42
  43 func (s *ClientSuite) TestShellGateway(c *check.C) {
  44         defer func() {
  45                 c.Check(arvados.NewClientFromEnv().RequestAndDecode(nil, "POST", "database/reset", nil, nil), check.IsNil)
  46         }()
  47         uuid := arvadostest.QueuedContainerUUID
  48         h := hmac.New(sha256.New, []byte(arvadostest.SystemRootToken))
  49         fmt.Fprint(h, uuid)
  50         authSecret := fmt.Sprintf("%x", h.Sum(nil))
  51         dcid := "theperthcountyconspiracy"
  52         gw := crunchrun.Gateway{
  53                 DockerContainerID: &dcid,
  54                 ContainerUUID:     uuid,
  55                 Address:           "0.0.0.0:0",
  56                 AuthSecret:        authSecret,
  57                 // Just forward connections to localhost instead of a
  58                 // container, so we can test without running a
  59                 // container.
  60                 ContainerIPAddress: func() (string, error) { return "0.0.0.0", nil },
  61         }
  62         err := gw.Start()
  63         c.Assert(err, check.IsNil)
  64
  65         rpcconn := rpc.NewConn("",
  66                 &url.URL{
  67                         Scheme: "https",
  68                         Host:   os.Getenv("ARVADOS_API_HOST"),
  69                 },
  70                 true,
  71                 func(context.Context) ([]string, error) {
  72                         return []string{arvadostest.SystemRootToken}, nil
  73                 })
  74         _, err = rpcconn.ContainerUpdate(context.TODO(), arvados.UpdateOptions{UUID: uuid, Attrs: map[string]interface{}{
  75                 "state": arvados.ContainerStateLocked,
  76         }})
  77         c.Assert(err, check.IsNil)
  78         _, err = rpcconn.ContainerUpdate(context.TODO(), arvados.UpdateOptions{UUID: uuid, Attrs: map[string]interface{}{
  79                 "state":           arvados.ContainerStateRunning,
  80                 "gateway_address": gw.Address,
  81         }})
  82         c.Assert(err, check.IsNil)
  83
  84         var stdout, stderr bytes.Buffer
  85         cmd := exec.Command("go", "run", ".", "shell", uuid, "-o", "controlpath=none", "-o", "userknownhostsfile="+c.MkDir()+"/known_hosts", "echo", "ok")
  86         cmd.Env = append(cmd.Env, os.Environ()...)
  87         cmd.Env = append(cmd.Env, "ARVADOS_API_TOKEN="+arvadostest.ActiveTokenV2)
  88         cmd.Stdout = &stdout
  89         cmd.Stderr = &stderr
  90         c.Check(cmd.Run(), check.NotNil)
  91         c.Log(stderr.String())
  92         c.Check(stderr.String(), check.Matches, `(?ms).*(No such container: theperthcountyconspiracy|exec: \"docker\": executable file not found in \$PATH).*`)
  93
  94         // Set up an http server, and try using "arvados-client shell"
  95         // to forward traffic to it.
  96         httpTarget := &httpserver.Server{}
  97         httpTarget.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  98                 c.Logf("httpTarget.Handler: incoming request: %s %s", r.Method, r.URL)
  99                 if r.URL.Path == "/foo" {
 100                         fmt.Fprintln(w, "bar baz")
 101                 } else {
 102                         w.WriteHeader(http.StatusNotFound)
 103                 }
 104         })
 105         err = httpTarget.Start()
 106         c.Assert(err, check.IsNil)
 107
 108         ln, err := net.Listen("tcp", ":0")
 109         c.Assert(err, check.IsNil)
 110         _, forwardedPort, _ := net.SplitHostPort(ln.Addr().String())
 111         ln.Close()
 112
 113         stdout.Reset()
 114         stderr.Reset()
 115         ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(10*time.Second))
 116         defer cancel()
 117         cmd = exec.CommandContext(ctx,
 118                 "go", "run", ".", "shell", uuid,
 119                 "-L", forwardedPort+":"+httpTarget.Addr,
 120                 "-o", "controlpath=none",
 121                 "-o", "userknownhostsfile="+c.MkDir()+"/known_hosts",
 122                 "-N",
 123         )
 124         c.Logf("cmd.Args: %s", cmd.Args)
 125         cmd.Env = append(cmd.Env, os.Environ()...)
 126         cmd.Env = append(cmd.Env, "ARVADOS_API_TOKEN="+arvadostest.ActiveTokenV2)
 127         cmd.Stdout = &stdout
 128         cmd.Stderr = &stderr
 129         go cmd.Run()
 130
 131         forwardedURL := fmt.Sprintf("http://localhost:%s/foo", forwardedPort)
 132
 133         for range time.NewTicker(time.Second / 20).C {
 134                 resp, err := http.Get(forwardedURL)
 135                 if err != nil {
 136                         if !strings.Contains(err.Error(), "connect") {
 137                                 c.Fatal(err)
 138                         } else if ctx.Err() != nil {
 139                                 c.Fatal("timed out")
 140                         }
 141                         // Retry until OpenSSH starts listening
 142                         continue
 143                 }
 144                 c.Check(resp.StatusCode, check.Equals, http.StatusOK)
 145                 body, err := ioutil.ReadAll(resp.Body)
 146                 c.Check(err, check.IsNil)
 147                 c.Check(string(body), check.Equals, "bar baz\n")
 148                 break
 149         }
 150
 151         var wg sync.WaitGroup
 152         for i := 0; i < 10; i++ {
 153                 wg.Add(1)
 154                 go func() {
 155                         defer wg.Done()
 156                         resp, err := http.Get(forwardedURL)
 157                         if !c.Check(err, check.IsNil) {
 158                                 return
 159                         }
 160                         body, err := ioutil.ReadAll(resp.Body)
 161                         c.Check(err, check.IsNil)
 162                         c.Check(string(body), check.Equals, "bar baz\n")
 163                 }()
 164         }
 165         wg.Wait()
 166 }