Merge branch 'id-implement-client-keep-alive' into 'main'

Implement ClientKeepAlive option

See merge request gitlab-org/gitlab-shell!622

Author: Patrick Bajao

config.yml.example

@@ -76,6 +76,8 @@ sshd:
   web_listen: "localhost:9122"
   # Maximum number of concurrent sessions allowed on a single SSH connection. Defaults to 10.
   concurrent_sessions_limit: 10
+  # Sets an interval after which server will send keepalive message to a client
+  client_alive_interval: 15
   # The server waits for this time (in seconds) for the ongoing connections to complete before shutting down. Defaults to 10.
   grace_period: 10
   # The endpoint that returns 200 OK if the server is ready to receive incoming connections; otherwise, it returns 503 Service Unavailable. Defaults to "/start".

internal/config/config.go

@@ -22,15 +22,16 @@ const (
 )
 
 type ServerConfig struct {
-	Listen                  string   `yaml:"listen,omitempty"`
-	ProxyProtocol           bool     `yaml:"proxy_protocol,omitempty"`
-	ProxyPolicy             string   `yaml:"proxy_policy,omitempty"`
-	WebListen               string   `yaml:"web_listen,omitempty"`
-	ConcurrentSessionsLimit int64    `yaml:"concurrent_sessions_limit,omitempty"`
-	GracePeriodSeconds      uint64   `yaml:"grace_period"`
-	ReadinessProbe          string   `yaml:"readiness_probe"`
-	LivenessProbe           string   `yaml:"liveness_probe"`
-	HostKeyFiles            []string `yaml:"host_key_files,omitempty"`
+	Listen                     string   `yaml:"listen,omitempty"`
+	ProxyProtocol              bool     `yaml:"proxy_protocol,omitempty"`
+	ProxyPolicy                string   `yaml:"proxy_policy,omitempty"`
+	WebListen                  string   `yaml:"web_listen,omitempty"`
+	ConcurrentSessionsLimit    int64    `yaml:"concurrent_sessions_limit,omitempty"`
+	ClientAliveIntervalSeconds int64    `yaml:"client_alive_interval,omitempty"`
+	GracePeriodSeconds         uint64   `yaml:"grace_period"`
+	ReadinessProbe             string   `yaml:"readiness_probe"`
+	LivenessProbe              string   `yaml:"liveness_probe"`
+	HostKeyFiles               []string `yaml:"host_key_files,omitempty"`
 }
 
 type HttpSettingsConfig struct {
@@ -75,12 +76,13 @@ var (
 	}
 
 	DefaultServerConfig = ServerConfig{
-		Listen:                  "[::]:22",
-		WebListen:               "localhost:9122",
-		ConcurrentSessionsLimit: 10,
-		GracePeriodSeconds:      10,
-		ReadinessProbe:          "/start",
-		LivenessProbe:           "/health",
+		Listen:                     "[::]:22",
+		WebListen:                  "localhost:9122",
+		ConcurrentSessionsLimit:    10,
+		GracePeriodSeconds:         10,
+		ClientAliveIntervalSeconds: 15,
+		ReadinessProbe:             "/start",
+		LivenessProbe:              "/health",
 		HostKeyFiles: []string{
 			"/run/secrets/ssh-hostkeys/ssh_host_rsa_key",
 			"/run/secrets/ssh-hostkeys/ssh_host_ecdsa_key",
@@ -89,6 +91,10 @@ var (
 	}
 )
 
+func (sc *ServerConfig) ClientAliveInterval() time.Duration {
+	return time.Duration(sc.ClientAliveIntervalSeconds) * time.Second
+}
+
 func (sc *ServerConfig) GracePeriod() time.Duration {
 	return time.Duration(sc.GracePeriodSeconds) * time.Second
 }

internal/sshd/connection.go

@@ -7,28 +7,41 @@ import (
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/sync/semaphore"
 
+	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/metrics"
 
 	"gitlab.com/gitlab-org/labkit/log"
 )
 
+const KeepAliveMsg = "[email protected]"
+
 type connection struct {
+	cfg                *config.Config
 	concurrentSessions *semaphore.Weighted
 	remoteAddr         string
+	sconn              *ssh.ServerConn
 }
 
 type channelHandler func(context.Context, ssh.Channel, <-chan *ssh.Request)
 
-func newConnection(maxSessions int64, remoteAddr string) *connection {
+func newConnection(cfg *config.Config, remoteAddr string, sconn *ssh.ServerConn) *connection {
 	return &connection{
-		concurrentSessions: semaphore.NewWeighted(maxSessions),
+		cfg:                cfg,
+		concurrentSessions: semaphore.NewWeighted(cfg.Server.ConcurrentSessionsLimit),
 		remoteAddr:         remoteAddr,
+		sconn:              sconn,
 	}
 }
 
 func (c *connection) handle(ctx context.Context, chans <-chan ssh.NewChannel, handler channelHandler) {
 	ctxlog := log.WithContextFields(ctx, log.Fields{"remote_addr": c.remoteAddr})
 
+	if c.cfg.Server.ClientAliveIntervalSeconds > 0 {
+		ticker := time.NewTicker(c.cfg.Server.ClientAliveInterval())
+		defer ticker.Stop()
+		go c.sendKeepAliveMsg(ctx, ticker)
+	}
+
 	for newChannel := range chans {
 		ctxlog.WithField("channel_type", newChannel.ChannelType()).Info("connection: handle: new channel requested")
 		if newChannel.ChannelType() != "session" {
@@ -68,3 +81,18 @@ func (c *connection) handle(ctx context.Context, chans <-chan ssh.NewChannel, ha
 		}()
 	}
 }
+
+func (c *connection) sendKeepAliveMsg(ctx context.Context, ticker *time.Ticker) {
+	ctxlog := log.WithContextFields(ctx, log.Fields{"remote_addr": c.remoteAddr})
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			ctxlog.Debug("session: handleShell: send keepalive message to a client")
+
+			c.sconn.SendRequest(KeepAliveMsg, true, nil)
+		}
+	}
+}

internal/sshd/connection_test.go

@@ -3,10 +3,14 @@ package sshd
 import (
 	"context"
 	"errors"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 	"golang.org/x/crypto/ssh"
+
+	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
 )
 
 type rejectCall struct {
@@ -47,8 +51,32 @@ func (f *fakeNewChannel) ExtraData() []byte {
 	return f.extraData
 }
 
+type fakeConn struct {
+	ssh.Conn
+
+	sentRequestName string
+	mu              sync.Mutex
+}
+
+func (f *fakeConn) SentRequestName() string {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	return f.sentRequestName
+}
+
+func (f *fakeConn) SendRequest(name string, wantReply bool, payload []byte) (bool, []byte, error) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	f.sentRequestName = name
+
+	return true, nil, nil
+}
+
 func setup(sessionsNum int64, newChannel *fakeNewChannel) (*connection, chan ssh.NewChannel) {
-	conn := newConnection(sessionsNum, "127.0.0.1:50000")
+	cfg := &config.Config{Server: config.ServerConfig{ConcurrentSessionsLimit: sessionsNum, ClientAliveIntervalSeconds: 1}}
+	conn := newConnection(cfg, "127.0.0.1:50000", &ssh.ServerConn{&fakeConn{}, nil})
 
 	chans := make(chan ssh.NewChannel, 1)
 	chans <- newChannel
@@ -145,3 +173,16 @@ func TestAcceptSessionFails(t *testing.T) {
 
 	require.False(t, channelHandled)
 }
+
+func TestClientAliveInterval(t *testing.T) {
+	f := &fakeConn{}
+
+	conn := newConnection(&config.Config{}, "127.0.0.1:50000", &ssh.ServerConn{f, nil})
+
+	ticker := time.NewTicker(time.Millisecond)
+	defer ticker.Stop()
+
+	go conn.sendKeepAliveMsg(context.Background(), ticker)
+
+	require.Eventually(t, func() bool { return KeepAliveMsg == f.SentRequestName() }, time.Second, time.Millisecond)
+}

internal/sshd/sshd.go

@@ -180,7 +180,7 @@ func (s *Server) handleConn(ctx context.Context, nconn net.Conn) {
 
 	started := time.Now()
 	var establishSessionDuration float64
-	conn := newConnection(s.Config.Server.ConcurrentSessionsLimit, remoteAddr)
+	conn := newConnection(s.Config, remoteAddr, sconn)
 	conn.handle(ctx, chans, func(ctx context.Context, channel ssh.Channel, requests <-chan *ssh.Request) {
 		establishSessionDuration = time.Since(started).Seconds()
 		metrics.SshdSessionEstablishedDuration.Observe(establishSessionDuration)

internal/sshd/sshd_test.go

@@ -265,6 +265,7 @@ func setupServerWithConfig(t *testing.T, cfg *config.Config) *Server {
 	cfg.User = user
 	cfg.Server.Listen = serverUrl
 	cfg.Server.ConcurrentSessionsLimit = 1
+	cfg.Server.ClientAliveIntervalSeconds = 15
 	cfg.Server.HostKeyFiles = []string{path.Join(testhelper.TestRoot, "certs/valid/server.key")}
 
 	s, err := NewServer(cfg)