17170: Test controller->crunch-run tunnel.
authorTom Clegg <tom@curii.com>
Wed, 13 Jan 2021 04:22:46 +0000 (23:22 -0500)
committerTom Clegg <tom@curii.com>
Wed, 13 Jan 2021 04:22:46 +0000 (23:22 -0500)
Arvados-DCO-1.1-Signed-off-by: Tom Clegg <tom@curii.com>

lib/controller/localdb/container_gateway_test.go [new file with mode: 0644]

diff --git a/lib/controller/localdb/container_gateway_test.go b/lib/controller/localdb/container_gateway_test.go
new file mode 100644 (file)
index 0000000..336d5b6
--- /dev/null
@@ -0,0 +1,123 @@
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
+package localdb
+
+import (
+       "context"
+       "crypto/hmac"
+       "crypto/sha256"
+       "fmt"
+       "io"
+       "time"
+
+       "git.arvados.org/arvados.git/lib/config"
+       "git.arvados.org/arvados.git/lib/crunchrun"
+       "git.arvados.org/arvados.git/sdk/go/arvados"
+       "git.arvados.org/arvados.git/sdk/go/arvadostest"
+       "git.arvados.org/arvados.git/sdk/go/auth"
+       "git.arvados.org/arvados.git/sdk/go/ctxlog"
+       check "gopkg.in/check.v1"
+)
+
+var _ = check.Suite(&ContainerGatewaySuite{})
+
+type ContainerGatewaySuite struct {
+       cluster *arvados.Cluster
+       localdb *Conn
+       ctx     context.Context
+       ctrUUID string
+       gw      *crunchrun.Gateway
+}
+
+func (s *ContainerGatewaySuite) TearDownSuite(c *check.C) {
+       // Undo any changes/additions to the user database so they
+       // don't affect subsequent tests.
+       arvadostest.ResetEnv()
+       c.Check(arvados.NewClientFromEnv().RequestAndDecode(nil, "POST", "database/reset", nil, nil), check.IsNil)
+}
+
+func (s *ContainerGatewaySuite) SetUpSuite(c *check.C) {
+       cfg, err := config.NewLoader(nil, ctxlog.TestLogger(c)).Load()
+       c.Assert(err, check.IsNil)
+       s.cluster, err = cfg.GetCluster("")
+       c.Assert(err, check.IsNil)
+       s.localdb = NewConn(s.cluster)
+       s.ctx = auth.NewContext(context.Background(), &auth.Credentials{Tokens: []string{arvadostest.ActiveTokenV2}})
+
+       s.ctrUUID = arvadostest.QueuedContainerUUID
+
+       h := hmac.New(sha256.New, []byte(s.cluster.SystemRootToken))
+       fmt.Fprint(h, s.ctrUUID)
+       authKey := fmt.Sprintf("%x", h.Sum(nil))
+
+       s.gw = &crunchrun.Gateway{
+               DockerContainerID: new(string),
+               ContainerUUID:     s.ctrUUID,
+               AuthSecret:        authKey,
+               Address:           "localhost:0",
+               Log:               ctxlog.TestLogger(c),
+       }
+       c.Assert(s.gw.Start(), check.IsNil)
+       rootctx := auth.NewContext(context.Background(), &auth.Credentials{Tokens: []string{s.cluster.SystemRootToken}})
+       _, err = s.localdb.ContainerUpdate(rootctx, arvados.UpdateOptions{
+               UUID: s.ctrUUID,
+               Attrs: map[string]interface{}{
+                       "state": arvados.ContainerStateLocked}})
+       c.Assert(err, check.IsNil)
+       _, err = s.localdb.ContainerUpdate(rootctx, arvados.UpdateOptions{
+               UUID: s.ctrUUID,
+               Attrs: map[string]interface{}{
+                       "state":           arvados.ContainerStateRunning,
+                       "gateway_address": s.gw.Address}})
+       c.Assert(err, check.IsNil)
+}
+
+func (s *ContainerGatewaySuite) TestConnect(c *check.C) {
+       c.Logf("connecting to %s", s.gw.Address)
+       sshconn, err := s.localdb.ContainerSSH(s.ctx, arvados.ContainerSSHOptions{UUID: s.ctrUUID})
+       c.Assert(err, check.IsNil)
+       c.Assert(sshconn.Conn, check.NotNil)
+       defer sshconn.Conn.Close()
+
+       done := make(chan struct{})
+       go func() {
+               defer close(done)
+
+               // Receive text banner
+               buf := make([]byte, 12)
+               _, err := io.ReadFull(sshconn.Conn, buf)
+               c.Check(err, check.IsNil)
+               c.Check(string(buf), check.Equals, "SSH-2.0-Go\r\n")
+
+               // Send text banner
+               _, err = sshconn.Conn.Write([]byte("SSH-2.0-Fake\r\n"))
+               c.Check(err, check.IsNil)
+
+               // Receive binary
+               _, err = io.ReadFull(sshconn.Conn, buf[:4])
+               c.Check(err, check.IsNil)
+               c.Check(buf[:4], check.DeepEquals, []byte{0, 0, 1, 0xfc})
+
+               // If we can get this far into an SSH handshake...
+               c.Log("success, tunnel is working")
+       }()
+       select {
+       case <-done:
+       case <-time.After(time.Second):
+               c.Fail()
+       }
+}
+
+func (s *ContainerGatewaySuite) TestConnectFail(c *check.C) {
+       c.Log("trying with no token")
+       ctx := auth.NewContext(context.Background(), &auth.Credentials{})
+       _, err := s.localdb.ContainerSSH(ctx, arvados.ContainerSSHOptions{UUID: s.ctrUUID})
+       c.Check(err, check.ErrorMatches, `.* 401 .*`)
+
+       c.Log("trying with anonymous token")
+       ctx = auth.NewContext(context.Background(), &auth.Credentials{Tokens: []string{arvadostest.AnonymousToken}})
+       _, err = s.localdb.ContainerSSH(ctx, arvados.ContainerSSHOptions{UUID: s.ctrUUID})
+       c.Check(err, check.ErrorMatches, `.* 404 .*`)
+}