lib/dispatchcloud/ssh_executor/executor_test.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package ssh_executor
   6
   7 import (
   8         "bytes"
   9         "fmt"
  10         "io"
  11         "io/ioutil"
  12         "net"
  13         "sync"
  14         "testing"
  15         "time"
  16
  17         "git.curoverse.com/arvados.git/lib/dispatchcloud/test"
  18         "golang.org/x/crypto/ssh"
  19         check "gopkg.in/check.v1"
  20 )
  21
  22 // Gocheck boilerplate
  23 func Test(t *testing.T) {
  24         check.TestingT(t)
  25 }
  26
  27 var _ = check.Suite(&ExecutorSuite{})
  28
  29 type testTarget struct {
  30         test.SSHService
  31 }
  32
  33 func (*testTarget) VerifyHostKey(ssh.PublicKey, *ssh.Client) error {
  34         return nil
  35 }
  36
  37 // Address returns the wrapped SSHService's host, with the port
  38 // stripped. This ensures the executor won't work until
  39 // SetTargetPort() is called -- see (*testTarget)Port().
  40 func (tt *testTarget) Address() string {
  41         h, _, err := net.SplitHostPort(tt.SSHService.Address())
  42         if err != nil {
  43                 panic(err)
  44         }
  45         return h
  46 }
  47
  48 func (tt *testTarget) Port() string {
  49         _, p, err := net.SplitHostPort(tt.SSHService.Address())
  50         if err != nil {
  51                 panic(err)
  52         }
  53         return p
  54 }
  55
  56 type mitmTarget struct {
  57         test.SSHService
  58 }
  59
  60 func (*mitmTarget) VerifyHostKey(key ssh.PublicKey, client *ssh.Client) error {
  61         return fmt.Errorf("host key failed verification: %#v", key)
  62 }
  63
  64 type ExecutorSuite struct{}
  65
  66 func (s *ExecutorSuite) TestBadHostKey(c *check.C) {
  67         _, hostpriv := test.LoadTestKey(c, "../test/sshkey_vm")
  68         clientpub, clientpriv := test.LoadTestKey(c, "../test/sshkey_dispatch")
  69         target := &mitmTarget{
  70                 SSHService: test.SSHService{
  71                         Exec: func(map[string]string, string, io.Reader, io.Writer, io.Writer) uint32 {
  72                                 c.Error("Target Exec func called even though host key verification failed")
  73                                 return 0
  74                         },
  75                         HostKey:        hostpriv,
  76                         AuthorizedKeys: []ssh.PublicKey{clientpub},
  77                 },
  78         }
  79
  80         err := target.Start()
  81         c.Check(err, check.IsNil)
  82         c.Logf("target address %q", target.Address())
  83         defer target.Close()
  84
  85         exr := New(target)
  86         exr.SetSigners(clientpriv)
  87
  88         _, _, err = exr.Execute(nil, "true", nil)
  89         c.Check(err, check.ErrorMatches, "host key failed verification: .*")
  90 }
  91
  92 func (s *ExecutorSuite) TestExecute(c *check.C) {
  93         command := `foo 'bar' "baz"`
  94         stdinData := "foobar\nbaz\n"
  95         _, hostpriv := test.LoadTestKey(c, "../test/sshkey_vm")
  96         clientpub, clientpriv := test.LoadTestKey(c, "../test/sshkey_dispatch")
  97         for _, exitcode := range []int{0, 1, 2} {
  98                 target := &testTarget{
  99                         SSHService: test.SSHService{
 100                                 Exec: func(env map[string]string, cmd string, stdin io.Reader, stdout, stderr io.Writer) uint32 {
 101                                         c.Check(env["TESTVAR"], check.Equals, "test value")
 102                                         c.Check(cmd, check.Equals, command)
 103                                         var wg sync.WaitGroup
 104                                         wg.Add(2)
 105                                         go func() {
 106                                                 io.WriteString(stdout, "stdout\n")
 107                                                 wg.Done()
 108                                         }()
 109                                         go func() {
 110                                                 io.WriteString(stderr, "stderr\n")
 111                                                 wg.Done()
 112                                         }()
 113                                         buf, err := ioutil.ReadAll(stdin)
 114                                         wg.Wait()
 115                                         c.Check(err, check.IsNil)
 116                                         if err != nil {
 117                                                 return 99
 118                                         }
 119                                         _, err = stdout.Write(buf)
 120                                         c.Check(err, check.IsNil)
 121                                         return uint32(exitcode)
 122                                 },
 123                                 HostKey:        hostpriv,
 124                                 AuthorizedKeys: []ssh.PublicKey{clientpub},
 125                         },
 126                 }
 127                 err := target.Start()
 128                 c.Check(err, check.IsNil)
 129                 c.Logf("target address %q", target.Address())
 130                 defer target.Close()
 131
 132                 exr := New(target)
 133                 exr.SetSigners(clientpriv)
 134
 135                 // Use the default target port (ssh). Execute will
 136                 // return a connection error or an authentication
 137                 // error, depending on whether the test host is
 138                 // running an SSH server.
 139                 _, _, err = exr.Execute(nil, command, nil)
 140                 c.Check(err, check.ErrorMatches, `.*(unable to authenticate|connection refused).*`)
 141
 142                 // Use a bogus target port. Execute will return a
 143                 // connection error.
 144                 exr.SetTargetPort("0")
 145                 _, _, err = exr.Execute(nil, command, nil)
 146                 c.Check(err, check.ErrorMatches, `.*connection refused.*`)
 147
 148                 // Use the test server's listening port.
 149                 exr.SetTargetPort(target.Port())
 150
 151                 done := make(chan bool)
 152                 go func() {
 153                         stdout, stderr, err := exr.Execute(map[string]string{"TESTVAR": "test value"}, command, bytes.NewBufferString(stdinData))
 154                         if exitcode == 0 {
 155                                 c.Check(err, check.IsNil)
 156                         } else {
 157                                 c.Check(err, check.NotNil)
 158                                 err, ok := err.(*ssh.ExitError)
 159                                 c.Assert(ok, check.Equals, true)
 160                                 c.Check(err.ExitStatus(), check.Equals, exitcode)
 161                         }
 162                         c.Check(stdout, check.DeepEquals, []byte("stdout\n"+stdinData))
 163                         c.Check(stderr, check.DeepEquals, []byte("stderr\n"))
 164                         close(done)
 165                 }()
 166
 167                 timeout := time.NewTimer(time.Second)
 168                 select {
 169                 case <-done:
 170                 case <-timeout.C:
 171                         c.Fatal("timed out")
 172                 }
 173         }
 174 }