jacob-mathison/query-orchestration
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity

wip

Browse Source
This commit is contained in:
jay brown
2025-03-17 16:36:02 -07:00
parent 18db7e5193
commit f531d06651
9 changed files with 271 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/cognito_test/jwtstuff/main.go
+1 -1
View File
@@ -12,7 +12,7 @@ func main() {
// Generate a mock JWT with multiple Cognito groups
groups := []string{"admins", "developers", "querybuilders", "uploaders", "exporters"}
tokenString, err := rbac.GenerateMockJWT(localProvider, "test-user-id", "test@example.com", groups)
tokenString, err := rbac.GenerateTestJWT(localProvider, "test-user-id", "test@example.com", groups)
if err != nil {
fmt.Println("Error generating token:", err)
return
internal/rbac/cognito.go
+2 -2
View File
@@ -203,8 +203,8 @@ func NewCognitoKeyProvider(region, userPoolID string) *CognitoKeyProvider {
}
}
// GenerateMockJWT creates a test JWT with given groups
func GenerateMockJWT(keyProvider *LocalKeyProvider, userID, email string, groups []string) (string, error) {
// GenerateTestJWT creates a test JWT with given groups
func GenerateTestJWT(keyProvider *LocalKeyProvider, userID, email string, groups []string) (string, error) {
privateKey, err := keyProvider.GetPrivateKey()
if err != nil {
return "", err
internal/rbac/cognito_test.go
+43 -60
View File
@@ -1,17 +1,21 @@
package rbac
package rbac_test
import (
"fmt"
"reflect"
"testing"
"time"
"queryorchestration/internal/rbac"
"github.com/golang-jwt/jwt/v5"
"github.com/stretchr/testify/assert"
)
// TestLocalKeyProvider validates that the LocalKeyProvider can successfully
// load both private and public keys from the specified PEM files.
func TestLocalKeyProvider(t *testing.T) {
provider := NewLocalKeyProvider("private_key.pem", "public_key.pem")
provider := rbac.NewLocalKeyProvider("private_key.pem", "public_key.pem")
// Test loading private key
privateKey, err := provider.GetPrivateKey()
@@ -38,7 +42,7 @@ func TestLocalKeyProvider(t *testing.T) {
// 3. Verifying that all claims are correctly preserved and can be retrieved
// This test covers both tokens with and without group information.
func TestGenerateAndValidateJWT(t *testing.T) {
provider := NewLocalKeyProvider("private_key.pem", "public_key.pem")
provider := rbac.NewLocalKeyProvider("private_key.pem", "public_key.pem")
testCases := []struct {
name string
@@ -66,13 +70,13 @@ func TestGenerateAndValidateJWT(t *testing.T) {
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
// Generate token
tokenString, err := GenerateMockJWT(provider, tc.userID, tc.email, tc.groups)
tokenString, err := rbac.GenerateTestJWT(provider, tc.userID, tc.email, tc.groups)
if err != nil {
t.Fatalf("Failed to generate token: %v", err)
}
// Validate token
token, claims, err := ValidateJWT(tokenString, provider)
token, claims, err := rbac.ValidateJWT(tokenString, provider)
if err != nil {
t.Fatalf("Failed to validate token: %v", err)
}
@@ -115,7 +119,7 @@ func TestGenerateAndValidateJWT(t *testing.T) {
// The actual expired token test is commented out because jwt.Parse
// automatically validates expiration times.
func TestTokenExpiration(t *testing.T) {
provider := NewLocalKeyProvider("private_key.pem", "public_key.pem")
provider := rbac.NewLocalKeyProvider("private_key.pem", "public_key.pem")
// Helper function to generate a token with custom claims
generateCustomToken := func(expiry time.Time) (string, error) {
@@ -152,7 +156,7 @@ func TestTokenExpiration(t *testing.T) {
t.Fatalf("Failed to generate token: %v", err)
}
token, _, err := ValidateJWT(tokenString, provider)
token, _, err := rbac.ValidateJWT(tokenString, provider)
if err != nil {
t.Fatalf("Expected valid token, got error: %v", err)
}
@@ -171,7 +175,7 @@ func TestTokenExpiration(t *testing.T) {
t.Fatalf("Failed to generate token: %v", err)
}
_, _, err = ValidateJWT(tokenString, provider)
_, _, err = rbac.ValidateJWT(tokenString, provider)
if err == nil {
t.Fatal("Expected error for expired token, got nil")
}
@@ -187,23 +191,23 @@ func TestTokenExpiration(t *testing.T) {
// - Token expiration time
// This ensures user context can be properly reconstructed from the token.
func TestExtractTokenDetails(t *testing.T) {
provider := NewLocalKeyProvider("private_key.pem", "public_key.pem")
provider := rbac.NewLocalKeyProvider("private_key.pem", "public_key.pem")
userID := "test-user-456"
email := "test456@example.com"
groups := []string{"admins", "developers", "testers"}
tokenString, err := GenerateMockJWT(provider, userID, email, groups)
tokenString, err := rbac.GenerateTestJWT(provider, userID, email, groups)
if err != nil {
t.Fatalf("Failed to generate token: %v", err)
}
token, claims, err := ValidateJWT(tokenString, provider)
token, claims, err := rbac.ValidateJWT(tokenString, provider)
if err != nil {
t.Fatalf("Failed to validate token: %v", err)
}
details := ExtractTokenDetails(token, claims)
details := rbac.ExtractTokenDetails(token, claims)
// Check extracted details
if details.UserID != userID {
@@ -248,15 +252,15 @@ func TestGroupMembershipFunctions(t *testing.T) {
{"viewers", true},
{"editors", false},
{"ADMINS", true}, // Case-insensitive match
{"Developers", true}, // Case-insensitive match
{"nonexistent", false},
}
for _, tc := range testCases {
result := HasGroup(claims, tc.group)
t.Run(tc.group, func(t *testing.T) {
result := rbac.HasGroup(claims, tc.group)
if result != tc.expected {
t.Errorf("HasGroup(%s) = %v, expected %v", tc.group, result, tc.expected)
t.Errorf("HasGroup(%q) = %v; want %v", tc.group, result, tc.expected)
}
})
}
})
@@ -266,18 +270,18 @@ func TestGroupMembershipFunctions(t *testing.T) {
groups []string
expected bool
}{
{[]string{"admins"}, true},
{[]string{"editors"}, false},
{[]string{"editors", "admins"}, true},
{[]string{"admins", "editors"}, true},
{[]string{"editors", "writers"}, false},
{[]string{"ADMINS", "editors"}, true}, // Case-insensitive match
}
for _, tc := range testCases {
result := HasAnyGroup(claims, tc.groups)
t.Run(fmt.Sprintf("%v", tc.groups), func(t *testing.T) {
result := rbac.HasAnyGroup(claims, tc.groups)
if result != tc.expected {
t.Errorf("HasAnyGroup(%v) = %v, expected %v", tc.groups, result, tc.expected)
t.Errorf("HasAnyGroup(%v) = %v; want %v", tc.groups, result, tc.expected)
}
})
}
})
@@ -287,18 +291,18 @@ func TestGroupMembershipFunctions(t *testing.T) {
groups []string
expected bool
}{
{[]string{"admins"}, true},
{[]string{"admins", "developers"}, true},
{[]string{"admins", "developers", "viewers"}, true},
{[]string{"admins", "editors"}, false},
{[]string{"ADMINS", "DEVELOPERS"}, true}, // Case-insensitive match
}
for _, tc := range testCases {
result := HasAllGroups(claims, tc.groups)
t.Run(fmt.Sprintf("%v", tc.groups), func(t *testing.T) {
result := rbac.HasAllGroups(claims, tc.groups)
if result != tc.expected {
t.Errorf("HasAllGroups(%v) = %v, expected %v", tc.groups, result, tc.expected)
t.Errorf("HasAllGroups(%v) = %v; want %v", tc.groups, result, tc.expected)
}
})
}
})
}
@@ -311,58 +315,42 @@ func TestGroupMembershipFunctions(t *testing.T) {
// 3. Case-insensitive group matching
// 4. Error handling when the cognito:groups claim is missing entirely
func TestVerifyGroupMembership(t *testing.T) {
// Create test claims
claims := jwt.MapClaims{
"cognito:groups": []interface{}{"admins", "developers", "viewers"},
}
testCases := []struct {
name string
groups []string
expectError bool
errorSubstring string
expectedGroups []string
wantErr bool
}{
{
name: "All groups present",
groups: []string{"admins", "developers"},
expectError: false,
expectedGroups: []string{"admins", "developers"},
wantErr: false,
},
{
name: "Missing group",
groups: []string{"admins", "editors"},
expectError: true,
errorSubstring: "expected group editors not found",
expectedGroups: []string{"admins", "editors"},
wantErr: true,
},
{
name: "Case insensitive match",
groups: []string{"ADMINS", "Developers"},
expectError: false,
expectedGroups: []string{"ADMINS", "DEVELOPERS"},
wantErr: false,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
err := VerifyGroupMembership(claims, tc.groups)
if tc.expectError {
if err == nil {
t.Errorf("Expected error, got nil")
} else if tc.errorSubstring != "" && !contains(err.Error(), tc.errorSubstring) {
t.Errorf("Error message '%s' doesn't contain expected substring '%s'", err.Error(), tc.errorSubstring)
}
} else if err != nil {
t.Errorf("Expected no error, got %v", err)
err := rbac.VerifyGroupMembership(claims, tc.expectedGroups)
if tc.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
// Test with missing cognito:groups claim
t.Run("Missing groups claim", func(t *testing.T) {
emptyClaims := jwt.MapClaims{}
err := VerifyGroupMembership(emptyClaims, []string{"admins"})
if err == nil {
t.Errorf("Expected error for missing cognito:groups claim, got nil")
}
})
}
// TestIsValidJWKSURL tests the URL validation for JWKS endpoints.
@@ -411,7 +399,7 @@ func TestIsValidJWKSURL(t *testing.T) {
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
err := IsValidJWKSURL(tc.url)
err := rbac.IsValidJWKSURL(tc.url)
if tc.expectError && err == nil {
t.Errorf("Expected error, got nil")
} else if !tc.expectError && err != nil {
@@ -420,8 +408,3 @@ func TestIsValidJWKSURL(t *testing.T) {
})
}
}
// Helper function to check if a string contains a substring
func contains(s, substr string) bool {
return s != "" && substr != "" && s != substr && len(s) > len(substr) && s[len(s)-1] != ' ' && s[:len(s)-1] != substr
}
internal/rbac/middleware_test.go
+59 -66
View File
@@ -1,14 +1,51 @@
package rbac
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
@@ -17,23 +54,13 @@ import (
// 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()
e, keyProvider := testSetup(t)
// 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{
createTestEndpoint(e, rbac.JWTConfig{
KeyProvider: keyProvider,
RequiredGroups: []string{"foo"},
}))
})
// Test cases
tests := []struct {
@@ -61,18 +88,13 @@ func TestJWTMiddleware(t *testing.T) {
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)
token, err := rbac.GenerateTestJWT(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
// 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())
@@ -90,15 +112,10 @@ func TestJWTMiddleware(t *testing.T) {
// 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)
_, keyProvider := testSetup(t)
// Generate JWT token with 'foo' group
token, err := GenerateMockJWT(keyProvider, "test-user-id", "test@example.com", []string{"foo"})
token, err := rbac.GenerateTestJWT(keyProvider, "test-user-id", "test@example.com", []string{"foo"})
if err != nil {
t.Fatalf("Failed to generate JWT: %v", err)
}
@@ -106,13 +123,13 @@ func TestTokenExtraction(t *testing.T) {
// Test different token extraction methods
tests := []struct {
name string
config JWTConfig
config rbac.JWTConfig
setupReq func(req *http.Request)
wantStatus int
}{
{
name: "Extract from header",
config: JWTConfig{
config: rbac.JWTConfig{
KeyProvider: keyProvider,
TokenLookup: "header:Authorization",
AuthScheme: "Bearer",
@@ -125,7 +142,7 @@ func TestTokenExtraction(t *testing.T) {
},
{
name: "Extract from query parameter",
config: JWTConfig{
config: rbac.JWTConfig{
KeyProvider: keyProvider,
TokenLookup: "query:token",
RequiredGroups: []string{"foo"},
@@ -139,7 +156,7 @@ func TestTokenExtraction(t *testing.T) {
},
{
name: "Extract from cookie",
config: JWTConfig{
config: rbac.JWTConfig{
KeyProvider: keyProvider,
TokenLookup: "cookie:jwt",
RequiredGroups: []string{"foo"},
@@ -156,21 +173,12 @@ func TestTokenExtraction(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// Initialize Echo instance
// Create new Echo instance for each test to avoid config conflicts
e := echo.New()
createTestEndpoint(e, tt.config)
// 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
// 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())
@@ -190,15 +198,7 @@ func TestTokenExtraction(t *testing.T) {
// 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()
e, keyProvider := testSetup(t)
// Custom error handler
customErrorHandler := func(c echo.Context, err error) error {
@@ -208,21 +208,14 @@ func TestCustomErrorHandler(t *testing.T) {
}
// 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{
createTestEndpoint(e, rbac.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
// 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")
}
internal/rbac/testutil.go
+10
View File
@@ -0,0 +1,10 @@
package rbac
import (
"testing"
)
// TestKeyProvider creates a test key provider for testing
func TestKeyProvider(t *testing.T, privateKeyPath, publicKeyPath string) *LocalKeyProvider {
return NewLocalKeyProvider(privateKeyPath, publicKeyPath)
}
internal/serviceconfig/common.go
+3
View File
@@ -8,6 +8,8 @@ import (
"strings"
"sync"
"queryorchestration/internal/serviceconfig/rbac"
"queryorchestration/internal/serviceconfig/aws"
"queryorchestration/internal/serviceconfig/database"
"queryorchestration/internal/serviceconfig/logger"
@@ -26,6 +28,7 @@ import (
// BaseConfig provides common configuration fields and functionality
// that can be embedded in service-specific configs.
type BaseConfig struct {
rbac.AuthConfig
logger.LogConfig
observability.ObsConfig
database.DBConfig
internal/serviceconfig/rbac/config.go
+108
View File
@@ -0,0 +1,108 @@
package rbac
import (
"fmt"
"log/slog"
"queryorchestration/internal/rbac"
)
// AuthConfig holds the RBAC configuration and KeyProvider
type AuthConfig struct {
Provider rbac.KeyProvider
ProviderType string `env:"RBAC_PROVIDER" envDefault:"cognito"`
// Cognito Config
Region string `env:"COGNITO_REGION" envDefault:"us-east-1"`
UserPoolID string `env:"COGNITO_USER_POOL_ID"`
// Local Key Config
PrivateKeyPath string `env:"RBAC_PRIVATE_KEY_PATH"`
PublicKeyPath string `env:"RBAC_PUBLIC_KEY_PATH"`
}
// NewAuthConfig creates a new AuthConfig with the specified provider type
func NewAuthConfig(providerType string) (*AuthConfig, error) {
config := &AuthConfig{
ProviderType: providerType,
}
if err := config.Initialize(); err != nil {
return nil, fmt.Errorf("failed to initialize auth config: %w", err)
}
return config, nil
}
// NewLocalAuthConfig creates a new AuthConfig configured for local key provider
func NewLocalAuthConfig(privateKeyPath, publicKeyPath string) (*AuthConfig, error) {
config := &AuthConfig{
ProviderType: "local",
PrivateKeyPath: privateKeyPath,
PublicKeyPath: publicKeyPath,
}
if err := config.Initialize(); err != nil {
return nil, fmt.Errorf("failed to initialize local auth config: %w", err)
}
return config, nil
}
// ConfigProvider interface defines the methods required for RBAC configuration
type ConfigProvider interface {
//Initialize() error
GetKeyProvider() rbac.KeyProvider
PrintConfig()
}
// GetKeyProvider returns the configured KeyProvider
func (r *AuthConfig) GetKeyProvider() rbac.KeyProvider {
return r.Provider
}
// Initialize initializes the appropriate KeyProvider based on configuration
func (r *AuthConfig) Initialize() error {
switch r.ProviderType {
case "local":
if r.PrivateKeyPath == "" || r.PublicKeyPath == "" {
return fmt.Errorf("local key provider requires both private and public key paths")
}
r.Provider = rbac.NewLocalKeyProvider(r.PrivateKeyPath, r.PublicKeyPath)
slog.Info("RBAC initialized with local key provider")
case "cognito":
if r.UserPoolID == "" {
return fmt.Errorf("cognito provider requires user pool ID")
}
r.Provider = rbac.NewCognitoKeyProvider(r.Region, r.UserPoolID)
slog.Info("RBAC initialized with Cognito key provider")
default:
return fmt.Errorf("unknown RBAC provider type: %s", r.ProviderType)
}
return nil
}
// PrintConfig logs the current configuration, masking sensitive values
func (r *AuthConfig) PrintConfig() {
slog.Info("RBAC Configuration",
"provider_type", r.ProviderType,
"region", r.Region,
"user_pool_id", r.UserPoolID,
"private_key_path", r.PrivateKeyPath,
"public_key_path", r.PublicKeyPath,
)
}
// Validate checks if the configuration is valid
func (r *AuthConfig) Validate() error {
switch r.ProviderType {
case "local":
if r.PrivateKeyPath == "" || r.PublicKeyPath == "" {
return fmt.Errorf("local key provider requires both private and public key paths")
}
case "cognito":
if r.UserPoolID == "" {
return fmt.Errorf("cognito provider requires user pool ID")
}
default:
return fmt.Errorf("invalid provider type: %s", r.ProviderType)
}
return nil
}
internal/serviceconfig/rbac/config_test.go
+15
View File
@@ -0,0 +1,15 @@
package rbac
import (
"testing"
)
// TestAuthConfig creates an initialized AuthConfig for testing
func TestAuthConfig(t *testing.T, privateKeyPath, publicKeyPath string) *AuthConfig {
config, err := NewLocalAuthConfig(privateKeyPath, publicKeyPath)
if err != nil {
t.Fatalf("Failed to initialize auth config: %v", err)
}
return config
}
internal/serviceconfig/rbac/testutil.go
+21
View File
@@ -0,0 +1,21 @@
package rbac
import (
"testing"
"queryorchestration/internal/rbac"
)
// TestKeyProvider creates a test key provider for testing
func TestKeyProvider(t *testing.T, privateKeyPath, publicKeyPath string) rbac.KeyProvider {
config := &AuthConfig{
ProviderType: "local",
PrivateKeyPath: privateKeyPath,
PublicKeyPath: publicKeyPath,
}
err := config.Initialize()
if err != nil {
t.Fatalf("Failed to initialize auth config: %v", err)
}
return config.GetKeyProvider()
}