github.com/rudderlabs/rudder-go-kit@v0.30.0/testhelper/docker/resource/sshserver/sshserver.go (about)

     1  package sshserver
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"strings"
     7  	"time"
     8  
     9  	"github.com/ory/dockertest/v3"
    10  	dc "github.com/ory/dockertest/v3/docker"
    11  
    12  	kithelper "github.com/rudderlabs/rudder-go-kit/testhelper"
    13  	"github.com/rudderlabs/rudder-go-kit/testhelper/docker/resource"
    14  )
    15  
    16  const exposedPort = "2222"
    17  
    18  type Option func(*config)
    19  
    20  type config struct {
    21  	publicKeyPath      string
    22  	username, password string
    23  	network            *dc.Network
    24  }
    25  
    26  // WithCredentials sets the username and password to use for the SSH server.
    27  func WithCredentials(username, password string) Option {
    28  	return func(c *config) {
    29  		c.username = username
    30  		c.password = password
    31  	}
    32  }
    33  
    34  // WithPublicKeyPath sets the public key path to use for the SSH server.
    35  func WithPublicKeyPath(publicKeyPath string) Option {
    36  	return func(c *config) {
    37  		c.publicKeyPath = publicKeyPath
    38  	}
    39  }
    40  
    41  // WithDockerNetwork sets the Docker network to use for the SSH server.
    42  func WithDockerNetwork(network *dc.Network) Option {
    43  	return func(c *config) {
    44  		c.network = network
    45  	}
    46  }
    47  
    48  type Resource struct {
    49  	Port int
    50  
    51  	container *dockertest.Resource
    52  }
    53  
    54  func Setup(pool *dockertest.Pool, cln resource.Cleaner, opts ...Option) (*Resource, error) {
    55  	var c config
    56  	for _, opt := range opts {
    57  		opt(&c)
    58  	}
    59  
    60  	network := c.network
    61  	if c.network == nil {
    62  		var err error
    63  		network, err = pool.Client.CreateNetwork(dc.CreateNetworkOptions{Name: "sshserver_network"})
    64  		if err != nil {
    65  			return nil, fmt.Errorf("could not create docker network: %w", err)
    66  		}
    67  		cln.Cleanup(func() {
    68  			if err := pool.Client.RemoveNetwork(network.ID); err != nil {
    69  				cln.Log(fmt.Sprintf("could not remove sshserver_network: %v", err))
    70  			}
    71  		})
    72  	}
    73  
    74  	port, err := kithelper.GetFreePort()
    75  	if err != nil {
    76  		return nil, err
    77  	}
    78  
    79  	var (
    80  		mounts  []string
    81  		envVars = []string{
    82  			"SUDO_ACCESS=false",
    83  			"DOCKER_MODS=linuxserver/mods:openssh-server-ssh-tunnel",
    84  		}
    85  	)
    86  	if c.username != "" {
    87  		envVars = append(envVars, "USER_NAME="+c.username)
    88  		if c.password != "" {
    89  			envVars = append(envVars, []string{
    90  				"USER_PASSWORD=" + c.password,
    91  				"PASSWORD_ACCESS=true",
    92  			}...)
    93  		}
    94  	}
    95  	if c.publicKeyPath != "" {
    96  		envVars = append(envVars, "PUBLIC_KEY_FILE=/test_key.pub")
    97  		mounts = []string{c.publicKeyPath + ":/test_key.pub"}
    98  	}
    99  	container, err := pool.RunWithOptions(&dockertest.RunOptions{
   100  		Repository: "lscr.io/linuxserver/openssh-server",
   101  		Tag:        "9.3_p2-r1-ls145",
   102  		NetworkID:  network.ID,
   103  		Hostname:   "sshserver",
   104  		PortBindings: map[dc.Port][]dc.PortBinding{
   105  			exposedPort + "/tcp": {
   106  				{HostIP: "sshserver", HostPort: fmt.Sprintf("%d", port)},
   107  			},
   108  		},
   109  		Env:    envVars,
   110  		Mounts: mounts,
   111  	})
   112  	if err != nil {
   113  		return nil, err
   114  	}
   115  	cln.Cleanup(func() {
   116  		if err := pool.Purge(container); err != nil {
   117  			cln.Log("Could not purge resource", err)
   118  		}
   119  	})
   120  
   121  	var (
   122  		buf     *bytes.Buffer
   123  		timeout = time.After(30 * time.Second)
   124  		ticker  = time.NewTicker(200 * time.Millisecond)
   125  	)
   126  loop:
   127  	for {
   128  		select {
   129  		case <-ticker.C:
   130  			buf = bytes.NewBuffer(nil)
   131  			exitCode, err := container.Exec([]string{"cat", "/config/logs/openssh/current"}, dockertest.ExecOptions{
   132  				StdOut: buf,
   133  			})
   134  			if err != nil {
   135  				cln.Log("could not exec into SSH server:", err)
   136  				continue
   137  			}
   138  			if exitCode != 0 {
   139  				cln.Log("invalid exit code while exec-ing into SSH server:", exitCode)
   140  				continue
   141  			}
   142  			if buf.String() == "" {
   143  				cln.Log("SSH server not ready yet")
   144  				continue
   145  			}
   146  			if !strings.Contains(buf.String(), "Server listening on :: port "+exposedPort) {
   147  				cln.Log("SSH server not listening on port yet")
   148  				continue
   149  			}
   150  			cln.Log("SSH server is ready:", exposedPort, "=>", container.GetPort(exposedPort+"/tcp"))
   151  			break loop
   152  		case <-timeout:
   153  			return nil, fmt.Errorf("ssh server not health within timeout")
   154  		}
   155  	}
   156  
   157  	return &Resource{
   158  		Port:      port,
   159  		container: container,
   160  	}, nil
   161  }