// Copyright (C) The Arvados Authors. All rights reserved. // // SPDX-License-Identifier: AGPL-3.0 package test import ( "bytes" "fmt" "io" "io/ioutil" "log" "net" "strings" "sync" "golang.org/x/crypto/ssh" check "gopkg.in/check.v1" ) func LoadTestKey(c *check.C, fnm string) (ssh.PublicKey, ssh.Signer) { rawpubkey, err := ioutil.ReadFile(fnm + ".pub") c.Assert(err, check.IsNil) pubkey, _, _, _, err := ssh.ParseAuthorizedKey(rawpubkey) c.Assert(err, check.IsNil) rawprivkey, err := ioutil.ReadFile(fnm) c.Assert(err, check.IsNil) privkey, err := ssh.ParsePrivateKey(rawprivkey) c.Assert(err, check.IsNil) return pubkey, privkey } // An SSHExecFunc handles an "exec" session on a multiplexed SSH // connection. type SSHExecFunc func(env map[string]string, command string, stdin io.Reader, stdout, stderr io.Writer) uint32 // An SSHService accepts SSH connections on an available TCP port and // passes clients' "exec" sessions to the provided SSHExecFunc. type SSHService struct { Exec SSHExecFunc HostKey ssh.Signer AuthorizedUser string AuthorizedKeys []ssh.PublicKey listener net.Listener conn *ssh.ServerConn setup sync.Once mtx sync.Mutex started chan bool closed bool err error } // Address returns the host:port where the SSH server is listening. It // returns "" if called before the server is ready to accept // connections. func (ss *SSHService) Address() string { ss.setup.Do(ss.start) ss.mtx.Lock() ln := ss.listener ss.mtx.Unlock() if ln == nil { return "" } return ln.Addr().String() } // RemoteUser returns the username that will be accepted. func (ss *SSHService) RemoteUser() string { return ss.AuthorizedUser } // Close shuts down the server and releases resources. Established // connections are unaffected. func (ss *SSHService) Close() { ss.Start() ss.mtx.Lock() ln := ss.listener ss.closed = true ss.mtx.Unlock() if ln != nil { ln.Close() } } // Start returns when the server is ready to accept connections. func (ss *SSHService) Start() error { ss.setup.Do(ss.start) <-ss.started return ss.err } func (ss *SSHService) start() { ss.started = make(chan bool) go ss.run() } func (ss *SSHService) run() { defer close(ss.started) config := &ssh.ServerConfig{ PublicKeyCallback: func(c ssh.ConnMetadata, pubKey ssh.PublicKey) (*ssh.Permissions, error) { for _, ak := range ss.AuthorizedKeys { if bytes.Equal(ak.Marshal(), pubKey.Marshal()) { return &ssh.Permissions{}, nil } } return nil, fmt.Errorf("unknown public key for %q", c.User()) }, } config.AddHostKey(ss.HostKey) listener, err := net.Listen("tcp", "127.0.0.1:") if err != nil { ss.err = err return } ss.mtx.Lock() ss.listener = listener ss.mtx.Unlock() go func() { for { nConn, err := listener.Accept() if err != nil && strings.Contains(err.Error(), "use of closed network connection") && ss.closed { return } else if err != nil { log.Printf("accept: %s", err) return } go ss.serveConn(nConn, config) } }() } func (ss *SSHService) serveConn(nConn net.Conn, config *ssh.ServerConfig) { defer nConn.Close() conn, newchans, reqs, err := ssh.NewServerConn(nConn, config) if err != nil { log.Printf("ssh.NewServerConn: %s", err) return } defer conn.Close() go ssh.DiscardRequests(reqs) for newch := range newchans { if newch.ChannelType() != "session" { newch.Reject(ssh.UnknownChannelType, "unknown channel type") continue } ch, reqs, err := newch.Accept() if err != nil { log.Printf("accept channel: %s", err) return } didExec := false sessionEnv := map[string]string{} go func() { for req := range reqs { switch { case didExec: // Reject anything after exec req.Reply(false, nil) case req.Type == "exec": var execReq struct { Command string } req.Reply(true, nil) ssh.Unmarshal(req.Payload, &execReq) go func() { var resp struct { Status uint32 } resp.Status = ss.Exec(sessionEnv, execReq.Command, ch, ch, ch.Stderr()) ch.SendRequest("exit-status", false, ssh.Marshal(&resp)) ch.Close() }() didExec = true case req.Type == "env": var envReq struct { Name string Value string } req.Reply(true, nil) ssh.Unmarshal(req.Payload, &envReq) sessionEnv[envReq.Name] = envReq.Value } } }() } }