jay brown
327 lines
10 KiB
Go
327 lines
10 KiB
Go
package cognitoauth
|
|
|
|
import (
|
|
"crypto/rand"
|
|
"crypto/rsa"
|
|
"encoding/json"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/lestrrat-go/jwx/v2/jwa"
|
|
"github.com/lestrrat-go/jwx/v2/jwk"
|
|
)
|
|
|
|
// generateTestJWKS creates a test JSON Web Key Set (JWKS) for testing purposes.
|
|
// It performs the following operations:
|
|
// - Generates a 2048-bit RSA key pair for cryptographic operations
|
|
// - Creates a JWK (JSON Web Key) from the public key portion
|
|
// - Sets key properties including ID ("test-key-id"), algorithm (RS256), and usage ("sig")
|
|
// - Creates a JWK Set and adds the key to it
|
|
// - Returns both the JWK Set object and its JSON string representation
|
|
func generateTestJWKS(t *testing.T) (jwk.Set, string) {
|
|
// Generate RSA key for testing
|
|
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
|
|
if err != nil {
|
|
t.Fatalf("failed to generate RSA key: %v", err)
|
|
}
|
|
|
|
// Create JWK from RSA key
|
|
key, err := jwk.FromRaw(privateKey.PublicKey)
|
|
if err != nil {
|
|
t.Fatalf("failed to create JWK: %v", err)
|
|
}
|
|
|
|
// Set key properties
|
|
if err := key.Set(jwk.KeyIDKey, "test-key-id"); err != nil {
|
|
t.Fatalf("failed to set key ID: %v", err)
|
|
}
|
|
if err := key.Set(jwk.AlgorithmKey, jwa.RS256); err != nil {
|
|
t.Fatalf("failed to set algorithm: %v", err)
|
|
}
|
|
if err := key.Set(jwk.KeyUsageKey, "sig"); err != nil {
|
|
t.Fatalf("failed to set key usage: %v", err)
|
|
}
|
|
|
|
// Create JWK Set with the key
|
|
keySet := jwk.NewSet()
|
|
if err := keySet.AddKey(key); err != nil {
|
|
t.Fatalf("failed to add key to key set: %v", err)
|
|
}
|
|
|
|
// Serialize to JSON
|
|
jwksJSON, err := json.Marshal(keySet)
|
|
if err != nil {
|
|
t.Fatalf("failed to marshal JWK Set: %v", err)
|
|
}
|
|
|
|
return keySet, string(jwksJSON)
|
|
}
|
|
|
|
// TestGetJWKS tests the GetJWKS function through multiple scenarios including
|
|
// initial fetching, caching, cache expiration, and error handling.
|
|
func TestGetJWKS(t *testing.T) {
|
|
// Setup logger
|
|
logHandler := slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
|
|
Level: slog.LevelDebug,
|
|
})
|
|
logger := slog.New(logHandler)
|
|
|
|
// Generate test JWK set for mock responses
|
|
testKeySet, jwksJSON := generateTestJWKS(t)
|
|
|
|
// Mock server to return JWKS
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, err := w.Write([]byte(jwksJSON))
|
|
if err != nil {
|
|
t.Errorf("Failed to write response: %v", err)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
|
|
// Reset JWKS cache before tests
|
|
jwksCache = &JWKSCache{
|
|
ExpiresAt: time.Now(), // Initial state is expired
|
|
}
|
|
|
|
// Test: First fetch - should make HTTP request
|
|
// Tests initial JWKS retrieval and verifies:
|
|
// - Successful HTTP request and response processing
|
|
// - Correct number of keys in the retrieved set
|
|
// - Proper caching of the retrieved key set
|
|
// - Correct cache expiration time setting
|
|
t.Run("first fetch", func(t *testing.T) {
|
|
keySet, err := GetJWKS(server.URL, logger)
|
|
if err != nil {
|
|
t.Fatalf("GetJWKS() error = %v", err)
|
|
}
|
|
|
|
// Verify key set contains expected keys
|
|
if keySet.Len() != testKeySet.Len() {
|
|
t.Errorf("GetJWKS() returned %d keys, want %d", keySet.Len(), testKeySet.Len())
|
|
}
|
|
|
|
// Check that it's cached
|
|
if jwksCache.KeySet == nil {
|
|
t.Errorf("KeySet was not cached after fetch")
|
|
}
|
|
if time.Now().After(jwksCache.ExpiresAt) {
|
|
t.Errorf("Cache expiration not set correctly")
|
|
}
|
|
})
|
|
|
|
// Test: Subsequent fetch - should use cache
|
|
// Verifies caching behavior:
|
|
// - Tests that subsequent requests use cached data instead of making HTTP requests
|
|
// - Tests the cache works even when the server is unavailable/failing
|
|
// - Validates that the cached key count matches expected values
|
|
t.Run("cached fetch", func(t *testing.T) {
|
|
// Create a mock server that fails - to verify we're using the cache
|
|
failServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
_, err := w.Write([]byte("This server should not be called"))
|
|
if err != nil {
|
|
t.Errorf("Failed to write response: %v", err)
|
|
}
|
|
}))
|
|
defer failServer.Close()
|
|
|
|
// We should still get valid keys because we use the cache
|
|
keySet, err := GetJWKS(failServer.URL, logger)
|
|
if err != nil {
|
|
t.Fatalf("GetJWKS() with cache error = %v", err)
|
|
}
|
|
|
|
// Verify key set contains expected keys
|
|
if keySet.Len() != testKeySet.Len() {
|
|
t.Errorf("GetJWKS() with cache returned %d keys, want %d", keySet.Len(), testKeySet.Len())
|
|
}
|
|
})
|
|
|
|
// Test: Cache expiry - should fetch again
|
|
// Tests cache expiration behavior:
|
|
// - Verifies new keys are fetched after cache expiration
|
|
// - Ensures the system properly handles refreshing expired cache data
|
|
// - Validates that updated keys are received correctly
|
|
t.Run("cache expiry", func(t *testing.T) {
|
|
// Force cache expiration
|
|
jwksCache.mutex.Lock()
|
|
jwksCache.ExpiresAt = time.Now().Add(-1 * time.Hour) // Expired one hour ago
|
|
jwksCache.mutex.Unlock()
|
|
|
|
// New server with updated keys
|
|
updatedKeySet, updatedJSON := generateTestJWKS(t)
|
|
updateServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, err := w.Write([]byte(updatedJSON))
|
|
if err != nil {
|
|
t.Errorf("Failed to write response: %v", err)
|
|
}
|
|
}))
|
|
defer updateServer.Close()
|
|
|
|
// Now we should fetch new keys
|
|
keySet, err := GetJWKS(updateServer.URL, logger)
|
|
if err != nil {
|
|
t.Fatalf("GetJWKS() after expiry error = %v", err)
|
|
}
|
|
|
|
// Verify key set contains expected keys (should match updated set)
|
|
if keySet.Len() != updatedKeySet.Len() {
|
|
t.Errorf("GetJWKS() after expiry returned %d keys, want %d", keySet.Len(), updatedKeySet.Len())
|
|
}
|
|
})
|
|
|
|
// Test: Server error
|
|
// Tests error handling when server fails:
|
|
// - Validates behavior with 500 Internal Server Error responses
|
|
// - Verifies appropriate error is returned to caller
|
|
// - Tests system behavior with empty/reset cache
|
|
t.Run("server error", func(t *testing.T) {
|
|
// Reset cache
|
|
jwksCache = &JWKSCache{
|
|
ExpiresAt: time.Now(), // Initial state is expired
|
|
}
|
|
|
|
// Create a mock server that returns an error
|
|
errorServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
_, err := w.Write([]byte("Internal Server Error"))
|
|
if err != nil {
|
|
t.Errorf("Failed to write response: %v", err)
|
|
}
|
|
}))
|
|
defer errorServer.Close()
|
|
|
|
// Should return an error
|
|
_, err := GetJWKS(errorServer.URL, logger)
|
|
if err == nil {
|
|
t.Errorf("GetJWKS() with server error should return error, got nil")
|
|
}
|
|
})
|
|
|
|
// Test: Invalid JSON response
|
|
// Tests handling of malformed responses:
|
|
// - Verifies error handling for invalid JSON responses
|
|
// - Ensures system stability when receiving corrupted data
|
|
// - Validates error reporting for bad responses
|
|
t.Run("invalid json", func(t *testing.T) {
|
|
// Reset cache
|
|
jwksCache = &JWKSCache{
|
|
ExpiresAt: time.Now(), // Initial state is expired
|
|
}
|
|
|
|
// Create a mock server that returns invalid JSON
|
|
invalidServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, err := w.Write([]byte("This is not valid JSON"))
|
|
if err != nil {
|
|
t.Errorf("Failed to write response: %v", err)
|
|
}
|
|
}))
|
|
defer invalidServer.Close()
|
|
|
|
// Should return an error
|
|
_, err := GetJWKS(invalidServer.URL, logger)
|
|
if err == nil {
|
|
t.Errorf("GetJWKS() with invalid JSON should return error, got nil")
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestFetchJWKS tests the fetchJWKS function directly with various scenarios
|
|
// including successful fetches, timeouts, invalid URLs, and error responses.
|
|
func TestFetchJWKS(t *testing.T) {
|
|
// Generate test JWK set for mock responses
|
|
_, jwksJSON := generateTestJWKS(t)
|
|
|
|
// Test: Successful fetch
|
|
// Verifies basic JWKS retrieval functionality:
|
|
// - Tests successful HTTP request and response handling
|
|
// - Validates that response body matches expected JSON
|
|
// - Confirms proper return of fetched content
|
|
t.Run("successful fetch", func(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, err := w.Write([]byte(jwksJSON))
|
|
if err != nil {
|
|
t.Errorf("Failed to write response: %v", err)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
|
|
result, err := fetchJWKS(server.URL)
|
|
if err != nil {
|
|
t.Fatalf("fetchJWKS() error = %v", err)
|
|
}
|
|
|
|
if result != jwksJSON {
|
|
t.Errorf("fetchJWKS() = %v, want %v", result, jwksJSON)
|
|
}
|
|
})
|
|
|
|
// Test: Timeout
|
|
// Validates timeout handling:
|
|
// - Tests behavior when server response exceeds client timeout limit
|
|
// - Verifies appropriate error handling for timed-out requests
|
|
// - Ensures the function fails properly in timeout scenarios
|
|
t.Run("timeout", func(t *testing.T) {
|
|
// Create a server that sleeps longer than the client timeout
|
|
slowServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
time.Sleep(JWKS_CLIENT_TIMEOUT + 1*time.Second)
|
|
_, err := w.Write([]byte(jwksJSON))
|
|
if err != nil {
|
|
t.Errorf("Failed to write response: %v", err)
|
|
}
|
|
}))
|
|
defer slowServer.Close()
|
|
|
|
_, err := fetchJWKS(slowServer.URL)
|
|
if err == nil {
|
|
t.Errorf("fetchJWKS() with timeout should return error, got nil")
|
|
}
|
|
})
|
|
|
|
// Test: Invalid URL
|
|
// Tests error handling for bad URLs:
|
|
// - Verifies behavior with non-existent or invalid URLs
|
|
// - Ensures appropriate error responses for network failures
|
|
// - Validates error handling for connection issues
|
|
t.Run("invalid url", func(t *testing.T) {
|
|
_, err := fetchJWKS("http://invalid-url-that-does-not-exist.example.com")
|
|
if err == nil {
|
|
t.Errorf("fetchJWKS() with invalid URL should return error, got nil")
|
|
}
|
|
})
|
|
|
|
// Test: HTTP error status
|
|
// Validates handling of HTTP error responses:
|
|
// - Tests behavior with various HTTP error status codes (e.g., 404)
|
|
// - Verifies response body handling regardless of status code
|
|
// - Tests that function returns body content even for error status codes
|
|
t.Run("http error status", func(t *testing.T) {
|
|
errorServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusNotFound)
|
|
_, err := w.Write([]byte("Not Found"))
|
|
if err != nil {
|
|
t.Errorf("Failed to write response: %v", err)
|
|
}
|
|
}))
|
|
defer errorServer.Close()
|
|
|
|
result, err := fetchJWKS(errorServer.URL)
|
|
if err != nil {
|
|
t.Fatalf("fetchJWKS() unexpected error = %v", err)
|
|
}
|
|
|
|
// Function doesn't check status code, just reads body
|
|
if result != "Not Found" {
|
|
t.Errorf("fetchJWKS() = %v, want %v", result, "Not Found")
|
|
}
|
|
})
|
|
}
|