at-container-registry/pkg/appview/crypto_keys_test.go

package appview

import (
	"bytes"
	"database/sql"
	"path/filepath"
	"sync"
	"testing"

	"atcr.io/pkg/appview/db"
)

// cryptoTestDB returns a file-backed database.
//
// Not ":memory:" — go-libsql gives each connection to an in-memory DSN its own
// private database, so concurrent loaders would not even see each other's writes
// and the test would pass without testing anything.
func cryptoTestDB(t *testing.T) *sql.DB {
	t.Helper()
	database, err := db.InitDB(filepath.Join(t.TempDir(), "keys.db"), db.LibsqlConfig{})
	if err != nil {
		t.Fatalf("InitDB: %v", err)
	}
	t.Cleanup(func() { database.Close() })
	return database
}

// TestLoadOAuthKeyConcurrentBootAgreesOnOneKey covers two instances booting
// against a fresh database at the same time.
//
// Both find no key, both generate one, and both write. The write used to be
// last-writer-wins, so the loser kept its own key in memory while the database
// held the other's: it then signed OAuth client assertions with a key absent
// from the published JWKS, and every one of them failed verification.
func TestLoadOAuthKeyConcurrentBootAgreesOnOneKey(t *testing.T) {
	database := cryptoTestDB(t)

	const loaders = 6
	keys := make([][]byte, loaders)
	errs := make([]error, loaders)

	var wg sync.WaitGroup
	start := make(chan struct{})
	for i := range loaders {
		wg.Go(func() {
			<-start
			key, err := loadOAuthKey(database)
			if err != nil {
				errs[i] = err
				return
			}
			keys[i] = key.Bytes()
		})
	}
	close(start)
	wg.Wait()

	for i, err := range errs {
		if err != nil {
			t.Fatalf("loader %d: %v", i, err)
		}
	}

	stored, err := db.GetCryptoKey(database, "oauth_p256")
	if err != nil {
		t.Fatalf("GetCryptoKey: %v", err)
	}
	for i, k := range keys {
		if !bytes.Equal(k, stored) {
			t.Errorf("loader %d returned a key that is not the one in the database; "+
				"it would sign with a key missing from the published JWKS", i)
		}
	}
}

// TestLoadRSAKeyConcurrentBootAgreesOnOneKey is the same race for the JWT
// signing key. Losing it means issuing registry JWTs that do not match the
// advertised certificate.
func TestLoadRSAKeyConcurrentBootAgreesOnOneKey(t *testing.T) {
	database := cryptoTestDB(t)

	const loaders = 6
	fingerprints := make([][]byte, loaders)
	errs := make([]error, loaders)

	var wg sync.WaitGroup
	start := make(chan struct{})
	for i := range loaders {
		wg.Go(func() {
			<-start
			key, err := loadRSAKey(database)
			if err != nil {
				errs[i] = err
				return
			}
			fingerprints[i] = key.N.Bytes()
		})
	}
	close(start)
	wg.Wait()

	for i, err := range errs {
		if err != nil {
			t.Fatalf("loader %d: %v", i, err)
		}
	}

	for i := 1; i < loaders; i++ {
		if !bytes.Equal(fingerprints[i], fingerprints[0]) {
			t.Fatalf("loaders disagreed on the JWT signing key: loader %d differs from loader 0", i)
		}
	}
}

// TestPutCryptoKeyKeepsTheFirstWrite pins the storage semantics the loaders rely
// on. If this ever goes back to last-writer-wins, a second instance booting
// would silently replace the key the first one is already signing with.
func TestPutCryptoKeyKeepsTheFirstWrite(t *testing.T) {
	database := cryptoTestDB(t)

	first := []byte("first-key")
	second := []byte("second-key")

	if err := db.PutCryptoKey(database, "probe", first); err != nil {
		t.Fatalf("first put: %v", err)
	}
	if err := db.PutCryptoKey(database, "probe", second); err != nil {
		t.Fatalf("second put: %v", err)
	}

	got, err := db.GetCryptoKey(database, "probe")
	if err != nil {
		t.Fatalf("GetCryptoKey: %v", err)
	}
	if !bytes.Equal(got, first) {
		t.Errorf("stored key = %q, want the first write %q", got, first)
	}
}

// TestLoadOAuthKeyIsStableAcrossRestarts: a restart must reuse the stored key,
// not generate a new one, or every previously issued token breaks.
func TestLoadOAuthKeyIsStableAcrossRestarts(t *testing.T) {
	database := cryptoTestDB(t)

	first, err := loadOAuthKey(database)
	if err != nil {
		t.Fatalf("first load: %v", err)
	}
	second, err := loadOAuthKey(database)
	if err != nil {
		t.Fatalf("second load: %v", err)
	}
	if !bytes.Equal(first.Bytes(), second.Bytes()) {
		t.Error("loadOAuthKey generated a new key instead of reusing the stored one")
	}
}