Merge branch 'id-fix-proxy-header-timeout' into 'main'

Make ProxyHeaderTimeout configurable

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

9 files changed, 84 insertions, 33 deletions

diff --git a/cmd/gitlab-sshd/main.go b/cmd/gitlab-sshd/main.go
index 8040a54..16fbca0 100644
--- a/cmd/gitlab-sshd/main.go
+++ b/cmd/gitlab-sshd/main.go
@@ -99,11 +99,12 @@ func main() {
 		sig := <-done
 		signal.Reset(syscall.SIGINT, syscall.SIGTERM)
 
-		log.WithContextFields(ctx, log.Fields{"shutdown_timeout_s": cfg.Server.GracePeriodSeconds, "signal": sig.String()}).Info("Shutdown initiated")
+		gracePeriod := time.Duration(cfg.Server.GracePeriod)
+		log.WithContextFields(ctx, log.Fields{"shutdown_timeout_s": gracePeriod.Seconds(), "signal": sig.String()}).Info("Shutdown initiated")
 
 		server.Shutdown()
 
-		<-time.After(cfg.Server.GracePeriod())
+		<-time.After(gracePeriod)
 
 		cancel()
 
diff --git a/config.yml.example b/config.yml.example
index 6c4bf0f..0e75d75 100644
--- a/config.yml.example
+++ b/config.yml.example
@@ -80,6 +80,8 @@ sshd:
   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
+  # A short timeout to decide to abort the connection if the protocol header is not seen within it. Defaults to 500ms
+  proxy_header_timeout: 500ms
   # The endpoint that returns 200 OK if the server is ready to receive incoming connections; otherwise, it returns 503 Service Unavailable. Defaults to "/start".
   readiness_probe: "/start"
   # The endpoint that returns 200 OK if the server is alive. Defaults to "/health".
diff --git a/internal/config/config.go b/internal/config/config.go
index 79f687d..9e95931 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -21,20 +21,23 @@ const (
 	defaultSecretFileName = ".gitlab_shell_secret"
 )
 
+type yamlDuration time.Duration
+
 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"`
-	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"`
-	MACs                       []string `yaml:"macs"`
-	KexAlgorithms              []string `yaml:"kex_algorithms"`
-	Ciphers                    []string `yaml:"ciphers"`
+	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"`
+	ClientAliveInterval     yamlDuration `yaml:"client_alive_interval,omitempty"`
+	GracePeriod             yamlDuration `yaml:"grace_period"`
+	ProxyHeaderTimeout      yamlDuration `yaml:"proxy_header_timeout"`
+	ReadinessProbe          string       `yaml:"readiness_probe"`
+	LivenessProbe           string       `yaml:"liveness_probe"`
+	HostKeyFiles            []string     `yaml:"host_key_files,omitempty"`
+	MACs                    []string     `yaml:"macs"`
+	KexAlgorithms           []string     `yaml:"kex_algorithms"`
+	Ciphers                 []string     `yaml:"ciphers"`
 }
 
 type HttpSettingsConfig struct {
@@ -79,13 +82,14 @@ var (
 	}
 
 	DefaultServerConfig = ServerConfig{
-		Listen:                     "[::]:22",
-		WebListen:                  "localhost:9122",
-		ConcurrentSessionsLimit:    10,
-		GracePeriodSeconds:         10,
-		ClientAliveIntervalSeconds: 15,
-		ReadinessProbe:             "/start",
-		LivenessProbe:              "/health",
+		Listen:                  "[::]:22",
+		WebListen:               "localhost:9122",
+		ConcurrentSessionsLimit: 10,
+		GracePeriod:             yamlDuration(10 * time.Second),
+		ClientAliveInterval:     yamlDuration(15 * time.Second),
+		ProxyHeaderTimeout:      yamlDuration(500 * time.Millisecond),
+		ReadinessProbe:          "/start",
+		LivenessProbe:           "/health",
 		HostKeyFiles: []string{
 			"/run/secrets/ssh-hostkeys/ssh_host_rsa_key",
 			"/run/secrets/ssh-hostkeys/ssh_host_ecdsa_key",
@@ -94,12 +98,15 @@ var (
 	}
 )
 
-func (sc *ServerConfig) ClientAliveInterval() time.Duration {
-	return time.Duration(sc.ClientAliveIntervalSeconds) * time.Second
-}
+func (d *yamlDuration) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var intDuration int
+	if err := unmarshal(&intDuration); err != nil {
+		return unmarshal((*time.Duration)(d))
+	}
 
-func (sc *ServerConfig) GracePeriod() time.Duration {
-	return time.Duration(sc.GracePeriodSeconds) * time.Second
+	*d = yamlDuration(time.Duration(intDuration) * time.Second)
+
+	return nil
 }
 
 func (c *Config) ApplyGlobalState() {
diff --git a/internal/config/config_test.go b/internal/config/config_test.go
index 32580b8..9d9e20a 100644
--- a/internal/config/config_test.go
+++ b/internal/config/config_test.go
@@ -3,9 +3,11 @@ package config
 import (
 	"os"
 	"testing"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/require"
+	yaml "gopkg.in/yaml.v2"
 
 	"gitlab.com/gitlab-org/gitlab-shell/client/testserver"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/testhelper"
@@ -58,3 +60,40 @@ func TestCustomPrometheusMetrics(t *testing.T) {
 
 	require.Equal(t, expectedMetricNames, actualNames)
 }
+
+func TestNewFromDir(t *testing.T) {
+	testhelper.PrepareTestRootDir(t)
+
+	cfg, err := NewFromDir(testhelper.TestRoot)
+	require.NoError(t, err)
+
+	require.Equal(t, 10*time.Second, time.Duration(cfg.Server.GracePeriod))
+	require.Equal(t, 1*time.Minute, time.Duration(cfg.Server.ClientAliveInterval))
+	require.Equal(t, 500*time.Millisecond, time.Duration(cfg.Server.ProxyHeaderTimeout))
+}
+
+func TestYAMLDuration(t *testing.T) {
+	testCases := []struct {
+		desc     string
+		data     string
+		duration time.Duration
+	}{
+		{"seconds assumed by default", "duration: 10", 10 * time.Second},
+		{"milliseconds are parsed", "duration: 500ms", 500 * time.Millisecond},
+		{"minutes are parsed", "duration: 1m", 1 * time.Minute},
+	}
+
+	type durationCfg struct {
+		Duration yamlDuration `yaml:"duration"`
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			var cfg durationCfg
+			err := yaml.Unmarshal([]byte(tc.data), &cfg)
+			require.NoError(t, err)
+
+			require.Equal(t, tc.duration, time.Duration(cfg.Duration))
+		})
+	}
+}
diff --git a/internal/sshd/connection.go b/internal/sshd/connection.go
index d38e3d8..7e56244 100644
--- a/internal/sshd/connection.go
+++ b/internal/sshd/connection.go
@@ -46,8 +46,8 @@ func newConnection(cfg *config.Config, remoteAddr string, sconn *ssh.ServerConn)
 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())
+	if c.cfg.Server.ClientAliveInterval > 0 {
+		ticker := time.NewTicker(time.Duration(c.cfg.Server.ClientAliveInterval))
 		defer ticker.Stop()
 		go c.sendKeepAliveMsg(ctx, ticker)
 	}
diff --git a/internal/sshd/connection_test.go b/internal/sshd/connection_test.go
index fe7e4be..942e9cd 100644
--- a/internal/sshd/connection_test.go
+++ b/internal/sshd/connection_test.go
@@ -80,7 +80,7 @@ func (f *fakeConn) SendRequest(name string, wantReply bool, payload []byte) (boo
 }
 
 func setup(sessionsNum int64, newChannel *fakeNewChannel) (*connection, chan ssh.NewChannel) {
-	cfg := &config.Config{Server: config.ServerConfig{ConcurrentSessionsLimit: sessionsNum, ClientAliveIntervalSeconds: 1}}
+	cfg := &config.Config{Server: config.ServerConfig{ConcurrentSessionsLimit: sessionsNum}}
 	conn := newConnection(cfg, "127.0.0.1:50000", &ssh.ServerConn{&fakeConn{}, nil})
 
 	chans := make(chan ssh.NewChannel, 1)
diff --git a/internal/sshd/sshd.go b/internal/sshd/sshd.go
index a9cd302..4d4d6d5 100644
--- a/internal/sshd/sshd.go
+++ b/internal/sshd/sshd.go
@@ -26,7 +26,6 @@ const (
 	StatusReady
 	StatusOnShutdown
 	StatusClosed
-	ProxyHeaderTimeout = 90 * time.Second
 )
 
 type Server struct {
@@ -97,7 +96,7 @@ func (s *Server) listen(ctx context.Context) error {
 		sshListener = &proxyproto.Listener{
 			Listener:          sshListener,
 			Policy:            s.requirePolicy,
-			ReadHeaderTimeout: ProxyHeaderTimeout,
+			ReadHeaderTimeout: time.Duration(s.Config.Server.ProxyHeaderTimeout),
 		}
 
 		log.ContextLogger(ctx).Info("Proxy protocol is enabled")
diff --git a/internal/sshd/sshd_test.go b/internal/sshd/sshd_test.go
index 80495f6..d725add 100644
--- a/internal/sshd/sshd_test.go
+++ b/internal/sshd/sshd_test.go
@@ -265,7 +265,6 @@ 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)
diff --git a/internal/testhelper/testdata/testroot/config.yml b/internal/testhelper/testdata/testroot/config.yml
index e69de29..89d7b73 100644
--- a/internal/testhelper/testdata/testroot/config.yml
+++ b/internal/testhelper/testdata/testroot/config.yml
@@ -0,0 +1,4 @@
+sshd:
+  grace_period: 10
+  client_alive_interval: 1m
+  proxy_header_timeout: 500ms