AOS/src/middleware.ts

import { NextRequest, NextResponse } from 'next/server'

// Middleware: protects all routes except public ones.
//
// Session = HMAC-signed token (base64url(payload).hmac) issued by the login
// flow. The signing key is AUTH_SECRET (required). Verified here with a
// constant-time HMAC comparison. The payload carries { sub, email, role,
// iat, exp } with a 7-day TTL.
//
// Auth flow (email + password, optional 2FA):
//   /login           — email + password step (public)
//   /login/verify    — 2FA OTP step (public, uses signed challenge)
//   /register        — create account (public)
//   /verify-email    — confirm email OTP (public)
//   /resend-otp      — re-send login OTP (public, challenge-gated)
//   /2fa/*           — manage 2FA (session-protected)
//   /me              — current user (session-protected)
//
// Also applies global rate limiting to API routes.

// --- Simple in-memory rate limiter (per-IP, per-minute) ---
const rateLimitStore = new Map<string, { count: number; resetAt: number }>()
const RATE_LIMIT_MAX = 120 // 120 requests per minute per IP
const RATE_LIMIT_WINDOW = 60_000 // 1 minute

function checkRateLimit(ip: string): boolean {
  const now = Date.now()
  const entry = rateLimitStore.get(ip)
  if (!entry || entry.resetAt < now) {
    rateLimitStore.set(ip, { count: 1, resetAt: now + RATE_LIMIT_WINDOW })
    return true
  }
  if (entry.count >= RATE_LIMIT_MAX) return false
  entry.count++
  return true
}

function getClientIp(req: NextRequest): string {
  return req.headers.get('x-forwarded-for')?.split(',')[0]?.trim()
    ?? req.headers.get('x-real-ip')
    ?? 'unknown'
}

const PUBLIC_PATHS = [
  '/api/health',
  '/api/aos/seed-prod',
  '/api/aos/licenses/check',
  '/api/aos/licenses/validate',
  '/api/aos/stripe/webhook',
  '/api/auth/login',
  '/api/auth/register',
  '/api/auth/verify-email',
  '/api/auth/resend-otp',
  '/login',
  '/register',
  '/_next',
  '/favicon.ico',
  '/manifest.json',
  '/icons/',
  '/apple-touch-icon',
]

const PUBLIC_API_PREFIXES = [
  '/api/health',
  '/api/aos/seed-prod',
  '/api/aos/licenses/check',
  '/api/aos/licenses/validate',
  '/api/aos/stripe/webhook',
  '/api/auth/login',
  '/api/auth/login/verify',
  '/api/auth/register',
  '/api/auth/verify-email',
  '/api/auth/resend-otp',
  '/api/v1/',  // public REST API (token-auth'd separately)
  '/api/mcp',  // MCP (token-auth'd)
]

function isPublic(pathname: string): boolean {
  // Exact matches
  if (PUBLIC_PATHS.includes(pathname)) return true
  // Prefix matches
  for (const prefix of PUBLIC_API_PREFIXES) {
    if (pathname.startsWith(prefix)) return true
  }
  // Static assets
  if (pathname.startsWith('/_next/') || pathname.startsWith('/icons/')) return true
  return false
}

// --- Session token verification (Web Crypto — Edge-runtime compatible) ---
//
// Token format: base64url(payload).base64url(hmac_sha256(payload))
// The HMAC key is AUTH_SECRET. We verify with SubtleCrypto.verify, which
// is constant-time in practice and runs in the Edge runtime (unlike Node's
// `crypto` module, which produced the earlier "not supported in Edge
// Runtime" warning). Payload carries { sub, email, role, iat, exp }.
function base64urlToBytes(s: string): Uint8Array {
  const pad = s.length % 4 === 0 ? '' : '='.repeat(4 - (s.length % 4))
  const b64 = (s + pad).replace(/-/g, '+').replace(/_/g, '/')
  const bin = atob(b64)
  const bytes = new Uint8Array(bin.length)
  for (let i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i)
  return bytes
}

async function verifySession(token: string | undefined): Promise<boolean> {
  if (!token) return false
  const secret = process.env.AUTH_SECRET
  if (!secret) {
    // In dev without AUTH_SECRET, allow the dev-ephemeral fallback tokens.
    // The route-handler verifySessionToken() enforces the same key.
    if (process.env.NODE_ENV !== 'production') return parseAndCheckExpiry(token)
    return false
  }
  try {
    const [payloadB64, sigB64] = token.split('.')
    if (!payloadB64 || !sigB64) return false
    // Expiry check first (cheap, avoids crypto work for expired tokens).
    const payload = JSON.parse(new TextDecoder().decode(base64urlToBytes(payloadB64)))
    if (payload.exp && Date.now() > payload.exp) return false
    const enc = new TextEncoder()
    const key = await crypto.subtle.importKey(
      'raw',
      enc.encode(secret),
      { name: 'HMAC', hash: 'SHA-256' },
      false,
      ['verify']
    )
    return await crypto.subtle.verify(
      'HMAC',
      key,
      base64urlToBytes(sigB64),
      enc.encode(payloadB64)
    )
  } catch {
    return false
  }
}

// Dev-only fallback: check expiry without HMAC (the login route still does
// full HMAC verification before issuing the token).
function parseAndCheckExpiry(token: string): boolean {
  try {
    const [payloadB64] = token.split('.')
    if (!payloadB64) return false
    const payload = JSON.parse(new TextDecoder().decode(base64urlToBytes(payloadB64)))
    if (payload.exp && Date.now() > payload.exp) return false
    return true
  } catch {
    return false
  }
}

export async function middleware(req: NextRequest) {
  const { pathname } = req.nextUrl

  // --- Global rate limiting for API routes ---
  if (pathname.startsWith('/api/')) {
    const ip = getClientIp(req)
    if (!checkRateLimit(ip)) {
      return NextResponse.json(
        { error: 'rate_limit_exceeded', message: 'Too many requests. Try again in a minute.' },
        { status: 429, headers: { 'retry-after': '60' } }
      )
    }
  }

  // Allow public paths
  if (isPublic(pathname)) {
    return NextResponse.next()
  }

  // Verify the session cookie with full HMAC (Web Crypto, Edge-compatible).
  const sessionToken = req.cookies.get('aos-session')?.value
  const valid = await verifySession(sessionToken)

  if (!valid) {
    if (pathname.startsWith('/api/')) {
      return NextResponse.json(
        { error: 'unauthorized', message: 'Authentication required' },
        { status: 401 }
      )
    }
    // For pages, redirect to login with the intended destination.
    const loginUrl = new URL('/login', req.url)
    loginUrl.searchParams.set('redirect', pathname)
    return NextResponse.redirect(loginUrl)
  }

  return NextResponse.next()
}

export const config = {
  matcher: [
    /*
     * Match all request paths except:
     * - _next/static, _next/image (static files)
     * - favicon.ico, icons, manifest (public assets)
     */
    '/((?!_next/static|_next/image|favicon.ico|icons|manifest.json|apple-touch-icon).*)',
  ],
}