query-orchestration/internal/serviceconfig/auth/config.go

package auth

import (
	"fmt"
	"log"
	"log/slog"
	"os"
	"strings"
)

// CognitoConfig holds the configuration for AWS Cognito
// Contains all necessary parameters to interact with Cognito endpoints
type CognitoConfig struct {
	// annotated vars
	AuthRegion       string `env:"COGNITO_REGION" envDefault:"us-east-2"`
	AuthUserPoolID   string `env:"COGNITO_USER_POOL_ID,required,notEmpty"`
	AuthClientSecret string `env:"COGNITO_CLIENT_SECRET,required,notEmpty"`
	AuthDomain       string `env:"COGNITO_DOMAIN,required,notEmpty"`
	AuthClientID     string `env:"COGNITO_CLIENT_ID,required,notEmpty"`

	AuthRedirectURI      string              // URL where Cognito redirects after authentication
	AuthTokenURL         string              // Cognito endpoint for token operations
	AuthJwksURL          string              // URL for JSON Web Key Set (for token verification)
	AuthURL              string              // Cognito authorization endpoint for initiating login
	AuthLoginPath        string              // Path for initiating login
	AuthCallbackPath     string              // Path for OAuth callback
	AuthHomePath         string              // Path for home page
	AuthLogoutPath       string              // Path for logout
	AuthLogger           *slog.Logger        // Logger instance ( normally just the same logger as main)
	AuthRoutePermissions map[string][]string // Map of routes to required permissions
}

// InitializeConfig creates a new CognitoConfig from environment variables
// and sets the values that are computed from other values
func InitializeConfig(baseURL string, logger *slog.Logger) *CognitoConfig {
	// Set default logger if not provided
	if logger == nil {
		var logLevel slog.Level
		if os.Getenv("DEBUG") == "true" {
			logLevel = slog.LevelDebug
		} else {
			logLevel = slog.LevelInfo
		}

		logHandler := slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
			Level: logLevel,
		})
		logger = slog.New(logHandler)
	}

	// Set paths with defaults
	loginPath := os.Getenv("COGNITO_LOGIN_PATH")
	if loginPath == "" {
		loginPath = "/login"
	}

	callbackPath := os.Getenv("COGNITO_CALLBACK_PATH")
	if callbackPath == "" {
		callbackPath = "/login-callback"
	}

	homePath := os.Getenv("COGNITO_HOME_PATH")
	if homePath == "" {
		homePath = "/home"
	}

	logoutPath := os.Getenv("COGNITO_LOGOUT_PATH")
	if logoutPath == "" {
		logoutPath = "/logout"
	}

	// Get AWS region with fallback
	region := os.Getenv("AWS_REGION")
	if region == "" {
		region = "us-east-2" // Fallback to default value
	}

	// Get User Pool ID with fallback
	userPoolID := os.Getenv("COGNITO_USER_POOL_ID")
	if userPoolID == "" {
		userPoolID = "" // No default, will need to be provided
		logger.Warn("COGNITO_USER_POOL_ID environment variable is not set")
	}

	// Get Cognito domain
	domain := os.Getenv("COGNITO_DOMAIN")
	if domain == "" {
		//logger.Warn("COGNITO_DOMAIN environment variable is not set, using default")
		//// Use the user pool ID directly as part of the domain
		//domain = fmt.Sprintf("%s.auth.%s.amazoncognito.com", userPoolID, region)
		log.Fatalf("COGNITO_DOMAIN environment variable is not set")
	}

	// Ensure domain format is correct (no protocol prefix)
	domain = strings.TrimPrefix(domain, "https://")
	domain = strings.TrimPrefix(domain, "http://")

	// Sanitize base URL (remove trailing slash)
	baseURL = strings.TrimSuffix(baseURL, "/")

	config := &CognitoConfig{
		AuthClientID:         os.Getenv("COGNITO_CLIENT_ID"),
		AuthClientSecret:     os.Getenv("COGNITO_CLIENT_SECRET"),
		AuthRedirectURI:      baseURL + callbackPath,
		AuthTokenURL:         fmt.Sprintf("https://%s/oauth2/token", domain),
		AuthURL:              fmt.Sprintf("https://%s/oauth2/authorize", domain),
		AuthJwksURL:          fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s/.well-known/jwks.json", region, userPoolID),
		AuthUserPoolID:       userPoolID,
		AuthRegion:           region,
		AuthLoginPath:        loginPath,
		AuthCallbackPath:     callbackPath,
		AuthHomePath:         homePath,
		AuthLogoutPath:       logoutPath,
		AuthLogger:           logger,
		AuthRoutePermissions: make(map[string][]string),
	}

	config.PrettyPrint()
	return config
}

func (c *CognitoConfig) PrettyPrint() {
	fmt.Printf("CognitoConfig:\n")
	fmt.Printf("  AuthClientID: %s\n", c.AuthClientID)
	fmt.Printf("  AuthClientSecret: %s\n", c.AuthClientSecret)
	fmt.Printf("  AuthRedirectURI: %s\n", c.AuthRedirectURI)
	fmt.Printf("  AuthTokenURL: %s\n", c.AuthTokenURL)
	fmt.Printf("  AuthURL: %s\n", c.AuthURL)
	fmt.Printf("  AuthJwksURL: %s\n", c.AuthJwksURL)
	fmt.Printf("  AuthUserPoolID: %s\n", c.AuthUserPoolID)
	fmt.Printf("  AuthRegion: %s\n", c.AuthRegion)
	fmt.Printf("  AuthLoginPath: %s\n", c.AuthLoginPath)
	fmt.Printf("  AuthCallbackPath: %s\n", c.AuthCallbackPath)
	fmt.Printf("  AuthHomePath: %s\n", c.AuthHomePath)
	fmt.Printf("  AuthLogoutPath: %s\n", c.AuthLogoutPath)

	// pretty print the route permissions map
	for route, permissions := range c.AuthRoutePermissions {
		fmt.Printf("  Route: %s, Permissions: %v\n", route, permissions)
	}

}

// SetRoutePermissions sets the route permissions map
func (c *CognitoConfig) SetRoutePermissions(permissions map[string][]string) {
	c.AuthRoutePermissions = permissions
}