From 95e31283b5a0b05e23c841d6358199812895538e Mon Sep 17 00:00:00 2001
From: jay brown <jay.a.brown@gmail.com>
Date: Fri, 21 Mar 2025 13:39:51 -0700
Subject: [PATCH] convert harness to echo

---
 cmd/cognito_test/cognitotest/main.go | 64 +++++++++++++---------------
 1 file changed, 30 insertions(+), 34 deletions(-)

diff --git a/cmd/cognito_test/cognitotest/main.go b/cmd/cognito_test/cognitotest/main.go
index bd2adb61..1eef3b05 100644
--- a/cmd/cognito_test/cognitotest/main.go
+++ b/cmd/cognito_test/cognitotest/main.go
@@ -12,6 +12,8 @@ import (
 	"strings"
 	"time"
 
+	"github.com/labstack/echo/v4"
+	"github.com/labstack/echo/v4/middleware"
 	"github.com/lestrrat-go/jwx/v2/jwk"
 	"github.com/lestrrat-go/jwx/v2/jwt"
 )
@@ -225,21 +227,19 @@ func GetUserGroups(claims map[string]interface{}) ([]string, error) {
 	return nil, errors.New("no groups found in token claims")
 }
 
-// Example usage in an HTTP handler
-func HandleCallback(w http.ResponseWriter, r *http.Request) {
+// HandleCallback is the Echo handler for OAuth callback
+func HandleCallback(c echo.Context) error {
 	// Extract the authorization code from the query parameters
-	code := r.URL.Query().Get("code")
+	code := c.QueryParam("code")
 	if code == "" {
-		http.Error(w, "Missing code parameter", http.StatusBadRequest)
-		return
+		return c.String(http.StatusBadRequest, "Missing code parameter")
 	}
 
 	// Extract error if present
-	errorMsg := r.URL.Query().Get("error")
-	errorDescription := r.URL.Query().Get("error_description")
+	errorMsg := c.QueryParam("error")
+	errorDescription := c.QueryParam("error_description")
 	if errorMsg != "" {
-		http.Error(w, fmt.Sprintf("Authorization error: %s - %s", errorMsg, errorDescription), http.StatusBadRequest)
-		return
+		return c.String(http.StatusBadRequest, fmt.Sprintf("Authorization error: %s - %s", errorMsg, errorDescription))
 	}
 
 	// Get user pool ID and region
@@ -260,8 +260,7 @@ func HandleCallback(w http.ResponseWriter, r *http.Request) {
 
 	// 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
+		return c.String(http.StatusInternalServerError, "Invalid JWKS URL")
 	}
 
 	// Create a custom client with a timeout to avoid gosec G107 warning
@@ -270,27 +269,23 @@ func HandleCallback(w http.ResponseWriter, r *http.Request) {
 	}
 	req, err := http.NewRequest("GET", jwksURL, nil)
 	if err != nil {
-		http.Error(w, fmt.Sprintf("Failed to create request: %v", err), http.StatusInternalServerError)
-		return
+		return c.String(http.StatusInternalServerError, fmt.Sprintf("Failed to create request: %v", err))
 	}
 	resp, err := client.Do(req)
 	if err != nil {
-		http.Error(w, fmt.Sprintf("Failed to fetch JWKS: %v", err), http.StatusInternalServerError)
-		return
+		return c.String(http.StatusInternalServerError, fmt.Sprintf("Failed to fetch JWKS: %v", err))
 	}
 	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
+		return c.String(http.StatusInternalServerError, fmt.Sprintf("Failed to read JWKS response: %v", err))
 	}
 
 	// 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
+		return c.String(http.StatusUnauthorized, fmt.Sprintf("Failed to verify code: %v", err))
 	}
 
 	// Extract user groups
@@ -301,9 +296,7 @@ func HandleCallback(w http.ResponseWriter, r *http.Request) {
 		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")
+	// Prepare response
 	response := map[string]interface{}{
 		"authenticated": true,
 		"username":      claims["cognito:username"],
@@ -311,16 +304,21 @@ func HandleCallback(w http.ResponseWriter, r *http.Request) {
 		"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
-	}
+	return c.JSON(http.StatusOK, response)
 }
 
 func main() {
-	// Set up your HTTP server
-	http.HandleFunc("/query", HandleCallback)
+	// Create a new Echo instance
+	e := echo.New()
+
+	// Middleware
+	e.Use(middleware.Logger())
+	e.Use(middleware.Recover())
+
+	// Routes
+	e.GET("/query", HandleCallback)
+
+	// Print startup information
 	fmt.Println("Server starting on :8080")
 	fmt.Println("Using environment variables:")
 	fmt.Printf("  COGNITO_CLIENT_ID: %s\n", maskString(os.Getenv("COGNITO_CLIENT_ID")))
@@ -329,7 +327,7 @@ func main() {
 	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
+	// Configure server
 	server := &http.Server{
 		Addr:              ":8080",
 		ReadHeaderTimeout: 3 * time.Second,
@@ -338,10 +336,8 @@ func main() {
 		IdleTimeout:       120 * time.Second,
 	}
 
-	if err := server.ListenAndServe(); err != nil {
-		fmt.Printf("Server failed to start: %v\n", err)
-		os.Exit(1)
-	}
+	// Start server with custom server config
+	e.Logger.Fatal(e.StartServer(server))
 }
 
 // Helper function to mask sensitive values