From 703e587fa51bbc0e72862a569481cff672901171 Mon Sep 17 00:00:00 2001
From: Brandon Williams <brandon.williams@amazee.io>
Date: Wed, 12 Jun 2024 16:17:09 -0500
Subject: [PATCH] feat: increase number of public key attempts before failing

---
 internal/sshserver/serve.go      | 17 +++++++++++------
 internal/sshserver/serve_test.go | 25 ++++++++++++-------------
 2 files changed, 23 insertions(+), 19 deletions(-)

diff --git a/internal/sshserver/serve.go b/internal/sshserver/serve.go
index 3b8c6514..c663b503 100644
--- a/internal/sshserver/serve.go
+++ b/internal/sshserver/serve.go
@@ -19,12 +19,13 @@ import (
 // (e.g. via signal)
 const shutdownTimeout = 8 * time.Second
 
-// disableSHA1Kex returns a ServerConfig which relies on default for everything
-// except key exchange algorithms. There it removes the SHA1 based algorithms.
-//
-// This works around https://github.com/golang/go/issues/59593
-func disableSHA1Kex(_ ssh.Context) *gossh.ServerConfig {
+// serverConfig returns a ServerConfig of default values with overriden public
+// key algorithms and failure attempts.
+func serverConfig(_ ssh.Context) *gossh.ServerConfig {
 	c := gossh.ServerConfig{}
+
+	// Remove the SHA1 based key algorithms.
+	// This works around https://github.com/golang/go/issues/59593
 	c.Config.KeyExchanges = []string{
 		"curve25519-sha256",
 		"curve25519-sha256@libssh.org",
@@ -33,6 +34,10 @@ func disableSHA1Kex(_ ssh.Context) *gossh.ServerConfig {
 		"ecdh-sha2-nistp521",
 		"diffie-hellman-group14-sha256",
 	}
+
+	// Increase the number of public-key attempts before failure.
+	c.MaxAuthTries = 18
+
 	return &c
 }
 
@@ -53,7 +58,7 @@ func Serve(
 			"sftp": ssh.SubsystemHandler(sessionHandler(log, c, true, logAccessEnabled)),
 		},
 		PublicKeyHandler:     pubKeyAuth(log, nc, c),
-		ServerConfigCallback: disableSHA1Kex,
+		ServerConfigCallback: serverConfig,
 		Banner:               banner,
 	}
 	for _, hk := range hostKeys {
diff --git a/internal/sshserver/serve_test.go b/internal/sshserver/serve_test.go
index a727dfe1..29fba05e 100644
--- a/internal/sshserver/serve_test.go
+++ b/internal/sshserver/serve_test.go
@@ -8,17 +8,16 @@ import (
 )
 
 func TestDisableSHA1Kex(t *testing.T) {
-	var testCases = map[string]struct {
-		input  string
-		expect bool
-	}{
-		"no sha1": {input: "diffie-hellman-group14-sha1", expect: false},
-	}
-	for name, tc := range testCases {
-		t.Run(name, func(tt *testing.T) {
-			conf := disableSHA1Kex(nil)
-			assert.Equal(tt, tc.expect,
-				slices.Contains(conf.Config.KeyExchanges, tc.input), name)
-		})
-	}
+	t.Run("no sha1", func(tt *testing.T) {
+		conf := serverConfig(nil)
+		assert.Equal(tt, false,
+			slices.Contains(conf.Config.KeyExchanges, "diffie-hellman-group14-sha1"), "no sha1")
+	})
+}
+
+func TestMaxAuthTries(t *testing.T) {
+	t.Run("MaxAuthTries", func(tt *testing.T) {
+		conf := serverConfig(nil)
+		assert.Equal(tt, 18, conf.MaxAuthTries, "MaxAuthTries")
+	})
 }