Allow specifying formatted durations in config

- If an integer is specified, we assume that these are seconds
- A duration of format "500ms", "10s", "1m", etc... accepted

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()
 

internal/config/config.go

@@ -21,20 +21,22 @@ 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"`
+	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 +81,13 @@ 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),
+		ReadinessProbe:          "/start",
+		LivenessProbe:           "/health",
 		HostKeyFiles: []string{
 			"/run/secrets/ssh-hostkeys/ssh_host_rsa_key",
 			"/run/secrets/ssh-hostkeys/ssh_host_ecdsa_key",
@@ -94,12 +96,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() {

internal/config/config_test.go

@@ -3,7 +3,9 @@ package config
 import (
 	"os"
 	"testing"
+	"time"
 
+	yaml "gopkg.in/yaml.v2"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/require"
 
@@ -58,3 +60,39 @@ 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))
+}
+
+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))
+		})
+	}
+}

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)
 	}

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)

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)

internal/testhelper/testdata/testroot/config.yml

+sshd:
+  grace_period: 10
+  client_alive_interval: 1m