query-orchestration/internal/rbac/middleware_test.go

package rbac

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/labstack/echo/v4"
	"github.com/stretchr/testify/assert"
)

// TestJWTMiddleware tests the core JWT middleware functionality with group-based authorization.
// It validates:
// 1. A user with the required 'foo' group can access a protected endpoint
// 2. A user with only 'bar' group (and not 'foo') is properly denied access
// 3. A user with multiple groups including the required 'foo' group is granted access
// This test demonstrates the basic setup for protecting routes with the JWT middleware
// and ensuring only users with appropriate group membership can access restricted resources.
func TestJWTMiddleware(t *testing.T) {
	// Use local key files
	privateKeyPath := "./private_key.pem"
	publicKeyPath := "./public_key.pem"

	// Create LocalKeyProvider for testing
	keyProvider := NewLocalKeyProvider(privateKeyPath, publicKeyPath)

	// Initialize Echo instance
	e := echo.New()

	// Set up a restricted endpoint that requires the 'foo' group
	e.GET("/test", func(c echo.Context) error {
		return c.String(http.StatusOK, "authorized")
	}, JWTWithConfig(JWTConfig{
		KeyProvider:    keyProvider,
		RequiredGroups: []string{"foo"},
	}))

	// Test cases
	tests := []struct {
		name       string
		groups     []string
		wantStatus int
	}{
		{
			name:       "Authorized user with 'foo' group",
			groups:     []string{"foo"},
			wantStatus: http.StatusOK,
		},
		{
			name:       "Unauthorized user with 'bar' group",
			groups:     []string{"bar"},
			wantStatus: http.StatusForbidden,
		},
		{
			name:       "User with multiple groups including required",
			groups:     []string{"bar", "foo", "baz"},
			wantStatus: http.StatusOK,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Generate JWT token with test groups
			token, err := GenerateMockJWT(keyProvider, "test-user-id", "test@example.com", tt.groups)
			if err != nil {
				t.Fatalf("Failed to generate JWT: %v", err)
			}

			// Create test request
			req := httptest.NewRequest(http.MethodGet, "/test", nil)
			req.Header.Set("Authorization", "Bearer "+token)
			rec := httptest.NewRecorder()
			e.ServeHTTP(rec, req)

			// Verify response
			assert.Equal(t, tt.wantStatus, rec.Code)
			if tt.wantStatus == http.StatusOK {
				assert.Equal(t, "authorized", rec.Body.String())
			}
		})
	}
}

// TestTokenExtraction validates different methods for extracting JWT tokens:
// 1. From Authorization header (standard Bearer token)
// 2. From URL query parameters
// 3. From cookies
// Each method is configured through the TokenLookup field in JWTConfig.
// The test confirms that the middleware can successfully authenticate a user
// regardless of where the token is provided, as long as it's valid and
// contains the required group membership.
func TestTokenExtraction(t *testing.T) {
	// Use local key files
	privateKeyPath := "./private_key.pem"
	publicKeyPath := "./public_key.pem"

	// Create LocalKeyProvider for testing
	keyProvider := NewLocalKeyProvider(privateKeyPath, publicKeyPath)

	// Generate JWT token with 'foo' group
	token, err := GenerateMockJWT(keyProvider, "test-user-id", "test@example.com", []string{"foo"})
	if err != nil {
		t.Fatalf("Failed to generate JWT: %v", err)
	}

	// Test different token extraction methods
	tests := []struct {
		name       string
		config     JWTConfig
		setupReq   func(req *http.Request)
		wantStatus int
	}{
		{
			name: "Extract from header",
			config: JWTConfig{
				KeyProvider:    keyProvider,
				TokenLookup:    "header:Authorization",
				AuthScheme:     "Bearer",
				RequiredGroups: []string{"foo"},
			},
			setupReq: func(req *http.Request) {
				req.Header.Set("Authorization", "Bearer "+token)
			},
			wantStatus: http.StatusOK,
		},
		{
			name: "Extract from query parameter",
			config: JWTConfig{
				KeyProvider:    keyProvider,
				TokenLookup:    "query:token",
				RequiredGroups: []string{"foo"},
			},
			setupReq: func(req *http.Request) {
				q := req.URL.Query()
				q.Add("token", token)
				req.URL.RawQuery = q.Encode()
			},
			wantStatus: http.StatusOK,
		},
		{
			name: "Extract from cookie",
			config: JWTConfig{
				KeyProvider:    keyProvider,
				TokenLookup:    "cookie:jwt",
				RequiredGroups: []string{"foo"},
			},
			setupReq: func(req *http.Request) {
				req.AddCookie(&http.Cookie{
					Name:  "jwt",
					Value: token,
				})
			},
			wantStatus: http.StatusOK,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Initialize Echo instance
			e := echo.New()

			// Set up a restricted endpoint with specified config
			e.GET("/test", func(c echo.Context) error {
				return c.String(http.StatusOK, "authorized")
			}, JWTWithConfig(tt.config))

			// Create test request
			req := httptest.NewRequest(http.MethodGet, "/test", nil)
			tt.setupReq(req)
			rec := httptest.NewRecorder()
			e.ServeHTTP(rec, req)

			// Verify response
			assert.Equal(t, tt.wantStatus, rec.Code)
			if tt.wantStatus == http.StatusOK {
				assert.Equal(t, "authorized", rec.Body.String())
			}
		})
	}
}

// TestCustomErrorHandler verifies that custom error handling works correctly.
// It tests:
//  1. Setting up a custom error handler function in the JWTConfig
//  2. Sending an invalid JWT token in the request
//  3. Confirming that the custom error handler is invoked and provides
//     the expected response format and status code
//
// This allows for customized error responses when authentication fails,
// which is useful for providing more context or formatting errors
// according to API standards.
func TestCustomErrorHandler(t *testing.T) {
	// Use local key files
	privateKeyPath := "./private_key.pem"
	publicKeyPath := "./public_key.pem"

	// Create LocalKeyProvider for testing
	keyProvider := NewLocalKeyProvider(privateKeyPath, publicKeyPath)

	// Initialize Echo instance
	e := echo.New()

	// Custom error handler
	customErrorHandler := func(c echo.Context, err error) error {
		return c.JSON(http.StatusUnauthorized, map[string]string{
			"error": "Custom auth error: " + err.Error(),
		})
	}

	// Set up a restricted endpoint with custom error handler
	e.GET("/test", func(c echo.Context) error {
		return c.String(http.StatusOK, "authorized")
	}, JWTWithConfig(JWTConfig{
		KeyProvider:    keyProvider,
		RequiredGroups: []string{"foo"},
		ErrorHandler:   customErrorHandler,
	}))

	// Create test request with invalid token
	req := httptest.NewRequest(http.MethodGet, "/test", nil)
	req.Header.Set("Authorization", "Bearer invalidtoken")
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, req)

	// Verify custom error response
	assert.Equal(t, http.StatusUnauthorized, rec.Code)
	assert.Contains(t, rec.Body.String(), "Custom auth error")
}

// These tests require valid RSA key files (private_key.pem and public_key.pem)
// in the current directory. These keys are used to sign and verify JWT tokens
// during the tests. In a production environment, the public key would typically
// be retrieved from AWS Cognito's JWKS endpoint, but for testing, we use local files.
// Make sure to generate these key files before running the tests.