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) TestConnect(c *check.C) {
  78         c.Logf("connecting to %s", s.gw.Address)
  79         sshconn, err := s.localdb.ContainerSSH(s.ctx, arvados.ContainerSSHOptions{UUID: s.ctrUUID})
  80         c.Assert(err, check.IsNil)
  81         c.Assert(sshconn.Conn, check.NotNil)
  82         defer sshconn.Conn.Close()
  83
  84         done := make(chan struct{})
  85         go func() {
  86                 defer close(done)
  87
  88                 // Receive text banner
  89                 buf := make([]byte, 12)
  90                 _, err := io.ReadFull(sshconn.Conn, buf)
  91                 c.Check(err, check.IsNil)
  92                 c.Check(string(buf), check.Equals, "SSH-2.0-Go\r\n")
  93
  94                 // Send text banner
  95                 _, err = sshconn.Conn.Write([]byte("SSH-2.0-Fake\r\n"))
  96                 c.Check(err, check.IsNil)
  97
  98                 // Receive binary
  99                 _, err = io.ReadFull(sshconn.Conn, buf[:4])
 100                 c.Check(err, check.IsNil)
 101                 c.Check(buf[:4], check.DeepEquals, []byte{0, 0, 1, 0xfc})
 102
 103                 // If we can get this far into an SSH handshake...
 104                 c.Log("success, tunnel is working")
 105         }()
 106         select {
 107         case <-done:
 108         case <-time.After(time.Second):
 109                 c.Fail()
 110         }
 111 }
 112
 113 func (s *ContainerGatewaySuite) TestConnectFail(c *check.C) {
 114         c.Log("trying with no token")
 115         ctx := auth.NewContext(context.Background(), &auth.Credentials{})
 116         _, err := s.localdb.ContainerSSH(ctx, arvados.ContainerSSHOptions{UUID: s.ctrUUID})
 117         c.Check(err, check.ErrorMatches, `.* 401 .*`)
 118
 119         c.Log("trying with anonymous token")
 120         ctx = auth.NewContext(context.Background(), &auth.Credentials{Tokens: []string{arvadostest.AnonymousToken}})
 121         _, err = s.localdb.ContainerSSH(ctx, arvados.ContainerSSHOptions{UUID: s.ctrUUID})
 122         c.Check(err, check.ErrorMatches, `.* 404 .*`)
 123 }