query-orchestration/internal/cognitoauth

Cognito Auth

Package Documentation

For comprehensive package-level documentation, see README.md.

Cognito Auth

A reusable Go package for AWS Cognito authentication and Permit.io authorization with Echo framework using PKCE flow.

Features

• Complete OAuth 2.0 PKCE flow for AWS Cognito authentication
• Automatic JWT token refresh for seamless user experience
• JWT token validation and verification
• Dynamic route-based authorization using Permit.io Policy Decision Point (PDP)
• Automatic resource mapping from HTTP routes to Permit.io resources
• HTTP method to action mapping for authorization checks
• Middleware for token handling and permission enforcement
• Cookie-based token storage with refresh token support
• Default home page with authentication status
• Multi-environment user isolation via COGNITO_ENVIRONMENT_ID
• Simple integration with Echo framework

Feature flag for Auth

To disable all of these features for all endpoints you must set DISABLE_AUTH=true in your environment.

Package details

For a detailed explanation of all of the files in this package see the summary file.

Configuration

The package reads configuration from environment variables:

AWS Cognito (Authentication)

• COGNITO_CLIENT_ID: Your AWS Cognito App Client ID
• COGNITO_CLIENT_SECRET: Your AWS Cognito App Client Secret
• COGNITO_USER_POOL_ID: Your AWS Cognito User Pool ID
• COGNITO_DOMAIN: Your AWS Cognito domain
• AWS_REGION: AWS region where your Cognito User Pool is located (default: us-east-2)
• COGNITO_REGION: AWS region for Cognito service (default: us-east-2)
• COGNITO_ENVIRONMENT_ID: Optional. Tags this server instance with an environment identifier for multi-environment user isolation. See Multi-Environment User Isolation.

Permit.io (Authorization)

