feat(ssh-server): check auth key cache in public key callback

3 files changed, 63 insertions, 3 deletions

diff --git a/internal/ssh/server/authkeystore.go b/internal/ssh/server/authkeystore.go
index 45855ff..6e4ca37 100644
--- a/internal/ssh/server/authkeystore.go
+++ b/internal/ssh/server/authkeystore.go
@@ -17,6 +17,8 @@ type authKeyEntry struct {
 	registeredAt time.Time
 }
 
+var authKeyStore = NewAuthKeyStore(0, 0)
+
 // AuthKeyStore is an in-memory, per-user cache of SSH public keys.
 type AuthKeyStore struct {
 	mu             sync.RWMutex
diff --git a/internal/ssh/server/publickeycallback.go b/internal/ssh/server/publickeycallback.go
index bcc9004..ae6ee60 100644
--- a/internal/ssh/server/publickeycallback.go
+++ b/internal/ssh/server/publickeycallback.go
@@ -23,6 +23,11 @@ func PublicKeyCallback(c gossh.ConnMetadata,
 	}
 	dlog.Server.Info(user, "Incoming authorization")
 
+	if permissions := authKeyStorePermissions(user.Name, offeredPubKey); permissions != nil {
+		dlog.Server.Info(user, "Authorized by in-memory auth key store")
+		return permissions, nil
+	}
+
 	authorizedKeysFile, err := authorizedKeysFile(user)
 	if err != nil {
 		return nil, err
@@ -56,14 +61,26 @@ func verifyAuthorizedKeys(user *user.User, authorizedKeysBytes []byte,
 
 	dlog.Server.Debug(user, "Offered public key fingerprint", gossh.FingerprintSHA256(offeredPubKey))
 	if authorizedKeysMap[string(offeredPubKey.Marshal())] {
-		return &gossh.Permissions{
-			Extensions: map[string]string{"pubkey-fp": gossh.FingerprintSHA256(offeredPubKey)},
-		}, nil
+		return permissionsFromPublicKey(offeredPubKey), nil
 	}
 
 	return nil, fmt.Errorf("%s|public key of user not authorized", user)
 }
 
+func authKeyStorePermissions(userName string, offeredPubKey gossh.PublicKey) *gossh.Permissions {
+	if !authKeyStore.Has(userName, offeredPubKey) {
+		return nil
+	}
+
+	return permissionsFromPublicKey(offeredPubKey)
+}
+
+func permissionsFromPublicKey(offeredPubKey gossh.PublicKey) *gossh.Permissions {
+	return &gossh.Permissions{
+		Extensions: map[string]string{"pubkey-fp": gossh.FingerprintSHA256(offeredPubKey)},
+	}
+}
+
 func authorizedKeysFile(user *user.User) (string, error) {
 	if config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		// In this case, we expect a pub key in the current directory.
diff --git a/internal/ssh/server/publickeycallback_test.go b/internal/ssh/server/publickeycallback_test.go
new file mode 100644
index 0000000..7ded4f3
--- /dev/null
+++ b/internal/ssh/server/publickeycallback_test.go
@@ -0,0 +1,41 @@
+package server
+
+import (
+	"testing"
+	"time"
+
+	gossh "golang.org/x/crypto/ssh"
+)
+
+func TestAuthKeyStorePermissions(t *testing.T) {
+	previousStore := authKeyStore
+	authKeyStore = NewAuthKeyStore(time.Hour, 5)
+	t.Cleanup(func() {
+		authKeyStore = previousStore
+	})
+
+	key := testPublicKey(t, 21)
+
+	if permissions := authKeyStorePermissions("alice", key); permissions != nil {
+		t.Fatalf("Expected nil permissions when no key is cached")
+	}
+
+	authKeyStore.Add("alice", key)
+
+	permissions := authKeyStorePermissions("alice", key)
+	if permissions == nil {
+		t.Fatalf("Expected permissions when key is cached")
+	}
+	if fingerprint := permissions.Extensions["pubkey-fp"]; fingerprint != gossh.FingerprintSHA256(key) {
+		t.Fatalf("Unexpected fingerprint: %s", fingerprint)
+	}
+
+	if permissions := authKeyStorePermissions("bob", key); permissions != nil {
+		t.Fatalf("Expected nil permissions for different user")
+	}
+
+	unknownKey := testPublicKey(t, 22)
+	if permissions := authKeyStorePermissions("alice", unknownKey); permissions != nil {
+		t.Fatalf("Expected nil permissions for unknown key")
+	}
+}