test harness

2025-03-20 15:37:10 -07:00
parent 99e1a51b53
commit b9da2e73c0
1 changed files with 356 additions and 0 deletions
@@ -0,0 +1,356 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"os"
+	"strings"
+	"time"
+
+	"github.com/lestrrat-go/jwx/v2/jwk"
+	"github.com/lestrrat-go/jwx/v2/jwt"
+)
+
+// CognitoTokenResponse represents the response from the token endpoint
+type CognitoTokenResponse struct {
+	IDToken      string `json:"id_token"`
+	AccessToken  string `json:"access_token"`
+	RefreshToken string `json:"refresh_token"`
+	ExpiresIn    int    `json:"expires_in"`
+	TokenType    string `json:"token_type"`
+}
+
+// CognitoConfig holds the configuration for Cognito
+type CognitoConfig struct {
+	ClientID     string
+	ClientSecret string
+	RedirectURI  string
+	TokenURL     string
+	JwksURL      string
+	UserPoolID   string
+}
+
+// VerifyCognitoCode exchanges the authorization code for tokens and verifies them
+func VerifyCognitoCode(authCode string, jwksJSON string) (map[string]interface{}, error) {
+	// Get client credentials from environment variables
+	clientID := os.Getenv("COGNITO_CLIENT_ID")
+	clientSecret := os.Getenv("COGNITO_CLIENT_SECRET")
+	userPoolID := os.Getenv("COGNITO_USER_POOL_ID")
+	region := os.Getenv("AWS_REGION")
+
+	if clientID == "" {
+		return nil, errors.New("COGNITO_CLIENT_ID environment variable is not set")
+	}
+
+	if userPoolID == "" {
+		userPoolID = "us-east-2_1y6po8rR8" // Fallback to the hardcoded value
+	}
+
+	if region == "" {
+		region = "us-east-2" // Fallback to the hardcoded value
+	}
+
+	// Get the domain name from environment variable or use a default
+	domain := os.Getenv("COGNITO_DOMAIN")
+	if domain == "" {
+		// Log warning about missing domain
+		fmt.Println("Warning: COGNITO_DOMAIN environment variable is not set, using default")
+		domain = fmt.Sprintf("%s.auth.%s.amazoncognito.com", userPoolID, region)
+	}
+
+	// Ensure domain format is correct (no protocol prefix)
+	domain = strings.TrimPrefix(domain, "https://")
+	domain = strings.TrimPrefix(domain, "http://")
+
+	config := CognitoConfig{
+		ClientID:     clientID,
+		ClientSecret: clientSecret,
+		RedirectURI:  "http://localhost:8080/query",
+		TokenURL:     fmt.Sprintf("https://%s/oauth2/token", domain),
+		JwksURL:      fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s/.well-known/jwks.json", region, userPoolID),
+		UserPoolID:   userPoolID,
+	}
+
+	// Debug output
+	fmt.Printf("Configuration:\n")
+	fmt.Printf("  Client ID: %s\n", config.ClientID)
+	fmt.Printf("  Using client secret: %v\n", config.ClientSecret != "")
+	fmt.Printf("  Domain: %s\n", domain)
+	fmt.Printf("  Token URL: %s\n", config.TokenURL)
+	fmt.Printf("  JWKS URL: %s\n", config.JwksURL)
+
+	// Exchange auth code for tokens
+	tokens, err := exchangeAuthCodeForTokens(authCode, config)
+	if err != nil {
+		return nil, fmt.Errorf("failed to exchange auth code for tokens: %w", err)
+	}
+
+	// Parse JWKS JSON
+	keySet, err := jwk.ParseString(jwksJSON)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse JWKS: %w", err)
+	}
+
+	// Verify the ID token
+	idTokenClaims, err := verifyToken(tokens.IDToken, keySet, config)
+	if err != nil {
+		return nil, fmt.Errorf("failed to verify ID token: %w", err)
+	}
+
+	return idTokenClaims, nil
+}
+
+// exchangeAuthCodeForTokens exchanges the authorization code for tokens
+func exchangeAuthCodeForTokens(authCode string, config CognitoConfig) (*CognitoTokenResponse, error) {
+	data := url.Values{}
+	data.Set("grant_type", "authorization_code")
+	data.Set("client_id", config.ClientID)
+	data.Set("code", authCode)
+	data.Set("redirect_uri", config.RedirectURI)
+
+	// Debug output
+	fmt.Printf("Token request details:\n")
+	fmt.Printf("  URL: %s\n", config.TokenURL)
+	fmt.Printf("  client_id: %s\n", config.ClientID)
+	fmt.Printf("  redirect_uri: %s\n", config.RedirectURI)
+	fmt.Printf("  grant_type: authorization_code\n")
+
+	req, err := http.NewRequest("POST", config.TokenURL, strings.NewReader(data.Encode()))
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
+
+	// Add Authorization header if client secret is provided
+	if config.ClientSecret != "" {
+		req.SetBasicAuth(config.ClientID, config.ClientSecret)
+		fmt.Println("  Using Basic Auth authentication")
+	} else {
+		fmt.Println("  Using public client authentication (no secret)")
+	}
+
+	client := &http.Client{Timeout: 10 * time.Second}
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("token endpoint returned status %d: %s", resp.StatusCode, string(body))
+	}
+
+	var tokenResponse CognitoTokenResponse
+	if err := json.Unmarshal(body, &tokenResponse); err != nil {
+		return nil, err
+	}
+
+	return &tokenResponse, nil
+}
+
+// verifyToken verifies the JWT token and returns its claims
+func verifyToken(tokenString string, keySet jwk.Set, config CognitoConfig) (map[string]interface{}, error) {
+	region := "us-east-2"
+	if strings.Contains(config.JwksURL, "amazonaws.com/") {
+		parts := strings.Split(config.JwksURL, "amazonaws.com/")
+		if len(parts) > 1 {
+			regionParts := strings.Split(parts[0], ".")
+			if len(regionParts) > 0 {
+				lastPart := regionParts[len(regionParts)-1]
+				if lastPart != "" {
+					region = lastPart
+				}
+			}
+		}
+	}
+
+	issuer := fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s", region, config.UserPoolID)
+	fmt.Printf("Using issuer: %s\n", issuer)
+
+	// Verify the token with the keySet
+	verifiedToken, err := jwt.Parse(
+		[]byte(tokenString),
+		jwt.WithKeySet(keySet),
+		jwt.WithValidate(true),
+		jwt.WithIssuer(issuer),
+		// Add other validation options as needed
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	// Extract claims to a map
+	claims, err := verifiedToken.AsMap(context.Background())
+	if err != nil {
+		return nil, err
+	}
+
+	return claims, nil
+}
+
+// GetUserGroups extracts user groups from the token claims
+func GetUserGroups(claims map[string]interface{}) ([]string, error) {
+	// The claim containing the groups might have different names depending on your Cognito setup
+	// Common names are "cognito:groups", "groups", or a custom attribute
+	possibleGroupFields := []string{"cognito:groups", "groups", "custom:groups"}
+
+	for _, field := range possibleGroupFields {
+		if rawGroups, ok := claims[field]; ok {
+			switch groups := rawGroups.(type) {
+			case []interface{}:
+				result := make([]string, len(groups))
+				for i, g := range groups {
+					result[i] = fmt.Sprintf("%v", g)
+				}
+				return result, nil
+			case []string:
+				return groups, nil
+			case string:
+				return []string{groups}, nil
+			}
+		}
+	}
+
+	return nil, errors.New("no groups found in token claims")
+}
+
+// Example usage in an HTTP handler
+func HandleCallback(w http.ResponseWriter, r *http.Request) {
+	// Extract the authorization code from the query parameters
+	code := r.URL.Query().Get("code")
+	if code == "" {
+		http.Error(w, "Missing code parameter", http.StatusBadRequest)
+		return
+	}
+
+	// Extract error if present
+	errorMsg := r.URL.Query().Get("error")
+	errorDescription := r.URL.Query().Get("error_description")
+	if errorMsg != "" {
+		http.Error(w, fmt.Sprintf("Authorization error: %s - %s", errorMsg, errorDescription), http.StatusBadRequest)
+		return
+	}
+
+	// Get user pool ID and region
+	userPoolID := os.Getenv("COGNITO_USER_POOL_ID")
+	region := os.Getenv("AWS_REGION")
+
+	if userPoolID == "" {
+		userPoolID = "us-east-2_1y6po8rR8" // Fallback to hardcoded value
+	}
+
+	if region == "" {
+		region = "us-east-2" // Fallback to hardcoded value
+	}
+
+	// Fetch the JWKS from Cognito
+	jwksURL := fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s/.well-known/jwks.json", region, userPoolID)
+	fmt.Printf("Fetching JWKS from: %s\n", jwksURL)
+
+	// Fix G107: Use a safe, validated URL for the HTTP request
+	if !strings.HasPrefix(jwksURL, "https://cognito-idp.") || !strings.Contains(jwksURL, ".amazonaws.com/") {
+		http.Error(w, "Invalid JWKS URL", http.StatusInternalServerError)
+		return
+	}
+
+	// Create a custom client with a timeout to avoid gosec G107 warning
+	client := &http.Client{
+		Timeout: 10 * time.Second,
+	}
+	req, err := http.NewRequest("GET", jwksURL, nil)
+	if err != nil {
+		http.Error(w, fmt.Sprintf("Failed to create request: %v", err), http.StatusInternalServerError)
+		return
+	}
+	resp, err := client.Do(req)
+	if err != nil {
+		http.Error(w, fmt.Sprintf("Failed to fetch JWKS: %v", err), http.StatusInternalServerError)
+		return
+	}
+	defer resp.Body.Close()
+
+	jwksJSON, err := io.ReadAll(resp.Body)
+	if err != nil {
+		http.Error(w, fmt.Sprintf("Failed to read JWKS response: %v", err), http.StatusInternalServerError)
+		return
+	}
+
+	// Verify the code and get the token claims
+	claims, err := VerifyCognitoCode(code, string(jwksJSON))
+	if err != nil {
+		http.Error(w, fmt.Sprintf("Failed to verify code: %v", err), http.StatusUnauthorized)
+		return
+	}
+
+	// Extract user groups
+	groups, err := GetUserGroups(claims)
+	if err != nil {
+		// Handle the case where groups can't be found, but the token is valid
+		// You might still want to proceed with authenticated but ungrouped user
+		fmt.Printf("Warning: %v\n", err)
+	}
+
+	// Do something with the verified claims and groups
+	// For example, set session cookies, redirect the user, etc.
+	w.Header().Set("Content-Type", "application/json")
+	response := map[string]interface{}{
+		"authenticated": true,
+		"username":      claims["cognito:username"],
+		"email":         claims["email"],
+		"groups":        groups,
+	}
+
+	// Fix line 299: Check the error from Encode
+	if err := json.NewEncoder(w).Encode(response); err != nil {
+		http.Error(w, fmt.Sprintf("Failed to encode response: %v", err), http.StatusInternalServerError)
+		return
+	}
+}
+
+func main() {
+	// Set up your HTTP server
+	http.HandleFunc("/query", HandleCallback)
+	fmt.Println("Server starting on :8080")
+	fmt.Println("Using environment variables:")
+	fmt.Printf("  COGNITO_CLIENT_ID: %s\n", maskString(os.Getenv("COGNITO_CLIENT_ID")))
+	fmt.Printf("  COGNITO_CLIENT_SECRET: %s\n", maskString(os.Getenv("COGNITO_CLIENT_SECRET")))
+	fmt.Printf("  COGNITO_DOMAIN: %s\n", os.Getenv("COGNITO_DOMAIN"))
+	fmt.Printf("  COGNITO_USER_POOL_ID: %s\n", os.Getenv("COGNITO_USER_POOL_ID"))
+	fmt.Printf("  AWS_REGION: %s\n", os.Getenv("AWS_REGION"))
+
+	// Fix G114: Use a server with timeout support instead of http.ListenAndServe
+	server := &http.Server{
+		Addr:              ":8080",
+		ReadHeaderTimeout: 3 * time.Second,
+		ReadTimeout:       20 * time.Second,
+		WriteTimeout:      20 * time.Second,
+		IdleTimeout:       120 * time.Second,
+	}
+
+	if err := server.ListenAndServe(); err != nil {
+		fmt.Printf("Server failed to start: %v\n", err)
+		os.Exit(1)
+	}
+}
+
+// Helper function to mask sensitive values
+func maskString(s string) string {
+	if s == "" {
+		return "<not set>"
+	}
+	if len(s) <= 8 {
+		return "****"
+	}
+	return s[:4] + "..." + s[len(s)-4:]
+}