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

package auth

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

	"github.com/caarlos0/env/v11"

	"queryorchestration/internal/usermanagement"
)

// 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"`
	DisableAuth      string `env:"DISABLE_AUTH" envDefault:"false"`
	NoJWTValidation  bool   `env:"NO_JWT_VALIDATION" envDefault:"false"`

	// Permit.io configuration fields
	PermitIOAPIKey string `env:"PERMIT_IO_API_KEY"`
	PermitIOPDPURL string `env:"PERMIT_IO_PDP_URL" envDefault:"http://localhost:7766"`

	// CognitoEnvironmentID scopes this server instance to a specific environment within
	// a shared Cognito user pool. When set, only users tagged with this environment ID
	// (or untagged users) are visible. Empty string means no scoping (backward compat).
	CognitoEnvironmentID string `env:"COGNITO_ENVIRONMENT_ID" envDefault:""`

	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
	AuthLogoutURL        string              // Cognito logout endpoint for session termination
	AuthLoginPath        string              `env:"COGNITO_LOGIN_PATH" envDefault:"/login"`
	AuthCallbackPath     string              `env:"COGNITO_CALLBACK_PATH" envDefault:"/login-callback"`
	AuthHomePath         string              `env:"COGNITO_HOME_PATH" envDefault:"/home"`
	AuthLogoutPath       string              `env:"COGNITO_LOGOUT_PATH" envDefault:"/logout"`
	HealthPath           string              `env:"HEALTH_PATH" envDefault:"/health"`
	AuthLogger           *slog.Logger        // Logger instance ( normally just the same logger as main)
	AuthRoutePermissions map[string][]string // Map of routes to required permissions
}

type ConfigProvider interface {
	// Env annotated fields
	GetAuthRegion() string
	SetAuthRegion(string)
	GetAuthUserPoolID() string
	SetAuthUserPoolID(string)
	GetAuthClientSecret() string
	SetAuthClientSecret(string)
	GetAuthDomain() string
	SetAuthDomain(string)
	GetAuthClientID() string
	SetAuthClientID(string)
	GetNoJWTValidation() bool
	SetNoJWTValidation(bool)

	// Permit.io configuration methods
	GetPermitIOAPIKey() string
	SetPermitIOAPIKey(string)
	GetPermitIOPDPURL() string
	SetPermitIOPDPURL(string)

	// Environment isolation methods
	GetCognitoEnvironmentID() string
	SetCognitoEnvironmentID(string)

	// Non-env fields
	GetAuthRedirectURI() string
	SetAuthRedirectURI(string)
	GetAuthTokenURL() string
	SetAuthTokenURL(string)
	GetAuthJwksURL() string
	SetAuthJwksURL(string)
	GetAuthURL() string
	SetAuthURL(string)
	GetAuthLogoutURL() string
	SetAuthLogoutURL(string)
	GetAuthLoginPath() string
	SetAuthLoginPath(string)
	GetAuthCallbackPath() string
	SetAuthCallbackPath(string)
	GetAuthHomePath() string
	SetAuthHomePath(string)
	GetAuthLogoutPath() string
	SetAuthLogoutPath(string)
	GetHealthPath() string
	SetHealthPath(string)
	GetAuthLogger() *slog.Logger
	SetAuthLogger(*slog.Logger)
	GetAuthRoutePermissions() map[string][]string
	SetAuthRoutePermissions(map[string][]string)

	// Initialize method
	InitializeAuthConfig(baseURL string, logger *slog.Logger) error
}

// InitializeAuthConfig initializes the receiver CognitoConfig instance with values from a newly created config.
// It takes a baseURL string used for redirect URI construction and a logger instance for logging operations.
// The function creates a temporary config using InitializeConfig() and copies all fields to the receiver using reflection.
// Returns an error if initialization fails, nil otherwise.
func (c *CognitoConfig) InitializeAuthConfig(baseURL string, logger *slog.Logger) error {
	tempConfig := InitializeConfig(baseURL, logger)
	// fmt.Printf("base url after InitializeAuthConfig: %s\n", tempConfig.)
	if tempConfig == nil {
		return fmt.Errorf("could not initialize auth config")
	}

	// Use reflection to copy all fields
	dst := reflect.ValueOf(c).Elem()
	src := reflect.ValueOf(tempConfig).Elem()

	for i := 0; i < src.NumField(); i++ {
		dstField := dst.Field(i)
		srcField := src.Field(i)

		if dstField.CanSet() {
			dstField.Set(srcField)
		}
	}

	return nil
}

