query-orchestration/internal/rbac/middleware_test.go

package rbac_test

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

	"queryorchestration/internal/rbac"

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

// testSetup creates a new Echo instance and key provider for testing
func testSetup(t *testing.T) (*echo.Echo, *rbac.LocalKeyProvider) {
	// Use local key files
	privateKeyPath := "./private_key.pem"
	publicKeyPath := "./public_key.pem"

	// Create and initialize key provider
	keyProvider := rbac.TestKeyProvider(t, privateKeyPath, publicKeyPath)

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

	return e, keyProvider
}

// createTestEndpoint creates a test endpoint with the given JWT config
func createTestEndpoint(e *echo.Echo, config rbac.JWTConfig) {
	e.GET("/test", func(c echo.Context) error {
		return c.String(http.StatusOK, "authorized")
	}, rbac.JWTWithConfig(config))
}

// makeTestRequest creates and executes a test request with the given token
func makeTestRequest(t *testing.T, e *echo.Echo, token string, setupReq func(*http.Request)) *httptest.ResponseRecorder {
	req := httptest.NewRequest(http.MethodGet, "/test", nil)
	if setupReq != nil {
		setupReq(req)
	} else {
		req.Header.Set("Authorization", "Bearer "+token)
	}
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, req)
	return rec
}

// 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) {
	e, keyProvider := testSetup(t)

	// Set up a restricted endpoint that requires the 'foo' group
	createTestEndpoint(e, rbac.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 := rbac.GenerateTestJWT(keyProvider, "test-user-id", "test@example.com", tt.groups)
			if err != nil {
				t.Fatalf("Failed to generate JWT: %v", err)
			}

			// Make test request and verify response
			rec := makeTestRequest(t, e, token, nil)
			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) {
	_, keyProvider := testSetup(t)

	// Generate JWT token with 'foo' group
	token, err := rbac.GenerateTestJWT(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     rbac.JWTConfig
		setupReq   func(req *http.Request)
		wantStatus int
	}{
		{
			name: "Extract from header",
			config: rbac.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: rbac.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: rbac.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) {
			// Create new Echo instance for each test to avoid config conflicts
			e := echo.New()
			createTestEndpoint(e, tt.config)

			// Make test request and verify response
			rec := makeTestRequest(t, e, token, tt.setupReq)
			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) {
	e, keyProvider := testSetup(t)

	// 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
	createTestEndpoint(e, rbac.JWTConfig{
		KeyProvider:    keyProvider,
		RequiredGroups: []string{"foo"},
		ErrorHandler:   customErrorHandler,
	})

	// Create test request with invalid token and verify response
	rec := makeTestRequest(t, e, "invalidtoken", nil)
	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.