Abort long-running unauthenticated SSH connections

The config option is basically a copy of LoginGraceTime OpenSSH
option.

If an SSH connection is hanging unauthenticated, after some period
of time, the connection gets canceled. The value is configurable,
the server waits for 60 seconds by default.

config.yml.example

@@ -76,10 +76,12 @@ 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
+  # Sets an interval after which server will send keepalive message to a client. Defaults to 15s.
   client_alive_interval: 15
-  # The server waits for this time (in seconds) for the ongoing connections to complete before shutting down. Defaults to 10.
+  # The server waits for this time for the ongoing connections to complete before shutting down. Defaults to 10s.
   grace_period: 10
+  # The server disconnects after this time if the user has not successfully logged in. Defaults to 60s.
+  login_grace_time: 60
   # 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".

internal/config/config.go

@@ -21,7 +21,7 @@ const (
 	defaultSecretFileName = ".gitlab_shell_secret"
 )
 
-type yamlDuration time.Duration
+type YamlDuration time.Duration
 
 type ServerConfig struct {
 	Listen                  string       `yaml:"listen,omitempty"`
@@ -29,9 +29,10 @@ type ServerConfig struct {
 	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"`
+	ClientAliveInterval     YamlDuration `yaml:"client_alive_interval,omitempty"`
+	GracePeriod             YamlDuration `yaml:"grace_period"`
+	ProxyHeaderTimeout      YamlDuration `yaml:"proxy_header_timeout"`
+	LoginGraceTime          YamlDuration `yaml:"login_grace_time"`
 	ReadinessProbe          string       `yaml:"readiness_probe"`
 	LivenessProbe           string       `yaml:"liveness_probe"`
 	HostKeyFiles            []string     `yaml:"host_key_files,omitempty"`
@@ -85,9 +86,10 @@ var (
 		Listen:                  "[::]:22",
 		WebListen:               "localhost:9122",
 		ConcurrentSessionsLimit: 10,
-		GracePeriod:             yamlDuration(10 * time.Second),
-		ClientAliveInterval:     yamlDuration(15 * time.Second),
-		ProxyHeaderTimeout:      yamlDuration(500 * time.Millisecond),
+		GracePeriod:             YamlDuration(10 * time.Second),
+		ClientAliveInterval:     YamlDuration(15 * time.Second),
+		ProxyHeaderTimeout:      YamlDuration(500 * time.Millisecond),
+		LoginGraceTime:          YamlDuration(60 * time.Second),
 		ReadinessProbe:          "/start",
 		LivenessProbe:           "/health",
 		HostKeyFiles: []string{
@@ -98,13 +100,13 @@ var (
 	}
 )
 
-func (d *yamlDuration) UnmarshalYAML(unmarshal func(interface{}) error) error {
+func (d *YamlDuration) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	var intDuration int
 	if err := unmarshal(&intDuration); err != nil {
 		return unmarshal((*time.Duration)(d))
 	}
 
-	*d = yamlDuration(time.Duration(intDuration) * time.Second)
+	*d = YamlDuration(time.Duration(intDuration) * time.Second)
 
 	return nil
 }

internal/config/config_test.go

@@ -83,7 +83,7 @@ func TestYAMLDuration(t *testing.T) {
 	}
 
 	type durationCfg struct {
-		Duration yamlDuration `yaml:"duration"`
+		Duration YamlDuration `yaml:"duration"`
 	}
 
 	for _, tc := range testCases {

internal/sshd/connection.go

@@ -71,10 +71,14 @@ func (c *connection) handle(ctx context.Context, srvCfg *ssh.ServerConfig, handl
 }
 
 func (c *connection) initServerConn(ctx context.Context, srvCfg *ssh.ServerConfig) (*ssh.ServerConn, <-chan ssh.NewChannel, error) {
+	if c.cfg.Server.LoginGraceTime > 0 {
+		c.nconn.SetDeadline(time.Now().Add(time.Duration(c.cfg.Server.LoginGraceTime)))
+		defer c.nconn.SetDeadline(time.Time{})
+	}
+
 	sconn, chans, reqs, err := ssh.NewServerConn(c.nconn, srvCfg)
 	if err != nil {
 		msg := "connection: initServerConn: failed to initialize SSH connection"
-
 		logger := log.WithContextFields(ctx, log.Fields{"remote_addr": c.remoteAddr}).WithError(err)
 
 		if strings.Contains(err.Error(), "no common algorithm for host key") || err.Error() == "EOF" {

internal/sshd/sshd_test.go

@@ -251,6 +251,39 @@ func TestClosingHangedConnections(t *testing.T) {
 	verifyStatus(t, s, StatusClosed)
 }
 
+func TestLoginGraceTime(t *testing.T) {
+	cfg := &config.Config{
+		Server: config.ServerConfig{
+			LoginGraceTime: config.YamlDuration(50 * time.Millisecond),
+		},
+	}
+	s := setupServerWithConfig(t, cfg)
+
+	unauthenticatedRequestStatus := make(chan string)
+	completed := make(chan bool)
+
+	clientCfg := clientConfig(t)
+	clientCfg.HostKeyCallback = func(_ string, _ net.Addr, _ ssh.PublicKey) error {
+		unauthenticatedRequestStatus <- "authentication-started"
+		<-completed // Wait infinitely
+
+		return nil
+	}
+
+	go func() {
+		// Start an SSH connection that never ends
+		ssh.Dial("tcp", serverUrl, clientCfg)
+	}()
+
+	require.Equal(t, "authentication-started", <-unauthenticatedRequestStatus)
+
+	// Shutdown the server and verify that it's closed
+	// If LoginGraceTime doesn't work, then the connection that runs infinitely, will stop it from closing.
+	// The close won't happen until the context is canceled like in TestClosingHangedConnections
+	require.NoError(t, s.Shutdown())
+	verifyStatus(t, s, StatusClosed)
+}
+
 func setupServer(t *testing.T) *Server {
 	t.Helper()