lib/controller/localdb/container_gateway_test.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package localdb
   6
   7 import (
   8         "context"
   9         "crypto/hmac"
  10         "crypto/sha256"
  11         "fmt"
  12         "io"
  13         "time"
  14
  15         "git.arvados.org/arvados.git/lib/config"
  16         "git.arvados.org/arvados.git/lib/crunchrun"
  17         "git.arvados.org/arvados.git/sdk/go/arvados"
  18         "git.arvados.org/arvados.git/sdk/go/arvadostest"
  19         "git.arvados.org/arvados.git/sdk/go/auth"
  20         "git.arvados.org/arvados.git/sdk/go/ctxlog"
  21         check "gopkg.in/check.v1"
  22 )
  23
  24 var _ = check.Suite(&ContainerGatewaySuite{})
  25
  26 type ContainerGatewaySuite struct {
  27         cluster *arvados.Cluster
  28         localdb *Conn
  29         ctx     context.Context
  30         ctrUUID string
  31         gw      *crunchrun.Gateway
  32 }
  33
  34 func (s *ContainerGatewaySuite) TearDownSuite(c *check.C) {
  35         // Undo any changes/additions to the user database so they
  36         // don't affect subsequent tests.
  37         arvadostest.ResetEnv()
  38         c.Check(arvados.NewClientFromEnv().RequestAndDecode(nil, "POST", "database/reset", nil, nil), check.IsNil)
  39 }
  40
  41 func (s *ContainerGatewaySuite) SetUpSuite(c *check.C) {
  42         cfg, err := config.NewLoader(nil, ctxlog.TestLogger(c)).Load()
  43         c.Assert(err, check.IsNil)
  44         s.cluster, err = cfg.GetCluster("")
  45         c.Assert(err, check.IsNil)
  46         s.localdb = NewConn(s.cluster)
  47         s.ctx = auth.NewContext(context.Background(), &auth.Credentials{Tokens: []string{arvadostest.ActiveTokenV2}})
  48
  49         s.ctrUUID = arvadostest.QueuedContainerUUID
  50
  51         h := hmac.New(sha256.New, []byte(s.cluster.SystemRootToken))
  52         fmt.Fprint(h, s.ctrUUID)
  53         authKey := fmt.Sprintf("%x", h.Sum(nil))
  54
  55         s.gw = &crunchrun.Gateway{
  56                 DockerContainerID: new(string),
  57                 ContainerUUID:     s.ctrUUID,
  58                 AuthSecret:        authKey,
  59                 Address:           "localhost:0",
  60                 Log:               ctxlog.TestLogger(c),
  61         }
  62         c.Assert(s.gw.Start(), check.IsNil)
  63         rootctx := auth.NewContext(context.Background(), &auth.Credentials{Tokens: []string{s.cluster.SystemRootToken}})
  64         _, err = s.localdb.ContainerUpdate(rootctx, arvados.UpdateOptions{
  65                 UUID: s.ctrUUID,
  66                 Attrs: map[string]interface{}{
  67                         "state": arvados.ContainerStateLocked}})
  68         c.Assert(err, check.IsNil)
  69         _, err = s.localdb.ContainerUpdate(rootctx, arvados.UpdateOptions{
  70                 UUID: s.ctrUUID,
  71                 Attrs: map[string]interface{}{
  72                         "state":           arvados.ContainerStateRunning,
  73                         "gateway_address": s.gw.Address}})
  74         c.Assert(err, check.IsNil)
  75 }
  76
  77 func (s *ContainerGatewaySuite) SetUpTest(c *check.C) {
  78         s.cluster.Containers.ShellAccess.Admin = true
  79         s.cluster.Containers.ShellAccess.User = true
  80         _, err := arvadostest.DB(c, s.cluster).Exec(`update containers set interactive_session_started=$1 where uuid=$2`, false, s.ctrUUID)
  81         c.Check(err, check.IsNil)
  82 }
  83
  84 func (s *ContainerGatewaySuite) TestConfig(c *check.C) {
  85         for _, trial := range []struct {
  86                 configAdmin bool
  87                 configUser  bool
  88                 sendToken   string
  89                 errorCode   int
  90         }{
  91                 {true, true, arvadostest.ActiveTokenV2, 0},
  92                 {true, false, arvadostest.ActiveTokenV2, 503},
  93                 {false, true, arvadostest.ActiveTokenV2, 0},
  94                 {false, false, arvadostest.ActiveTokenV2, 503},
  95                 {true, true, arvadostest.AdminToken, 0},
  96                 {true, false, arvadostest.AdminToken, 0},
  97                 {false, true, arvadostest.AdminToken, 403},
  98                 {false, false, arvadostest.AdminToken, 503},
  99         } {
 100                 c.Logf("trial %#v", trial)
 101                 s.cluster.Containers.ShellAccess.Admin = trial.configAdmin
 102                 s.cluster.Containers.ShellAccess.User = trial.configUser
 103                 ctx := auth.NewContext(s.ctx, &auth.Credentials{Tokens: []string{trial.sendToken}})
 104                 sshconn, err := s.localdb.ContainerSSH(ctx, arvados.ContainerSSHOptions{UUID: s.ctrUUID})
 105                 if trial.errorCode == 0 {
 106                         if !c.Check(err, check.IsNil) {
 107                                 continue
 108                         }
 109                         if !c.Check(sshconn.Conn, check.NotNil) {
 110                                 continue
 111                         }
 112                         sshconn.Conn.Close()
 113                 } else {
 114                         c.Check(err, check.NotNil)
 115                         err, ok := err.(interface{ HTTPStatus() int })
 116                         if c.Check(ok, check.Equals, true) {
 117                                 c.Check(err.HTTPStatus(), check.Equals, trial.errorCode)
 118                         }
 119                 }
 120         }
 121 }
 122
 123 func (s *ContainerGatewaySuite) TestConnect(c *check.C) {
 124         c.Logf("connecting to %s", s.gw.Address)
 125         sshconn, err := s.localdb.ContainerSSH(s.ctx, arvados.ContainerSSHOptions{UUID: s.ctrUUID})
 126         c.Assert(err, check.IsNil)
 127         c.Assert(sshconn.Conn, check.NotNil)
 128         defer sshconn.Conn.Close()
 129
 130         done := make(chan struct{})
 131         go func() {
 132                 defer close(done)
 133
 134                 // Receive text banner
 135                 buf := make([]byte, 12)
 136                 _, err := io.ReadFull(sshconn.Conn, buf)
 137                 c.Check(err, check.IsNil)
 138                 c.Check(string(buf), check.Equals, "SSH-2.0-Go\r\n")
 139
 140                 // Send text banner
 141                 _, err = sshconn.Conn.Write([]byte("SSH-2.0-Fake\r\n"))
 142                 c.Check(err, check.IsNil)
 143
 144                 // Receive binary
 145                 _, err = io.ReadFull(sshconn.Conn, buf[:4])
 146                 c.Check(err, check.IsNil)
 147                 c.Check(buf[:4], check.DeepEquals, []byte{0, 0, 1, 0xfc})
 148
 149                 // If we can get this far into an SSH handshake...
 150                 c.Log("success, tunnel is working")
 151         }()
 152         select {
 153         case <-done:
 154         case <-time.After(time.Second):
 155                 c.Fail()
 156         }
 157         ctr, err := s.localdb.ContainerGet(s.ctx, arvados.GetOptions{UUID: s.ctrUUID})
 158         c.Check(err, check.IsNil)
 159         c.Check(ctr.InteractiveSessionStarted, check.Equals, true)
 160 }
 161
 162 func (s *ContainerGatewaySuite) TestConnectFail(c *check.C) {
 163         c.Log("trying with no token")
 164         ctx := auth.NewContext(context.Background(), &auth.Credentials{})
 165         _, err := s.localdb.ContainerSSH(ctx, arvados.ContainerSSHOptions{UUID: s.ctrUUID})
 166         c.Check(err, check.ErrorMatches, `.* 401 .*`)
 167
 168         c.Log("trying with anonymous token")
 169         ctx = auth.NewContext(context.Background(), &auth.Credentials{Tokens: []string{arvadostest.AnonymousToken}})
 170         _, err = s.localdb.ContainerSSH(ctx, arvados.ContainerSSHOptions{UUID: s.ctrUUID})
 171         c.Check(err, check.ErrorMatches, `.* 404 .*`)
 172 }