query-orchestration/internal/cognitoauth/envcheck.go

// envcheck.go provides environment access control for authenticated requests.
// When COGNITO_ENVIRONMENT_ID is configured, it verifies that the authenticated user
// belongs to the same environment as this server instance. Super-admins bypass this check.
package cognitoauth

import (
	"context"
	"net/http"
	"os"
	"sync"
	"time"

	"queryorchestration/internal/serviceconfig/auth"
	"queryorchestration/internal/usermanagement"

	"github.com/aws/aws-sdk-go-v2/service/cognitoidentityprovider"
	"github.com/labstack/echo/v4"
)

// EnvCheckCacheTTL is the duration that environment check results are cached per user.
const EnvCheckCacheTTL = 30 * time.Minute

// envCheckCacheEntry stores a cached environment check result for a single user.
type envCheckCacheEntry struct {
	allowed   bool
	expiresAt time.Time
}

// envCheckCache stores cached environment check results to avoid repeated Cognito lookups.
type envCheckCache struct {
	mu      sync.RWMutex
	entries map[string]envCheckCacheEntry
}

// globalEnvCheckCache is the package-level cache shared across all calls.
var globalEnvCheckCache = &envCheckCache{
	entries: make(map[string]envCheckCacheEntry),
}

// get retrieves a cached result for the given user sub. Returns the allowed value and
// true if a valid (non-expired) entry exists, or false if not found or expired.
func (c *envCheckCache) get(sub string) (bool, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	entry, ok := c.entries[sub]
	if !ok {
		return false, false
	}
	if time.Now().After(entry.expiresAt) {
		return false, false
	}
	return entry.allowed, true
}

// set stores a cached result for the given user sub with the configured TTL.
func (c *envCheckCache) set(sub string, allowed bool) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.entries[sub] = envCheckCacheEntry{
		allowed:   allowed,
		expiresAt: time.Now().Add(EnvCheckCacheTTL),
	}
}

// CheckUserEnvironmentAccess verifies that the authenticated user belongs to the server's
// configured environment. If COGNITO_ENVIRONMENT_ID is not set, all users are allowed.
// Results are cached per user sub for EnvCheckCacheTTL to avoid repeated API calls.
//
// The check flow is:
//  1. If no server environment ID is configured, allow all users (return nil)
//  2. Extract user sub from echo context ("user_claims")
//  3. Check cache for existing result
//  4. On cache miss, call GetCognitoUser to retrieve the user's custom:environment_id
//  5. If environment matches or user is untagged (legacy), allow and cache
//  6. If mismatch, check Permit.io for super_admin role (super_admins bypass env checks)
//  7. Cache and return the result
//
// Parameters:
//   - c: Echo context with user_claims set by auth middleware
//   - config: Auth configuration provider with environment ID and Permit.io API key
//   - cognitoClient: AWS Cognito client for user attribute lookup
//
// Returns:
//   - nil if the user is allowed access
//   - echo.NewHTTPError(403) if the user is denied
func CheckUserEnvironmentAccess(
	c echo.Context,
	config auth.ConfigProvider,
	cognitoClient *cognitoidentityprovider.Client,
) error {
	serverEnvID := config.GetCognitoEnvironmentID()
	if serverEnvID == "" {
		return nil
	}

	// Extract user sub from claims
	sub := extractSubFromContext(c)
	if sub == "" {
		return echo.NewHTTPError(http.StatusForbidden, "missing user identity")
	}

	// Check cache
	if allowed, found := globalEnvCheckCache.get(sub); found {
		if allowed {
			return nil
		}
		return echo.NewHTTPError(http.StatusForbidden, "environment access denied")
	}

	// Cache miss -- look up user's environment from Cognito
	user, err := usermanagement.GetCognitoUser(context.Background(), cognitoClient, config.GetAuthUserPoolID(), sub)
	if err != nil {
		// If we can't look up the user, deny access
		return echo.NewHTTPError(http.StatusForbidden, "failed to verify environment access")
	}

	// Check if user belongs to this environment
	if usermanagement.UserBelongsToEnvironment(user, serverEnvID) {
		globalEnvCheckCache.set(sub, true)
		return nil
	}

	// Environment mismatch -- check if user has super_admin role in Permit.io
	if hasSuperAdmin(config, user.SubjectID) {
		globalEnvCheckCache.set(sub, true)
		return nil
	}

	// Denied
	globalEnvCheckCache.set(sub, false)
	return echo.NewHTTPError(http.StatusForbidden, "environment access denied")
}

// extractSubFromContext extracts the user's subject ID from the echo context's user_claims.
func extractSubFromContext(c echo.Context) string {
	claims, ok := c.Get("user_claims").(map[string]interface{})
	if !ok {
		return ""
	}
	sub, ok := claims["sub"].(string)
	if !ok {
		return ""
	}
	return sub
}

// hasSuperAdmin checks if the user has the super_admin role in Permit.io.
// Returns true if the user is a super_admin, false otherwise (including on errors).
func hasSuperAdmin(config auth.ConfigProvider, userKey string) bool {
	permitConfig := &usermanagement.PermitConfig{
		APIKey:  config.GetPermitIOAPIKey(),
		BaseURL: getPermitBaseURL(),
	}

	httpClient := &http.Client{Timeout: 10 * time.Second}
	roles, err := usermanagement.GetPermitUserRoles(httpClient, permitConfig, userKey)
	if err != nil {
		return false
	}

	for _, role := range roles {
		if role == "super_admin" {
			return true
		}
	}
	return false
}

// getPermitBaseURL returns the Permit.io API base URL from environment or default.
func getPermitBaseURL() string {
	if url := os.Getenv("PERMIT_IO_BASE_URL"); url != "" {
		return url
	}
	return "https://api.permit.io"
}

// ResetEnvCheckCache clears the environment check cache. Intended for testing.
func ResetEnvCheckCache() {
	globalEnvCheckCache.mu.Lock()
	defer globalEnvCheckCache.mu.Unlock()
	globalEnvCheckCache.entries = make(map[string]envCheckCacheEntry)
}