--- /dev/null
+// 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 .*`)
+}