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

comments

Browse Source
This commit is contained in:
jay brown
2025-03-17 13:40:43 -07:00
parent 3c1a77edbb
commit 560433af6d
3 changed files with 77 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/cognito_test/jwtstuff/main.go
-185
View File
@@ -1,185 +0,0 @@
package main
import (
"crypto/rsa"
"encoding/json"
"errors"
"fmt"
"os"
"time"
"github.com/golang-jwt/jwt/v5"
)
// Load private key from file
func loadPrivateKey() (*rsa.PrivateKey, error) {
keyData, err := os.ReadFile("private_key.pem")
if err != nil {
return nil, err
}
return jwt.ParseRSAPrivateKeyFromPEM(keyData)
}
// Generate a mock JWT with multiple Cognito groups
func generateMockJWT() (string, error) {
privateKey, err := loadPrivateKey()
if err != nil {
return "", err
}
/* payload will look like this when verified
{
"cognito:groups": [
"admins",
"developers",
"querybuilders",
"uploaders",
"exporters"
],
"email": "test@example.com",
"exp": 1741820939,
"iat": 1741817339,
"iss": "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_XXXXXXXXX",
"sub": "test-user-id"
}
*/
token := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.MapClaims{
"sub": "test-user-id",
"email": "test@example.com",
"exp": time.Now().Add(time.Hour).Unix(),
"iat": time.Now().Unix(),
"iss": "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_XXXXXXXXX",
"cognito:groups": []string{"admins", "developers", "querybuilders", "uploaders", "exporters"}, // Multiple groups
})
signedToken, errorSigning := token.SignedString(privateKey)
if errorSigning != nil {
return "", errorSigning
}
fmt.Printf("raw token: %s\n", signedToken)
return signedToken, nil
}
// Load public key from file
func loadPublicKey() (*rsa.PublicKey, error) {
keyData, err := os.ReadFile("public_key.pem")
if err != nil {
return nil, err
}
return jwt.ParseRSAPublicKeyFromPEM(keyData)
}
func printTokenDetails(token *jwt.Token) {
fmt.Println("\n=== JWT Token Details ===")
// Print Header
headerJSON, _ := json.MarshalIndent(token.Header, "", " ")
fmt.Printf("\nHeader:\n%s\n", string(headerJSON))
// Print Claims
claimsJSON, _ := json.MarshalIndent(token.Claims, "", " ")
fmt.Printf("\nClaims:\n%s\n", string(claimsJSON))
// Print Signature
fmt.Printf("\nSignature: %s\n", token.Signature)
// Print Method
fmt.Printf("\nSigning Method: %s\n", token.Method.Alg())
// Print Valid status
fmt.Printf("\nValid: %v\n", token.Valid)
fmt.Println("\n=====================")
}
// Validate JWT and extract claims
func validateJWT(tokenString string) (*jwt.Token, jwt.MapClaims, error) {
publicKey, err := loadPublicKey()
if err != nil {
return nil, nil, err
}
token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
return nil, errors.New("unexpected signing method")
}
return publicKey, nil
})
// pretty print all of the details of the token
printTokenDetails(token)
if err != nil {
return nil, nil, err
}
if claims, ok := token.Claims.(jwt.MapClaims); ok && token.Valid {
return token, claims, nil
}
return nil, nil, errors.New("invalid token")
}
func main() {
tokenString, err := generateMockJWT()
if err != nil {
fmt.Println("Error generating token:", err)
return
}
token, claims, err := validateJWT(tokenString)
if err != nil {
fmt.Println("Token validation failed:", err)
return
}
if token.Valid {
fmt.Println("Token is valid.")
fmt.Println("User ID:", claims["sub"])
fmt.Println("Email:", claims["email"])
if groups, ok := claims["cognito:groups"].([]interface{}); ok {
fmt.Println("User Groups:")
for _, group := range groups {
fmt.Println("-", group)
}
} else {
fmt.Println("No groups found in token.")
}
} else {
fmt.Println("Invalid token")
}
}
//4. Unit Test to Verify Group Membership
//
//Add a unit test to verify if the JWT contains expected groups.
//
//func TestJWTGroups(t *testing.T) {
// tokenString, _ := generateMockJWT()
//
// _, claims, err := validateJWT(tokenString)
// if err != nil {
// t.Errorf("Token validation failed: %v", err)
// }
//
// expectedGroups := []string{"Admins", "Developers", "BetaTesters"}
// groups, ok := claims["cognito:groups"].([]interface{})
// if !ok {
// t.Errorf("Expected 'cognito:groups' claim, but not found")
// }
//
// for _, eg := range expectedGroups {
// found := false
// for _, g := range groups {
// if g.(string) == eg {
// found = true
// break
// }
// }
// if !found {
// t.Errorf("Expected group %s not found in token", eg)
// }
// }
//}
internal/rbac/cognito_test.go
+47
View File
@@ -8,6 +8,8 @@ import (
"github.com/golang-jwt/jwt/v5"
)
// 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")
@@ -30,6 +32,11 @@ func TestLocalKeyProvider(t *testing.T) {
}
}
// TestGenerateAndValidateJWT tests the complete JWT lifecycle:
// 1. Generating a mock JWT token with specific claims (userID, email, groups)
// 2. Validating the token with the same key provider
// 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")
@@ -102,6 +109,11 @@ func TestGenerateAndValidateJWT(t *testing.T) {
}
}
// TestTokenExpiration verifies that token expiration works correctly:
// 1. A token with a future expiration time should be considered valid
// 2. A token with a past expiration time would be rejected by the JWT library
// 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")
@@ -167,6 +179,13 @@ func TestTokenExpiration(t *testing.T) {
*/
}
// TestExtractTokenDetails verifies that the ExtractTokenDetails function
// correctly extracts and organizes user information from the JWT token:
// - UserID
// - Email
// - Group memberships
// - 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")
@@ -207,6 +226,11 @@ func TestExtractTokenDetails(t *testing.T) {
}
}
// TestGroupMembershipFunctions validates the three group membership check functions:
// 1. HasGroup - Tests if a user belongs to a specific group
// 2. HasAnyGroup - Tests if a user belongs to at least one of several groups
// 3. HasAllGroups - Tests if a user belongs to all specified groups
// Each test includes positive and negative cases, as well as case-insensitive matching.
func TestGroupMembershipFunctions(t *testing.T) {
// Create test claims
claims := jwt.MapClaims{
@@ -279,6 +303,13 @@ func TestGroupMembershipFunctions(t *testing.T) {
})
}
// TestVerifyGroupMembership validates the VerifyGroupMembership function,
// which ensures that a user belongs to all required groups.
// It tests:
// 1. Successful case where all required groups are present
// 2. Error case where a required group is missing
// 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{
@@ -334,6 +365,22 @@ func TestVerifyGroupMembership(t *testing.T) {
})
}
// TestIsValidJWKSURL tests the URL validation for JWKS endpoints.
//
// JWKS (JSON Web Key Set) is a standard format used to publish public keys as a set of JWKs (JSON Web Keys).
// In the AWS Cognito context, JWKS endpoints allow services to retrieve the public keys needed to verify
// JWT tokens issued by Cognito. These endpoints must be secured (HTTPS) and trusted to prevent
// token forgery attacks.
//
// This test ensures that:
// 1. Valid Cognito JWKS URLs are accepted
// 2. Insecure URLs (HTTP instead of HTTPS) are rejected
// 3. URLs from untrusted domains are rejected
// 4. Malformed URLs are rejected
//
// This validation is critical for preventing security issues like malicious redirects,
// SSRF (Server-Side Request Forgery) attacks, or using untrusted key sources for JWT verification.
// This is important for preventing malicious redirects or SSRF attacks.
func TestIsValidJWKSURL(t *testing.T) {
testCases := []struct {
name string
internal/rbac/middleware_test.go
+30 -5
View File
@@ -9,6 +9,13 @@ import (
"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"
@@ -74,7 +81,14 @@ func TestJWTMiddleware(t *testing.T) {
}
}
// Test different token extraction methods
// 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"
@@ -165,7 +179,16 @@ func TestTokenExtraction(t *testing.T) {
}
}
// Test the custom error handler
// 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"
@@ -204,6 +227,8 @@ func TestCustomErrorHandler(t *testing.T) {
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.
// 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.