package main

import (
	"context"
	"encoding/base64"
	"math/big"
	"testing"
	"time"

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

func TestRSAPublicKeyAcceptsPaddedAndUnpaddedBase64URL(t *testing.T) {
	t.Parallel()

	publicKey := &newPrivateKey(t).PublicKey
	exponent := base64.RawURLEncoding.EncodeToString(new(big.Int).SetInt64(int64(publicKey.E)).Bytes())

	tests := []struct {
		name    string
		modulus string
	}{
		{
			name:    "unpadded",
			modulus: base64.RawURLEncoding.EncodeToString(publicKey.N.Bytes()),
		},
		{
			name:    "padded",
			modulus: base64.URLEncoding.EncodeToString(publicKey.N.Bytes()),
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()

			decoded, err := rsaPublicKey(jsonWebKey{
				Exponent: exponent,
				Modulus:  test.modulus,
			})
			require.NoError(t, err)
			require.Equal(t, publicKey, decoded)
		})
	}
}

func TestJWKSRefreshOnUnknownKeyIsRateLimited(t *testing.T) {
	t.Parallel()

	privateKey := newPrivateKey(t)
	server, requests := newJWKSServer(t, &privateKey.PublicKey)
	cache := newJWKSCache(server.URL, server.Client(), time.Hour)

	now := time.Now()
	cache.now = func() time.Time { return now }

	_, err := cache.key(context.Background(), "unknown-key-1")
	require.Error(t, err)
	require.Equal(t, int32(1), requests.Load())

	// A known key remains usable from the fresh cache, while arbitrary unknown
	// kids cannot trigger another outbound request during the cooldown.
	key, err := cache.key(context.Background(), testKeyID)
	require.NoError(t, err)
	require.Equal(t, &privateKey.PublicKey, key)

	_, err = cache.key(context.Background(), "unknown-key-2")
	require.Error(t, err)
	require.Equal(t, int32(1), requests.Load())

	now = now.Add(defaultJWKSRefreshCooldown)
	_, err = cache.key(context.Background(), "unknown-key-2")
	require.Error(t, err)
	require.Equal(t, int32(2), requests.Load())
}