1 files changed, 106 insertions, 14 deletions

diff --git a/internal/sshd/sshd_test.go b/internal/sshd/sshd_test.go
index e5f6111..2923737 100644
--- a/internal/sshd/sshd_test.go
+++ b/internal/sshd/sshd_test.go
@@ -2,37 +2,71 @@ package sshd
 
 import (
 	"context"
+	"fmt"
+	"io/ioutil"
+	"net/http"
 	"net/http/httptest"
 	"path"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
+	"golang.org/x/crypto/ssh"
 
 	"gitlab.com/gitlab-org/gitlab-shell/client/testserver"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/testhelper"
 )
 
-const serverUrl = "127.0.0.1:50000"
-
-func TestShutdown(t *testing.T) {
-	s := setupServer(t)
+const (
+	serverUrl = "127.0.0.1:50000"
+	user      = "git"
+)
 
-	go func() { require.NoError(t, s.ListenAndServe(context.Background())) }()
+var (
+	correlationId = ""
+)
 
-	verifyStatus(t, s, StatusReady)
+func TestListenAndServe(t *testing.T) {
+	s := setupServer(t)
 
-	s.wg.Add(1)
+	client, err := ssh.Dial("tcp", serverUrl, clientConfig(t))
+	require.NoError(t, err)
+	defer client.Close()
 
 	require.NoError(t, s.Shutdown())
 	verifyStatus(t, s, StatusOnShutdown)
 
-	s.wg.Done()
+	holdSession(t, client)
+
+	_, err = ssh.Dial("tcp", serverUrl, clientConfig(t))
+	require.Equal(t, err.Error(), "dial tcp 127.0.0.1:50000: connect: connection refused")
+
+	client.Close()
 
 	verifyStatus(t, s, StatusClosed)
 }
 
+func TestCorrelationId(t *testing.T) {
+	setupServer(t)
+
+	client, err := ssh.Dial("tcp", serverUrl, clientConfig(t))
+	require.NoError(t, err)
+	defer client.Close()
+
+	holdSession(t, client)
+
+	previousCorrelationId := correlationId
+
+	client, err = ssh.Dial("tcp", serverUrl, clientConfig(t))
+	require.NoError(t, err)
+	defer client.Close()
+
+	holdSession(t, client)
+
+	require.NotEqual(t, previousCorrelationId, correlationId)
+}
+
 func TestReadinessProbe(t *testing.T) {
 	s := &Server{Config: &config.Config{Server: config.DefaultServerConfig}}
 
@@ -71,17 +105,75 @@ func TestLivenessProbe(t *testing.T) {
 }
 
 func setupServer(t *testing.T) *Server {
+	t.Helper()
+
+	requests := []testserver.TestRequestHandler{
+		{
+			Path: "/api/v4/internal/authorized_keys",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				correlationId = r.Header.Get("X-Request-Id")
+
+				require.NotEmpty(t, correlationId)
+
+				fmt.Fprint(w, `{"id": 1000, "key": "key"}`)
+			},
+		}, {
+			Path: "/api/v4/internal/discover",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				require.Equal(t, correlationId, r.Header.Get("X-Request-Id"))
+
+				fmt.Fprint(w, `{"id": 1000, "name": "Test User", "username": "test-user"}`)
+			},
+		},
+	}
+
 	testhelper.PrepareTestRootDir(t)
 
-	url := testserver.StartSocketHttpServer(t, []testserver.TestRequestHandler{})
+	url := testserver.StartSocketHttpServer(t, requests)
 	srvCfg := config.ServerConfig{
-		Listen:       serverUrl,
-		HostKeyFiles: []string{path.Join(testhelper.TestRoot, "certs/valid/server.key")},
+		Listen:                  serverUrl,
+		ConcurrentSessionsLimit: 1,
+		HostKeyFiles:            []string{path.Join(testhelper.TestRoot, "certs/valid/server.key")},
+	}
+
+	s, err := NewServer(&config.Config{User: user, RootDir: "/tmp", GitlabUrl: url, Server: srvCfg})
+	require.NoError(t, err)
+
+	go func() { require.NoError(t, s.ListenAndServe(context.Background())) }()
+	t.Cleanup(func() { s.Shutdown() })
+
+	verifyStatus(t, s, StatusReady)
+
+	return s
+}
+
+func clientConfig(t *testing.T) *ssh.ClientConfig {
+	keyRaw, err := ioutil.ReadFile(path.Join(testhelper.TestRoot, "certs/valid/server_authorized_key"))
+	pKey, _, _, _, err := ssh.ParseAuthorizedKey(keyRaw)
+	require.NoError(t, err)
+
+	key, err := ioutil.ReadFile(path.Join(testhelper.TestRoot, "certs/client/key.pem"))
+	require.NoError(t, err)
+	signer, err := ssh.ParsePrivateKey(key)
+	require.NoError(t, err)
+
+	return &ssh.ClientConfig{
+		User: user,
+		Auth: []ssh.AuthMethod{
+			ssh.PublicKeys(signer),
+		},
+		HostKeyCallback: ssh.FixedHostKey(pKey),
 	}
+}
 
-	cfg := &config.Config{RootDir: "/tmp", GitlabUrl: url, Server: srvCfg}
+func holdSession(t *testing.T, c *ssh.Client) {
+	session, err := c.NewSession()
+	require.NoError(t, err)
+	defer session.Close()
 
-	return &Server{Config: cfg}
+	output, err := session.Output("discover")
+	require.NoError(t, err)
+	require.Equal(t, "Welcome to GitLab, @test-user!\n", string(output))
 }
 
 func verifyStatus(t *testing.T, s *Server, st status) {
@@ -94,5 +186,5 @@ func verifyStatus(t *testing.T, s *Server, st status) {
 		time.Sleep(time.Duration(i) * time.Millisecond)
 	}
 
-	require.Equal(t, s.getStatus(), st)
+	require.Equal(t, st, s.getStatus())
 }