From 493add9e2a5273d403a3f606e86a85428392891f Mon Sep 17 00:00:00 2001
From: Igor Drozdov <idrozdov@gitlab.com>
Date: Wed, 30 Mar 2022 08:17:04 +0000
Subject: Abort long-running unauthenticated SSH connections

---
 config.yml.example         |  2 ++
 internal/config/config.go  |  6 ++++++
 internal/sshd/sshd.go      | 16 +++++++++++++++-
 internal/sshd/sshd_test.go | 29 +++++++++++++++++++++++++++++
 4 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/config.yml.example b/config.yml.example
index 3714d73..f3545ba 100644
--- a/config.yml.example
+++ b/config.yml.example
@@ -84,6 +84,8 @@ sshd:
   readiness_probe: "/start"
   # The endpoint that returns 200 OK if the server is alive. Defaults to "/health".
   liveness_probe: "/health"
+  # The server disconnects after this time (in seconds) if the user has not successfully logged in. Defaults to 60.
+  login_grace_time: 60
   # SSH host key files.
   host_key_files:
     - /run/secrets/ssh-hostkeys/ssh_host_rsa_key
diff --git a/internal/config/config.go b/internal/config/config.go
index c4dc073..2d5c29d 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -26,6 +26,7 @@ type ServerConfig struct {
 	ProxyProtocol           bool     `yaml:"proxy_protocol,omitempty"`
 	WebListen               string   `yaml:"web_listen,omitempty"`
 	ConcurrentSessionsLimit int64    `yaml:"concurrent_sessions_limit,omitempty"`
+	LoginGraceTimeSeconds   uint64   `yaml:"login_grace_time"`
 	GracePeriodSeconds      uint64   `yaml:"grace_period"`
 	ReadinessProbe          string   `yaml:"readiness_probe"`
 	LivenessProbe           string   `yaml:"liveness_probe"`
@@ -78,6 +79,7 @@ var (
 		Listen:                  "[::]:22",
 		WebListen:               "localhost:9122",
 		ConcurrentSessionsLimit: 10,
+		LoginGraceTimeSeconds:   60,
 		GracePeriodSeconds:      10,
 		ReadinessProbe:          "/start",
 		LivenessProbe:           "/health",
@@ -89,6 +91,10 @@ var (
 	}
 )
 
+func (sc *ServerConfig) LoginGraceTime() time.Duration {
+	return time.Duration(sc.LoginGraceTimeSeconds) * time.Second
+}
+
 func (sc *ServerConfig) GracePeriod() time.Duration {
 	return time.Duration(sc.GracePeriodSeconds) * time.Second
 }
diff --git a/internal/sshd/sshd.go b/internal/sshd/sshd.go
index d765faf..e8743ae 100644
--- a/internal/sshd/sshd.go
+++ b/internal/sshd/sshd.go
@@ -164,7 +164,7 @@ func (s *Server) handleConn(ctx context.Context, nconn net.Conn) {
 
 	ctxlog.Info("server: handleConn: start")
 
-	sconn, chans, reqs, err := ssh.NewServerConn(nconn, s.serverConfig.get(ctx))
+	sconn, chans, reqs, err := s.initSSHConnection(ctx, nconn)
 	if err != nil {
 		ctxlog.WithError(err).Error("server: handleConn: failed to initialize SSH connection")
 		return
@@ -187,6 +187,20 @@ func (s *Server) handleConn(ctx context.Context, nconn net.Conn) {
 	ctxlog.Info("server: handleConn: done")
 }
 
+func (s *Server) initSSHConnection(ctx context.Context, nconn net.Conn) (sconn *ssh.ServerConn, chans <-chan ssh.NewChannel, reqs <-chan *ssh.Request, err error) {
+	timer := time.AfterFunc(s.Config.Server.LoginGraceTime(), func() {
+		nconn.Close()
+	})
+	defer func() {
+		// If time.Stop() equals false, that means that AfterFunc has been executed
+		if !timer.Stop() {
+			err = fmt.Errorf("initSSHConnection: ssh handshake timeout")
+		}
+	}()
+
+	return ssh.NewServerConn(nconn, s.serverConfig.get(ctx))
+}
+
 func unconditionalRequirePolicy(_ net.Addr) (proxyproto.Policy, error) {
 	return proxyproto.REQUIRE, nil
 }
diff --git a/internal/sshd/sshd_test.go b/internal/sshd/sshd_test.go
index 0c6a8ec..ce395cf 100644
--- a/internal/sshd/sshd_test.go
+++ b/internal/sshd/sshd_test.go
@@ -3,6 +3,7 @@ package sshd
 import (
 	"context"
 	"fmt"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -134,6 +135,33 @@ func TestInvalidServerConfig(t *testing.T) {
 	require.Nil(t, s.Shutdown())
 }
 
+func TestLoginGraceTime(t *testing.T) {
+	s := setupServer(t)
+
+	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)
+
+	// Verify that the server shutdowns instantly without waiting for any connection to execute
+	// That means that the server doesn't have any connections to wait for
+	require.NoError(t, s.Shutdown())
+	verifyStatus(t, s, StatusClosed)
+}
+
 func setupServer(t *testing.T) *Server {
 	t.Helper()
 
@@ -183,6 +211,7 @@ func setupServerWithConfig(t *testing.T, cfg *config.Config) *Server {
 	cfg.User = user
 	cfg.Server.Listen = serverUrl
 	cfg.Server.ConcurrentSessionsLimit = 1
+	cfg.Server.LoginGraceTimeSeconds = 1
 	cfg.Server.HostKeyFiles = []string{path.Join(testhelper.TestRoot, "certs/valid/server.key")}
 
 	s, err := NewServer(cfg)
-- 
cgit v1.2.1