github.com/secure-build/gitlab-runner@v12.5.0+incompatible/executors/ssh/stub_ssh_server.go (about)

     1  package ssh
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"strconv"
     7  	"strings"
     8  
     9  	"github.com/tevino/abool"
    10  	cryptoSSH "golang.org/x/crypto/ssh"
    11  )
    12  
    13  type StubSSHServer struct {
    14  	User     string
    15  	Password string
    16  	Config   *cryptoSSH.ServerConfig
    17  
    18  	stop       chan bool
    19  	shouldExit *abool.AtomicBool
    20  }
    21  
    22  func NewStubServer(user, pass string, privateKey []byte) (*StubSSHServer, error) {
    23  	server := &StubSSHServer{
    24  		User:     user,
    25  		Password: pass,
    26  		Config: &cryptoSSH.ServerConfig{
    27  			PasswordCallback: func(conn cryptoSSH.ConnMetadata, password []byte) (*cryptoSSH.Permissions, error) {
    28  				if conn.User() == user && string(password) == pass {
    29  					return nil, nil
    30  				}
    31  				return nil, fmt.Errorf("wrong password for %q", conn.User())
    32  			},
    33  		},
    34  		stop:       make(chan bool),
    35  		shouldExit: abool.New(),
    36  	}
    37  
    38  	key, err := cryptoSSH.ParsePrivateKey(privateKey)
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  
    43  	server.Config.AddHostKey(key)
    44  
    45  	return server, nil
    46  }
    47  
    48  func (s *StubSSHServer) Start() (int, error) {
    49  	listener, err := net.Listen("tcp", "127.0.0.1:")
    50  	if err != nil {
    51  		return 0, err
    52  	}
    53  
    54  	go func() {
    55  		<-s.stop
    56  		s.shouldExit.Set()
    57  		listener.Close()
    58  	}()
    59  
    60  	address := strings.SplitN(listener.Addr().String(), ":", 2)
    61  	go s.mainLoop(listener)
    62  
    63  	return strconv.Atoi(address[1])
    64  }
    65  
    66  func (s *StubSSHServer) Stop() {
    67  	s.stop <- true
    68  }
    69  
    70  func (s *StubSSHServer) mainLoop(listener net.Listener) {
    71  	for {
    72  		if s.shouldExit.IsSet() {
    73  			return
    74  		}
    75  
    76  		conn, err := listener.Accept()
    77  		if err != nil {
    78  			continue
    79  		}
    80  
    81  		if s.shouldExit.IsSet() {
    82  			return
    83  		}
    84  
    85  		//upgrade to ssh connection
    86  		cryptoSSH.NewServerConn(conn, s.Config)
    87  		// This is enough just for handling incoming connections
    88  	}
    89  }