feat(ssh-server): add in-memory auth key store

1 files changed, 178 insertions, 0 deletions

diff --git a/internal/ssh/server/authkeystore_test.go b/internal/ssh/server/authkeystore_test.go
new file mode 100644
index 0000000..056db7b
--- /dev/null
+++ b/internal/ssh/server/authkeystore_test.go
@@ -0,0 +1,178 @@
+package server
+
+import (
+	"crypto/ed25519"
+	"sync"
+	"testing"
+	"time"
+
+	gossh "golang.org/x/crypto/ssh"
+)
+
+func TestAuthKeyStoreAddHasRemove(t *testing.T) {
+	store := NewAuthKeyStore(time.Hour, 5)
+	key := testPublicKey(t, 1)
+
+	if store.Has("alice", key) {
+		t.Fatalf("Store should not contain key before add")
+	}
+
+	store.Add("alice", key)
+	if !store.Has("alice", key) {
+		t.Fatalf("Store should contain key after add")
+	}
+
+	store.Remove("alice", key)
+	if store.Has("alice", key) {
+		t.Fatalf("Store should not contain key after remove")
+	}
+}
+
+func TestAuthKeyStoreHasExpiresKeysLazily(t *testing.T) {
+	now := time.Date(2026, 3, 3, 10, 0, 0, 0, time.UTC)
+	store := newAuthKeyStoreWithClock(10*time.Second, 5, func() time.Time { return now })
+	key := testPublicKey(t, 2)
+
+	store.Add("alice", key)
+	if !store.Has("alice", key) {
+		t.Fatalf("Store should contain fresh key")
+	}
+
+	now = now.Add(11 * time.Second)
+	if store.Has("alice", key) {
+		t.Fatalf("Store should expire key when ttl is exceeded")
+	}
+
+	store.mu.RLock()
+	defer store.mu.RUnlock()
+	if len(store.keysByUser["alice"]) != 0 {
+		t.Fatalf("Expired entries should be removed on Has call")
+	}
+}
+
+func TestAuthKeyStoreEnforcesPerUserKeyLimit(t *testing.T) {
+	now := time.Date(2026, 3, 3, 10, 0, 0, 0, time.UTC)
+	store := newAuthKeyStoreWithClock(time.Hour, 2, func() time.Time { return now })
+
+	keyOne := testPublicKey(t, 3)
+	keyTwo := testPublicKey(t, 4)
+	keyThree := testPublicKey(t, 5)
+
+	store.Add("alice", keyOne)
+	now = now.Add(1 * time.Second)
+	store.Add("alice", keyTwo)
+	now = now.Add(1 * time.Second)
+	store.Add("alice", keyThree)
+
+	if store.Has("alice", keyOne) {
+		t.Fatalf("Oldest key should be evicted once max key limit is reached")
+	}
+	if !store.Has("alice", keyTwo) {
+		t.Fatalf("Second key should remain in store")
+	}
+	if !store.Has("alice", keyThree) {
+		t.Fatalf("Newest key should remain in store")
+	}
+}
+
+func TestAuthKeyStoreAddRefreshesExistingKey(t *testing.T) {
+	now := time.Date(2026, 3, 3, 10, 0, 0, 0, time.UTC)
+	store := newAuthKeyStoreWithClock(10*time.Second, 5, func() time.Time { return now })
+	key := testPublicKey(t, 6)
+
+	store.Add("alice", key)
+	now = now.Add(9 * time.Second)
+	store.Add("alice", key)
+
+	now = now.Add(5 * time.Second)
+	if !store.Has("alice", key) {
+		t.Fatalf("Key should stay valid after it is refreshed")
+	}
+
+	now = now.Add(6 * time.Second)
+	if store.Has("alice", key) {
+		t.Fatalf("Refreshed key should expire once ttl is exceeded from latest add")
+	}
+}
+
+func TestAuthKeyStoreUserIsolation(t *testing.T) {
+	store := NewAuthKeyStore(time.Hour, 5)
+	key := testPublicKey(t, 7)
+
+	store.Add("alice", key)
+	if store.Has("bob", key) {
+		t.Fatalf("Key lookup must be isolated by user")
+	}
+}
+
+func TestAuthKeyStoreIgnoresInvalidInput(t *testing.T) {
+	store := NewAuthKeyStore(time.Hour, 5)
+	key := testPublicKey(t, 8)
+
+	store.Add("", key)
+	store.Add("alice", nil)
+	store.Remove("", key)
+	store.Remove("alice", nil)
+
+	if store.Has("", key) {
+		t.Fatalf("Empty user should not match")
+	}
+	if store.Has("alice", nil) {
+		t.Fatalf("Nil key should not match")
+	}
+}
+
+func TestAuthKeyStoreConcurrentAccess(t *testing.T) {
+	store := NewAuthKeyStore(time.Hour, 5)
+	users := []string{"alice", "bob", "carol"}
+	keys := []gossh.PublicKey{
+		testPublicKey(t, 11),
+		testPublicKey(t, 12),
+		testPublicKey(t, 13),
+		testPublicKey(t, 14),
+	}
+
+	var wg sync.WaitGroup
+	for worker := 0; worker < 32; worker++ {
+		wg.Add(1)
+		go func(workerID int) {
+			defer wg.Done()
+
+			user := users[workerID%len(users)]
+			for i := 0; i < 200; i++ {
+				key := keys[(workerID+i)%len(keys)]
+				store.Add(user, key)
+				_ = store.Has(user, key)
+				if i%3 == 0 {
+					store.Remove(user, key)
+				}
+			}
+		}(worker)
+	}
+	wg.Wait()
+
+	store.mu.RLock()
+	defer store.mu.RUnlock()
+	for user, userEntries := range store.keysByUser {
+		if len(userEntries) > store.maxKeysPerUser {
+			t.Fatalf("User %s exceeded max key limit: %d", user, len(userEntries))
+		}
+	}
+}
+
+func testPublicKey(t *testing.T, seedByte byte) gossh.PublicKey {
+	t.Helper()
+
+	seed := make([]byte, ed25519.SeedSize)
+	for i := range seed {
+		seed[i] = seedByte
+	}
+
+	privateKey := ed25519.NewKeyFromSeed(seed)
+	publicKey, err := gossh.NewPublicKey(privateKey.Public())
+	if err != nil {
+		t.Fatalf("Unable to build ssh public key: %s", err.Error())
+	}
+
+	return publicKey
+}