1  import { NextResponse } from 'next/server'
  2  import type { NextRequest } from 'next/server'
  3  import { AUTH_COOKIE_NAME } from '@/lib/auth'
  4  import {
  5    buildExtensionInstallCorsHeaders,
  6    isExtensionInstallCorsPath,
  7    resolveExtensionInstallCorsOrigin,
  8  } from '@/lib/extension-install-cors'
  9  import { isProductionRuntime } from '@/lib/runtime/runtime-env'
 10  import { hmrSingleton } from '@/lib/shared-utils'
 11  
 12  /* ------------------------------------------------------------------ */
 13  /*  Rate-limit state — HMR-safe via globalThis                        */
 14  /* ------------------------------------------------------------------ */
 15  
 16  interface RateLimitEntry {
 17    count: number
 18    lockedUntil: number
 19  }
 20  
 21  const rateLimitMap = hmrSingleton('__swarmclaw_rate_limit__', () => new Map<string, RateLimitEntry>())
 22  
 23  const MAX_ATTEMPTS = 5
 24  const LOCKOUT_MS = 15 * 60 * 1000 // 15 minutes
 25  const PRUNE_THRESHOLD = 1000
 26  
 27  function isRateLimitEnabled(): boolean {
 28    return isProductionRuntime()
 29  }
 30  
 31  /** Prune expired entries when the map grows too large. */
 32  function pruneRateLimitMap() {
 33    if (rateLimitMap.size <= PRUNE_THRESHOLD) return
 34    const now = Date.now()
 35    rateLimitMap.forEach((entry, ip) => {
 36      if (entry.lockedUntil < now && entry.count < MAX_ATTEMPTS) {
 37        rateLimitMap.delete(ip)
 38      }
 39    })
 40  }
 41  
 42  /** Extract client IP from the request. */
 43  function getClientIp(request: NextRequest): string {
 44    const forwarded = request.headers.get('x-forwarded-for')
 45    if (forwarded) {
 46      const first = forwarded.split(',')[0]?.trim()
 47      if (first) return first
 48    }
 49    return (request as unknown as { ip?: string }).ip ?? 'unknown'
 50  }
 51  
 52  function withExtensionInstallCorsHeaders(pathname: string, origin: string | null, headers?: HeadersInit): Headers {
 53    const merged = new Headers(headers)
 54    if (!isExtensionInstallCorsPath(pathname)) return merged
 55    const corsHeaders = buildExtensionInstallCorsHeaders(origin)
 56    new Headers(corsHeaders).forEach((value, key) => {
 57      merged.set(key, value)
 58    })
 59    return merged
 60  }
 61  
 62  /* ------------------------------------------------------------------ */
 63  /*  Proxy                                                              */
 64  /* ------------------------------------------------------------------ */
 65  
 66  /** Access key auth proxy with brute-force rate limiting.
 67   *  Checks X-Access-Key header or auth cookie on all /api/ routes except /api/auth.
 68   *  The key is validated against the ACCESS_KEY env var.
 69   *  After 5 failed attempts from a single IP the client is locked out for 15 minutes.
 70   */
 71  export function proxy(request: NextRequest) {
 72    const rateLimitEnabled = isRateLimitEnabled()
 73    const { pathname } = request.nextUrl
 74    const corsOrigin = resolveExtensionInstallCorsOrigin(request.headers.get('origin'))
 75    const isWebhookTrigger = request.method === 'POST'
 76      && /^\/api\/webhooks\/[^/]+\/?$/.test(pathname)
 77    const isConnectorWebhook = request.method === 'POST'
 78      && /^\/api\/connectors\/[^/]+\/webhook\/?$/.test(pathname)
 79  
 80    if (request.method === 'OPTIONS' && isExtensionInstallCorsPath(pathname)) {
 81      if (!corsOrigin) {
 82        return NextResponse.json({ error: 'Origin not allowed' }, { status: 403 })
 83      }
 84      return new NextResponse(null, {
 85        status: 204,
 86        headers: buildExtensionInstallCorsHeaders(corsOrigin),
 87      })
 88    }
 89  
 90    // A2A endpoints use their own authentication (Authorization: Bearer / x-a2a-access-key)
 91    const isA2ARoute = pathname === '/api/a2a'
 92      || pathname.startsWith('/api/a2a/')
 93      || pathname === '/api/.well-known/agent-card'
 94  
 95    // Only protect API routes (not auth, inbound webhooks, or A2A)
 96    if (
 97      !pathname.startsWith('/api/')
 98      || pathname === '/api/auth'
 99      || pathname === '/api/healthz'
100      || isWebhookTrigger
101      || isConnectorWebhook
102      || isA2ARoute
103    ) {
104      return NextResponse.next()
105    }
106  
107    const accessKey = process.env.ACCESS_KEY
108    if (!accessKey) {
109      // No key configured — allow all (dev mode)
110      return NextResponse.next()
111    }
112  
113    // --- Rate-limit housekeeping ---
114    if (rateLimitEnabled) pruneRateLimitMap()
115  
116    const clientIp = getClientIp(request)
117    const entry = rateLimitEnabled ? rateLimitMap.get(clientIp) : undefined
118  
119    // Check lockout before even validating the key
120    if (rateLimitEnabled && entry && entry.lockedUntil > Date.now()) {
121      const retryAfter = Math.ceil((entry.lockedUntil - Date.now()) / 1000)
122      return NextResponse.json(
123        { error: 'Too many failed attempts. Try again later.', retryAfter },
124        {
125          status: 429,
126          headers: withExtensionInstallCorsHeaders(pathname, corsOrigin, { 'Retry-After': String(retryAfter) }),
127        },
128      )
129    }
130  
131    const cookieKey = request.cookies.get(AUTH_COOKIE_NAME)?.value?.trim() || ''
132    const headerKey = request.headers.get('x-access-key')?.trim() || ''
133    const providedKey = cookieKey || headerKey
134  
135    if (providedKey !== accessKey) {
136      let remaining = MAX_ATTEMPTS
137      if (rateLimitEnabled) {
138        const current = rateLimitMap.get(clientIp) ?? { count: 0, lockedUntil: 0 }
139        current.count += 1
140  
141        if (current.count >= MAX_ATTEMPTS) {
142          current.lockedUntil = Date.now() + LOCKOUT_MS
143        }
144  
145        rateLimitMap.set(clientIp, current)
146        remaining = Math.max(0, MAX_ATTEMPTS - current.count)
147      }
148      return NextResponse.json(
149        { error: 'Unauthorized' },
150        {
151          status: 401,
152          headers: withExtensionInstallCorsHeaders(pathname, corsOrigin, { 'X-RateLimit-Remaining': String(remaining) }),
153        },
154      )
155    }
156  
157    // Successful auth — clear any prior failed-attempt tracking for this IP
158    if (rateLimitEnabled && entry) {
159      rateLimitMap.delete(clientIp)
160    }
161  
162    return NextResponse.next()
163  }
164  
165  export const config = {
166    matcher: '/api/:path*',
167  }