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"
)

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())
			}
		})
	}
}

// Test different token extraction methods
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())
			}
		})
	}
}

// Test the custom error handler
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")
}

// This test assumes that you have private_key.pem and public_key.pem
// files in the current directory for testing. You need to create these
// files before running the tests.