• PERMIT_IO_API_KEY: Your Permit.io API key for accessing the PDP
• PERMIT_IO_PDP_URL: URL of your Permit.io Policy Decision Point (default: http://localhost:7766)

General

• DEBUG: Set to "true" for debug logging

Environment Variable Loading

Environment variables are loaded in the following hierarchy (first found wins):

1. System environment variables
2. devbox.json "env" section - Development environment defaults
3. .env file - Local overrides and secrets (loaded via env_from in devbox.json)

Feature flags

Disable Authentication

You can disable all of the auth features with the following environment variable: DISABLE_AUTH=true

Test Mode (JWT Validation Bypass)

For testing purposes only, you can bypass JWT signature validation with: NO_JWT_VALIDATION=true

Warning: This flag should ONLY be used in testing environments. When enabled:

• JWT signature validation is bypassed
• Token expiration is ignored
• Unsigned test tokens can be used for authorization testing
• The system logs warnings when this mode is active

This allows testing of authorization logic without requiring valid AWS Cognito tokens.

Middleware Architecture

Middleware Chain Order

The authentication system uses two middleware components that must be registered in the correct order:

1. JWTAuthMiddleware (First)

   • Converts JWT tokens from cookies to Authorization headers
   • Skips processing for public paths (/login, /logout, /home, /swagger/*)
   • Handles logout by clearing auth cookies
   • Does NOT perform authentication - only prepares the request
   • Passes requests to the next middleware in the chain

2. TokenValidationMiddleware (Second)

   • Performs authentication: Validates JWT tokens and extracts user claims
   • Automatic token refresh: Detects expired tokens and refreshes them transparently
   • Performs authorization: Checks permissions via Permit.io PDP
   • Returns JSON error responses for authentication/authorization failures
   • Sets user context for subsequent handlers

Why order matters: The first middleware prepares the request for token validation, while the second middleware performs the actual authentication and authorization enforcement.

Automatic Token Refresh

The system provides seamless token refresh functionality to maintain user sessions without interruption:

How It Works

1. Proactive Refresh: Tokens are automatically refreshed 5 minutes before expiration
2. Transparent Process: Users never experience session interruptions
3. Dual Cookie Storage: Both access tokens (short-lived) and refresh tokens (30-day) are stored as HTTP-only cookies
4. Fallback Handling: When refresh fails, users are gracefully redirected to login

Token Storage Strategy

• Access Token Cookie: auth_token - expires with the JWT token (typically 1 hour)
• Refresh Token Cookie: refresh_token - longer expiration (30 days)
• Security: Both cookies are HTTP-only and secure in production environments

Refresh Process Flow

1. Token Expiration Check: Middleware checks if access token expires within 5 minutes
2. Refresh Token Retrieval: Extracts refresh token from secure cookie
3. Token Exchange: Uses OAuth 2.0 refresh grant to get new tokens from Cognito
4. Cookie Update: Sets new access and refresh tokens in cookies
5. Request Continuation: Processes the original request with fresh tokens

Error Handling

• Missing Refresh Token: Redirects to login with "Session expired" message
• Invalid Refresh Token: Clears cookies and redirects to login
• Cognito Service Issues: Logs errors and falls back to re-authentication
• Network Errors: Retries on next request, maintains user experience

Session Expiration Scenarios

• Access Token Expired + Valid Refresh Token: Automatic refresh (transparent)
• Both Tokens Expired: Redirect to login with clear messaging
• Refresh Token Invalid: Clear cookies and redirect to login
• User Logout: Both cookies are immediately cleared

Client Caching

The system implements efficient client caching for Permit.io connections:

• Singleton pattern: One client instance per unique (apiKey, pdpURL) combination
• Thread-safe: Uses sync.RWMutex for concurrent access protection
• Automatic cleanup: Clients are reused across requests to the same PDP

Startup Validation

The system performs comprehensive validation at startup to ensure proper configuration:

Configuration Validation (ValidatePermitIOConfiguration)

1. Skip validation: If DISABLE_AUTH=true (case-insensitive)
2. PDP connectivity test: HTTP health check to /healthy endpoint
3. API key validation: Quick permission check to verify API key validity
4. Error handling: Application halts if validation fails

Validation Process

• PDP URL: Uses PERMIT_IO_PDP_URL or defaults to http://localhost:7766
• API key: Requires PERMIT_IO_API_KEY environment variable
• Health check: Tests connectivity with 10-second timeout
• Permission test: Validates API key with dummy authorization request

Error Handling

Authentication vs Authorization Errors

Authentication Errors (No valid JWT token):

{
  "error": "Authorization required",
  "message": "This endpoint requires authentication. Please obtain a valid JWT token.",
  "login_url": "/login"
}

• Status Code: 401 Unauthorized
• When: Missing, invalid, or malformed JWT tokens

Session Expiration Errors (Token refresh failed):

{
  "error": "Session expired",
  "message": "Your session has expired. Please log in again.",
  "login_url": "/login"
}

• Status Code: 401 Unauthorized
• When: Access token expired and refresh token is missing or invalid

Authorization Errors (Valid token, insufficient permissions):

{
  "error": "Insufficient permissions",
  "message": "Access denied by authorization service",
  "user": "user-subject-id",
  "resource": "document",
  "action": "get"
}

• Status Code: 403 Forbidden
• When: Valid JWT but Permit.io denies access

Error Response Format

All error responses return JSON with consistent structure:

• error: Brief error identifier
• message: Human-readable description
• Additional fields: Context-specific information (user, resource, action, etc.)

Testing

Test Mode JWT Handling

When NO_JWT_VALIDATION=true is set:

• Pretty-printed JWTs: Handles formatted JSON payloads with newlines/tabs
• 2-part JWT format: Automatically handles header.payload format (adds empty signature)
• Unsigned tokens: Bypasses signature verification for testing
• Flexible parsing: More permissive token parsing for development

Running Integration Tests

# Requires local PDP and API key
task test:permitio

# Or directly with verbose output
go test -tags=permitio -parallel 2 -count=1 -v ./internal/cognitoauth

Test Environment Requirements

• Local PDP: Permit.io Policy Decision Point running on port 7766
• API key: Valid PERMIT_IO_API_KEY in environment or .env file
• Test data: Configured users, resources, and permissions in Permit.io

Authorization System

Permit.io Integration

This package uses Permit.io for dynamic authorization instead of static group-based permissions. The system:

1. Authenticates users via AWS Cognito (JWT tokens)
2. Authorizes requests via Permit.io Policy Decision Point (PDP)
3. Maps HTTP routes to Permit.io resources automatically
4. Converts HTTP methods to Permit.io actions

Automatic Resource Mapping

The system automatically maps HTTP routes to Permit.io resources using the following logic:

• Extracts the first path segment from the route as the resource name
• Removes leading slashes to match Permit.io naming conventions
• Ignores path parameters and nested segments

Examples:

• /document → document
• /document/123 → document
• /document/foo/123 → document
• /api/v1/users → api
• /query-endpoint/run → query-endpoint

HTTP Method to Action Mapping

HTTP methods are mapped to lowercase Permit.io actions:

• GET → get
• POST → post
• PUT → put
• PATCH → patch
• DELETE → delete

Authorization Flow

1. User makes request to protected endpoint
2. Middleware extracts JWT token and validates it with Cognito
3. User subject (sub claim) is extracted from JWT
4. Route is mapped to resource name (first path segment)
5. HTTP method is mapped to action name
6. Permit.io PDP is queried: CheckPermission(user, action, resource)
7. Request is allowed/denied based on PDP response

Multi-Environment User Isolation

Overview

The platform supports hosting multiple customer environments in a shared Cognito user pool. Each server instance is tagged with a COGNITO_ENVIRONMENT_ID to scope its user visibility. Users are tagged with a custom:environment_id attribute in their Cognito profile, and all user operations are filtered to only show users belonging to the current server's environment.

For example, a server configured with COGNITO_ENVIRONMENT_ID=acmehealth will only see users tagged acmehealth and untagged (legacy) users, but will never see users tagged fordmotorco.

Environment Variable

• COGNITO_ENVIRONMENT_ID: Optional. When set, scopes all user operations to this environment. When not set, the system operates in "unscoped" mode where all users are visible (backward compatible).

Naming Rules

Values for COGNITO_ENVIRONMENT_ID must follow these rules:

Rule	Detail
Start character	Must start with a lowercase alpha character (a-z)
End character	Must end with a lowercase alphanumeric character (a-z, 0-9)
Allowed characters	Lowercase alphanumeric (a-z, 0-9) and underscores (_)
Underscore position	Never in first or last position, no consecutive underscores
Case	All lowercase enforced
Length	1 to 50 characters
Empty string	Valid -- means "not set"

Valid examples: acme, client1, acme_health, a_b_c, a

Invalid examples: 1client (starts with digit), Acme (uppercase), _acme (starts with underscore), acme_ (ends with underscore), acme__health (consecutive underscores)

The value is validated at startup. If the format is invalid, the server logs an error and halts.

Access Control Rules

When COGNITO_ENVIRONMENT_ID is set on the server, the TokenValidationMiddleware performs an environment access check after JWT validation and before Permit.io authorization. The rules are:

User Condition	Result
User's custom:environment_id matches the server's value	Allowed
User has no custom:environment_id attribute (legacy/untagged)	Allowed (backward compatible)
User's custom:environment_id does not match the server's value	Blocked with 403 Forbidden
User has the super_admin role in Permit.io	Allowed regardless of environment tag
Server has no COGNITO_ENVIRONMENT_ID set	No check performed (backward compatible)

The super_admin bypass exists because super admins need to operate across environments for administrative tasks.

Caching

The environment access check uses an in-memory cache (keyed by user sub, TTL 30 minutes via EnvCheckCacheTTL) to avoid per-request Cognito and Permit.io API calls. This means changes to a user's custom:environment_id attribute or super_admin role assignment take up to 30 minutes to take effect.

Environment Mismatch Error

When a non-super-admin user with a mismatched environment ID attempts access:

{
  "message": "environment access denied"
}

• Status Code: 403 Forbidden

Data Filtering

All user list and query operations filter results to show only users belonging to the current environment:

• Admin user list endpoints: Results are filtered to include only users whose custom:environment_id matches the server's value or is empty (untagged)
• Admin user CRUD endpoints (get, update, delete, disable, enable): Return 404 Not Found (not 403) for users in other environments to avoid information leakage
• New user creation: Users created via admin API endpoints are automatically tagged with the server's COGNITO_ENVIRONMENT_ID
• EULA compliance reports: User lists are filtered by environment before generating compliance data

Users from other environments are never visible in any API response.

Health Endpoint

When COGNITO_ENVIRONMENT_ID is configured, the /health endpoint response includes an additional field:

{
  "status": "healthy",
  "version": "1.2.3",
  "buildTime": "2025-01-15T10:00:00Z",
  "commit": "abc1234",
  "cognitoEnvironmentId": "acmehealth"
}

The cognitoEnvironmentId field is omitted when the environment variable is not set.

Auto-Registration of Custom Attribute

At startup, the queryAPI automatically ensures the custom:environment_id attribute is registered on the Cognito user pool. This makes new environment deployments self-configuring with no manual AWS CLI commands needed.

The registration process:

1. Calls DescribeUserPool to check if custom:environment_id already exists in the pool's schema
2. If not found, calls AddCustomAttributes to register it (String type, Mutable, MaxLength 50)
3. If already registered, skips silently (debug-level log)
4. If registration fails, the server halts (the attribute is required for correct operation)

IAM permissions required: The IAM role or user running the queryAPI needs cognito-idp:DescribeUserPool and cognito-idp:AddCustomAttributes permissions on the user pool.

Permit.io Attribute Sync

When users are created via the admin API or the user creation tool, the environment_id value is also stored as a custom attribute on the corresponding Permit.io user object. This keeps Permit.io user records consistent with Cognito for audit and policy purposes.