/** * Tier enforcement middleware for SentryAgent.ai AgentIdP. * * Enforces per-tenant daily API call limits based on the tenant's tier. * Uses Redis keys `rate:tier:calls:` with TTL aligned to UTC midnight. * * Behaviour: * - Skipped entirely when TIER_ENFORCEMENT env var is 'false'. * - Skipped for enterprise tenants (no limits apply). * - On Redis unavailability: logs the error and proceeds (fail-open). * - Sets X-RateLimit-Limit, X-RateLimit-Remaining, X-RateLimit-Reset on every response. * - Returns HTTP 429 with Retry-After header when limit is exceeded. */ import { Request, Response, NextFunction, RequestHandler } from 'express'; import type { RedisClientType } from 'redis'; import { Pool } from 'pg'; import { TIER_CONFIG, TierName, isTierName } from '../config/tiers.js'; import { TierLimitError } from '../utils/errors.js'; /** Redis key prefix for daily API call counters. */ const CALLS_KEY_PREFIX = 'rate:tier:calls:'; /** * Returns the number of seconds remaining until the next UTC midnight. * Used as the Redis key TTL and the Retry-After value on rejection. * * @returns Seconds until next UTC midnight (minimum 1). */ function secondsUntilUtcMidnight(): number { const now = Date.now(); const midnight = new Date(); midnight.setUTCHours(24, 0, 0, 0); return Math.max(1, Math.ceil((midnight.getTime() - now) / 1000)); } /** * Returns the Unix timestamp (seconds) of the next UTC midnight. * Used for the X-RateLimit-Reset header. * * @returns Unix timestamp of next UTC midnight. */ function nextUtcMidnightTimestamp(): number { const midnight = new Date(); midnight.setUTCHours(24, 0, 0, 0); return Math.ceil(midnight.getTime() / 1000); } /** * Fetches the tenant's current tier from the organizations table. * Falls back to 'free' when the tenant row is not found. * * @param pool - PostgreSQL connection pool. * @param orgId - The organization ID. * @returns The tenant's current TierName. */ async function getTenantTier(pool: Pool, orgId: string): Promise { const result = await pool.query<{ tier: string }>( `SELECT tier FROM organizations WHERE organization_id = $1 LIMIT 1`, [orgId], ); if (result.rows.length === 0) return 'free'; const tier = result.rows[0].tier; return isTierName(tier) ? tier : 'free'; } /** * Creates the tier enforcement middleware. * * Designed to run after auth middleware (req.user must be populated). * Unauthenticated requests pass through unaffected. * * @param pool - PostgreSQL connection pool (used to look up tenant tier). * @param redis - Redis client (used for rate counter storage). * @returns Express RequestHandler. */ export function createTierEnforcementMiddleware( pool: Pool, redis: RedisClientType, ): RequestHandler { return (req: Request, res: Response, next: NextFunction): void => { // Feature flag: bypass all tier enforcement when disabled if (process.env['TIER_ENFORCEMENT'] === 'false') { // Still set headers reflecting unlimited limits res.setHeader('X-RateLimit-Limit', 'unlimited'); res.setHeader('X-RateLimit-Remaining', 'unlimited'); res.setHeader('X-RateLimit-Reset', nextUtcMidnightTimestamp()); next(); return; } // Only enforce for authenticated requests if (!req.user?.organization_id) { next(); return; } const orgId = req.user.organization_id; void (async (): Promise => { try { const tier = await getTenantTier(pool, orgId); // Enterprise tenants bypass all limits if (tier === 'enterprise') { res.setHeader('X-RateLimit-Limit', 'unlimited'); res.setHeader('X-RateLimit-Remaining', 'unlimited'); res.setHeader('X-RateLimit-Reset', nextUtcMidnightTimestamp()); next(); return; } const limit = TIER_CONFIG[tier].maxCallsPerDay; const redisKey = `${CALLS_KEY_PREFIX}${orgId}`; const ttl = secondsUntilUtcMidnight(); const resetAt = nextUtcMidnightTimestamp(); let currentCount: number; try { // Atomically increment and set TTL aligned to UTC midnight. // INCR returns the new value after increment. const newCount = await redis.incr(redisKey); currentCount = newCount; // Set TTL only on the first increment to avoid resetting the window // on every request. If the key was brand new, newCount === 1. if (newCount === 1) { await redis.expire(redisKey, ttl); } } catch (redisErr) { // Redis unavailable — fail-open: log and proceed without rate limiting // eslint-disable-next-line no-console console.error('[tierEnforcement] Redis error — proceeding fail-open:', redisErr); res.setHeader('X-RateLimit-Limit', limit); res.setHeader('X-RateLimit-Remaining', 0); res.setHeader('X-RateLimit-Reset', resetAt); next(); return; } const remaining = Math.max(0, limit - currentCount); // Set rate limit headers on all responses res.setHeader('X-RateLimit-Limit', limit); res.setHeader('X-RateLimit-Remaining', remaining); res.setHeader('X-RateLimit-Reset', resetAt); // Reject if the new count exceeds the limit if (currentCount > limit) { res.setHeader('Retry-After', ttl); next(new TierLimitError('API call', limit, { orgId, tier })); return; } next(); } catch (err) { next(err); } })(); }; }