// InitializeConfig creates a new CognitoConfig from environment variables
// and sets the values that are computed from other values and
// sets all default values for fields that will infrequently change.
func InitializeConfig(baseURL string, logger *slog.Logger) *CognitoConfig {

	// Set default logger to text slog 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)
	}

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

	tempConfig := &CognitoConfig{}
	errorParsingCognitoConfig := env.Parse(tempConfig)
	if errorParsingCognitoConfig != nil {
		logger.Error("Error parsing Cognito config", "error", errorParsingCognitoConfig)
		return nil
	}

	// Validate CognitoEnvironmentID if set
	if tempConfig.CognitoEnvironmentID != "" {
		if err := usermanagement.ValidateCognitoEnvironmentID(tempConfig.CognitoEnvironmentID); err != nil {
			logger.Error("Invalid COGNITO_ENVIRONMENT_ID", "error", err)
			return nil
		}
		logger.Info("Cognito environment isolation enabled", "environmentID", tempConfig.CognitoEnvironmentID)
	}

	tempConfig.AuthRedirectURI = baseURL + tempConfig.AuthCallbackPath
	tempConfig.AuthRedirectURI = baseURL + tempConfig.AuthCallbackPath
	tempConfig.AuthTokenURL = fmt.Sprintf("%s/oauth2/token", tempConfig.AuthDomain)
	tempConfig.AuthURL = fmt.Sprintf("%s/oauth2/authorize", tempConfig.AuthDomain)
	tempConfig.AuthLogoutURL = fmt.Sprintf("%s/logout", tempConfig.AuthDomain)
	tempConfig.AuthJwksURL = fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s/.well-known/jwks.json",
		tempConfig.AuthRegion,
		tempConfig.AuthUserPoolID)
	tempConfig.AuthLogger = logger
	tempConfig.AuthRoutePermissions = make(map[string][]string)

	tempConfig.PrettyPrintDebugOnly()
	return tempConfig
}

func (c *CognitoConfig) PrettyPrintDebugOnly() {
	logger := c.AuthLogger
	if logger != nil {
		logger.Debug("CognitoConfig details",
			"AuthClientID", c.AuthClientID,
			"AuthClientSecret", c.AuthClientSecret,
			"AuthRedirectURI", c.AuthRedirectURI,
			"AuthTokenURL", c.AuthTokenURL,
			"AuthURL", c.AuthURL,
			"AuthLogoutURL", c.AuthLogoutURL,
			"AuthJwksURL", c.AuthJwksURL,
			"AuthUserPoolID", c.AuthUserPoolID,
			"AuthRegion", c.AuthRegion,
			"AuthLoginPath", c.AuthLoginPath,
			"AuthCallbackPath", c.AuthCallbackPath,
			"AuthHomePath", c.AuthHomePath,
			"AuthLogoutPath", c.AuthLogoutPath,
			"HealthPath", c.HealthPath,
			"DisableAuth", c.DisableAuth,
			"AuthDomain", c.AuthDomain,
			"PermitIOAPIKey", c.PermitIOAPIKey,
			"PermitIOPDPURL", c.PermitIOPDPURL,
			"CognitoEnvironmentID", c.CognitoEnvironmentID)

		// pretty print the route permissions map
		for route, permissions := range c.AuthRoutePermissions {
			logger.Debug("Route permission",
				"route", route,
				"permissions", permissions)
		}

	}
}

func (c *CognitoConfig) GetAuthRegion() string {
	return c.AuthRegion
}

func (c *CognitoConfig) SetAuthRegion(val string) {
	c.AuthRegion = val
}

func (c *CognitoConfig) GetAuthUserPoolID() string {
	return c.AuthUserPoolID
}

func (c *CognitoConfig) SetAuthUserPoolID(val string) {
	c.AuthUserPoolID = val
}

