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