query-orchestration/internal/cognitoauth/readme.md

# Cognito Auth

## Package Documentation

For comprehensive package-level documentation, see [README.md](./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](./summary.md) 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](#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):
```json
{
  "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):
```json
{
  "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):
```json
{
  "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
```bash
# 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:

```json
{
  "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:

```json
{
  "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.