diff --git a/cmd/cognito_test/cognitotest/main.go b/cmd/cognito_test/cognitotest/main.go new file mode 100644 index 00000000..bd2adb61 --- /dev/null +++ b/cmd/cognito_test/cognitotest/main.go @@ -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 "" + } + if len(s) <= 8 { + return "****" + } + return s[:4] + "..." + s[len(s)-4:] +}