package traefikoidc

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

// TestJWKSRotationPickedUpOnKidMiss reproduces the key-rotation availability
// bug: after the IdP rotates signing keys and issues a NEW kid, GetPublicKey
// must pick the new key up promptly (bounded live refresh) instead of
// returning "no matching public key found for kid" for up to the 1h cache
// TTL because the stale cached keyset is reused.
func TestJWKSRotationPickedUpOnKidMiss(t *testing.T) {
	require := require.New(t)

	// Two RSA key pairs: keyA is the "pre-rotation" key, keyB the rotated one.
	keyA, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(err)
	keyB, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(err)

	// The server serves kidA until the atomic flag is set, then kidB.
	var rotated int32
	var mu sync.Mutex
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		defer mu.Unlock()
		if atomic.LoadInt32(&rotated) == 1 {
			writeJWKS(t, w, "kidB", &keyB.PublicKey)
			return
		}
		writeJWKS(t, w, "kidA", &keyA.PublicKey)
	}))
	defer srv.Close()

	cache := NewJWKCache()

	// Seed the raw cache with the pre-rotation keyset (as happens on a cold
	// start), WITHOUT taking the kid-miss force path.
	_, err = cache.GetJWKS(context.Background(), srv.URL, http.DefaultClient)
	require.NoError(err)

	// Pre-rotation token for kidA resolves fine.
	_, err = cache.GetPublicKey(context.Background(), srv.URL, "kidA", http.DefaultClient)
	require.NoError(err, "pre-rotation kid should resolve")

	// Simulate a long-lived process: the key rotation happens well after the
	// process start (and thus well after any earlier live fetch), so the
	// force-refresh cooldown has long expired.
	cache.forceMu.Lock()
	cache.lastForceRefresh[srv.URL] = time.Time{}
	cache.forceMu.Unlock()

	// Provider rotates keys; now issues a token with the NEW kid.
	atomic.StoreInt32(&rotated, 1)

	// The new kid must be resolved now (bounded live refresh on kid-miss),
	// not after the 1-hour cache TTL.
	pub, err := cache.GetPublicKey(context.Background(), srv.URL, "kidB", http.DefaultClient)
	require.NoError(err, "post-rotation new kid must be picked up promptly")
	require.NotNil(pub)
	require.Equal(&keyB.PublicKey, pub, "resolved key must be the rotated key")
}

// TestJWKSRotationBoundRefetch verifies the cooldown bounds repeated live
// fetches when an attacker presents bogus unknown kids, while a real rotation
// still succeeds on the first new-kid request.
func TestJWKSRotationBoundRefetch(t *testing.T) {
	require := require.New(t)

	keyA, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(err)

	var fetches int32
	var mu sync.Mutex
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		defer mu.Unlock()
		atomic.AddInt32(&fetches, 1)
		writeJWKS(t, w, "kidA", &keyA.PublicKey)
	}))
	defer srv.Close()

	cache := NewJWKCache()

	// First kid-miss forces exactly one live fetch and returns "no key".
	_, err = cache.GetPublicKey(context.Background(), srv.URL, "bogus1", http.DefaultClient)
	require.Error(err)

	time.Sleep(100 * time.Millisecond)
	// The first request did exactly one live fetch: the gated refresh
	// performed a fresh fetch and, finding no kid, did NOT re-fetch
	// (R124 refinement — an immediate second fetch is redundant).
	require.Equal(int32(1), atomic.LoadInt32(&fetches), "first request must make exactly one live fetch")

	// Within the cooldown window, a further bogus kid triggers the gated
	// (cached) refresh plus ONE ungated live pick-up for the lagging-
	// rotation case (R124), so amplification stays bounded per-request
	// (2 total across both requests, not 1 per unknown key).
	_, err = cache.GetPublicKey(context.Background(), srv.URL, "bogus2", http.DefaultClient)
	require.Error(err)

	time.Sleep(100 * time.Millisecond)
	require.Equal(int32(2), atomic.LoadInt32(&fetches), "within the cooldown window a new bogus kid costs exactly one ungated live fetch")
}