func (c *CognitoConfig) GetAuthClientSecret() string {
	return c.AuthClientSecret
}

func (c *CognitoConfig) SetAuthClientSecret(val string) {
	c.AuthClientSecret = val
}

func (c *CognitoConfig) GetAuthDomain() string {
	return c.AuthDomain
}

func (c *CognitoConfig) SetAuthDomain(val string) {
	c.AuthDomain = val
}

func (c *CognitoConfig) GetAuthClientID() string {
	return c.AuthClientID
}

func (c *CognitoConfig) SetAuthClientID(val string) {
	c.AuthClientID = val
}

func (c *CognitoConfig) GetAuthRedirectURI() string {
	return c.AuthRedirectURI
}

func (c *CognitoConfig) SetAuthRedirectURI(val string) {
	c.AuthRedirectURI = val
}

func (c *CognitoConfig) GetAuthTokenURL() string {
	return c.AuthTokenURL
}

func (c *CognitoConfig) SetAuthTokenURL(val string) {
	c.AuthTokenURL = val
}

func (c *CognitoConfig) GetAuthJwksURL() string {
	return c.AuthJwksURL
}

func (c *CognitoConfig) SetAuthJwksURL(val string) {
	c.AuthJwksURL = val
}

func (c *CognitoConfig) GetAuthURL() string {
	return c.AuthURL
}

func (c *CognitoConfig) SetAuthURL(val string) {
	c.AuthURL = val
}

func (c *CognitoConfig) GetAuthLogoutURL() string {
	return c.AuthLogoutURL
}

func (c *CognitoConfig) SetAuthLogoutURL(val string) {
	c.AuthLogoutURL = val
}

func (c *CognitoConfig) GetAuthLoginPath() string {
	return c.AuthLoginPath
}

func (c *CognitoConfig) SetAuthLoginPath(val string) {
	c.AuthLoginPath = val
}

func (c *CognitoConfig) GetAuthCallbackPath() string {
	return c.AuthCallbackPath
}

func (c *CognitoConfig) SetAuthCallbackPath(val string) {
	c.AuthCallbackPath = val
}

func (c *CognitoConfig) GetAuthHomePath() string {
	return c.AuthHomePath
}

func (c *CognitoConfig) SetAuthHomePath(val string) {
	c.AuthHomePath = val
}

func (c *CognitoConfig) GetAuthLogoutPath() string {
	return c.AuthLogoutPath
}

func (c *CognitoConfig) SetAuthLogoutPath(val string) {
	c.AuthLogoutPath = val
}

func (c *CognitoConfig) GetHealthPath() string {
	return c.HealthPath
}

func (c *CognitoConfig) SetHealthPath(val string) {
	c.HealthPath = val
}

func (c *CognitoConfig) GetAuthLogger() *slog.Logger {
	return c.AuthLogger
}

func (c *CognitoConfig) SetAuthLogger(val *slog.Logger) {
	c.AuthLogger = val
}

func (c *CognitoConfig) GetAuthRoutePermissions() map[string][]string {
	return c.AuthRoutePermissions
}

func (c *CognitoConfig) SetAuthRoutePermissions(val map[string][]string) {
	c.AuthRoutePermissions = val
}

func (c *CognitoConfig) GetPermitIOAPIKey() string {
	return c.PermitIOAPIKey
}

func (c *CognitoConfig) SetPermitIOAPIKey(val string) {
	c.PermitIOAPIKey = val
}

func (c *CognitoConfig) GetPermitIOPDPURL() string {
	return c.PermitIOPDPURL
}

func (c *CognitoConfig) SetPermitIOPDPURL(val string) {
	c.PermitIOPDPURL = val
}

func (c *CognitoConfig) GetCognitoEnvironmentID() string {
	return c.CognitoEnvironmentID
}

func (c *CognitoConfig) SetCognitoEnvironmentID(val string) {
	c.CognitoEnvironmentID = val
}

func (c *CognitoConfig) GetNoJWTValidation() bool {
	return c.NoJWTValidation
}

func (c *CognitoConfig) SetNoJWTValidation(val bool) {
	c.NoJWTValidation = val
}