drinktracker/src/lib/mcp/oauth.ts

import { createHash, randomBytes, timingSafeEqual } from "crypto"
import { NextResponse } from "next/server"
import { prisma } from "@/lib/prisma"
import { generateMcpToken } from "@/lib/mcp/tokens"
import { MCP_SCOPES, type McpScope } from "@/lib/mcp/config"

/**
 * OAuth 2.1 authorization server primitives.
 *
 * This app is both the authorization server and the resource server, and it
 * serves public clients only - no client secrets are stored anywhere, so there
 * is nothing here to leak. Clients authenticate at the token endpoint with PKCE,
 * which Claude always sends with S256.
 */

/** Access tokens are short; the refresh token is what carries the session. */
export const ACCESS_TOKEN_TTL_MS = 60 * 60 * 1000
export const REFRESH_TOKEN_TTL_MS = 90 * 24 * 60 * 60 * 1000
/** Long enough for a redirect and a token call, short enough to be useless if leaked. */
export const AUTH_CODE_TTL_MS = 60 * 1000

export const OFFLINE_ACCESS = "offline_access"

/** Advertised in discovery metadata and offered on the consent screen. */
export const SUPPORTED_SCOPES: string[] = [...MCP_SCOPES, OFFLINE_ACCESS]

export const SCOPE_DESCRIPTIONS: Record<string, string> = {
  "drinks:read": "See your drinks, ratings, wishlist and taste preferences",
  "drinks:write": "Add, edit and delete drinks, ratings and wishlist entries",
  "bar:read": "See your bar inventory and saved recipes",
  "bar:write": "Add, edit and delete bar items and recipes",
  [OFFLINE_ACCESS]: "Stay connected without asking you to sign in again",
}

// ─── Opaque secrets ──────────────────────────────────────────────

export function sha256(value: string): string {
  return createHash("sha256").update(value).digest("hex")
}

export function generateAuthCode(): { code: string; hash: string } {
  const code = randomBytes(32).toString("base64url")
  return { code, hash: sha256(code) }
}

export function generateRefreshToken(): { token: string; hash: string } {
  const token = randomBytes(32).toString("base64url")
  return { token, hash: sha256(token) }
}

export function generateClientId(): string {
  return `dtc_${randomBytes(16).toString("hex")}`
}

// ─── Scopes ──────────────────────────────────────────────────────

/**
 * Narrow a requested scope string to what this server actually grants. Unknown
 * scopes are dropped rather than rejected: an unrecognised scope is not an error
 * per RFC 6749, and rejecting would break clients that request extras.
 */
export function parseScopes(requested: string | null | undefined): string[] {
  const asked = (requested ?? "").split(/\s+/).filter(Boolean)
  const granted = asked.filter((s) => SUPPORTED_SCOPES.includes(s))
  // Default to read-only when nothing usable was asked for.
  if (granted.filter((s) => s !== OFFLINE_ACCESS).length === 0) {
    return ["drinks:read", "bar:read", ...granted.filter((s) => s === OFFLINE_ACCESS)]
  }
  return granted
}

/** The scopes actually stored on an access token: everything except the marker. */
export function accessScopes(scopes: string[]): McpScope[] {
  return scopes.filter((s): s is McpScope =>
    (MCP_SCOPES as readonly string[]).includes(s)
  )
}

export function coversScopes(granted: string[], requested: string[]): boolean {
  return requested.every((s) => granted.includes(s))
}

// ─── PKCE ────────────────────────────────────────────────────────

/** RFC 7636 unreserved set, 43-128 characters. */
const VERIFIER_RE = /^[A-Za-z0-9\-._~]{43,128}$/

export function verifyPkce(verifier: string, challenge: string): boolean {
  if (!VERIFIER_RE.test(verifier)) return false
  const computed = createHash("sha256").update(verifier).digest("base64url")
  const a = Buffer.from(computed)
  const b = Buffer.from(challenge)
  // Length check first: timingSafeEqual throws on a mismatch rather than
  // returning false.
  if (a.length !== b.length) return false
  return timingSafeEqual(a, b)
}

// ─── Redirect URIs ───────────────────────────────────────────────

function isLoopback(url: URL): boolean {
  return (
    url.protocol === "http:" &&
    (url.hostname === "localhost" ||
      url.hostname === "127.0.0.1" ||
      url.hostname === "[::1]" ||
      url.hostname === "::1")
  )
}

/**
 * A redirect URI is acceptable at registration if it is https, or an RFC 8252
 * loopback address. Fragments are forbidden by RFC 6749.
 */
export function isRegisterableRedirectUri(value: string): boolean {
  let url: URL
  try {
    url = new URL(value)
  } catch {
    return false
  }
  if (url.hash) return false
  return url.protocol === "https:" || isLoopback(url)
}

/**
 * Matching a redirect at authorize time.
 *
 * Exact string equality, except for loopback addresses where the PORT ALONE is
 * ignored. Claude Code registers http://localhost/callback and then redirects to
 * an ephemeral port such as http://localhost:51234/callback, so exact matching
 * would reject every native client. RFC 8252 section 7.3 requires the port to be
 * ignored for the IP-literal form; the same is applied to localhost so Claude
 * Code works at all.
 *
 * Everything else must match exactly. An earlier version compared only scheme,
 * host and path, which silently accepted a request that differed in query,
 * fragment or userinfo - and a callback that forwards its query is then a route
 * for delivering the authorization code somewhere the client never registered.
 * RFC 9700 section 2.1 requires exact matching apart from the loopback port.
 */
export function redirectUriMatches(registered: string, requested: string): boolean {
  if (registered === requested) return true

  let a: URL, b: URL
  try {
    a = new URL(registered)
    b = new URL(requested)
  } catch {
    return false
  }
  if (!isLoopback(a) || !isLoopback(b)) return false

  // A fragment is never legal on a redirect URI (RFC 6749 3.1.2).
  if (a.hash || b.hash) return false
  // Credentials in a redirect target are a spoofing aid and never needed here.
  if (a.username || a.password || b.username || b.password) return false

  // Normalise away only the port, then require the rest to be identical.
  const strip = (u: URL) => {
    const c = new URL(u.toString())
    c.port = ""
    return c.toString()
  }
  return strip(a) === strip(b)
}

export function findMatchingRedirectUri(
  registered: string[],
  requested: string
): string | null {
  return registered.find((uri) => redirectUriMatches(uri, requested)) ?? null
}

// ─── Errors ──────────────────────────────────────────────────────

/**
 * RFC 6749 error body. `invalid_grant` in particular is load-bearing: Claude
 * keys off exactly that code to decide a refresh token is dead and it should
 * restart the whole flow. Any other code leaves the connection stuck.
 */
export function oauthError(
  error: string,
  description?: string,
  status = 400
): NextResponse {
  return NextResponse.json(
    { error, ...(description ? { error_description: description } : {}) },
    { status, headers: { "Cache-Control": "no-store", Pragma: "no-cache" } }
  )
}

// ─── Token issuance ──────────────────────────────────────────────

export interface IssuedTokens {
  access_token: string
  token_type: "Bearer"
  expires_in: number
  scope: string
  refresh_token?: string
}

/** Thrown when a grant was revoked or re-consented while a request was in flight. */
export class GrantStateError extends Error {
  constructor(message = "This authorization is no longer valid.") {
    super(message)
    this.name = "GrantStateError"
  }
}

/**
 * Mints an access token against a grant.
 *
 * Two things this deliberately does NOT do:
 *
 *  - It does not write scopes back to the grant. Token issuance is not consent,
 *    and treating it as such let a stale authorization code redefine what the
 *    user had most recently agreed to. The consent endpoint is the only writer.
 *  - It does not touch refresh tokens. Those have their own lineage.
 *
 * The whole write is conditional on the grant still being active AT THE EPOCH
 * the caller validated. A revoke racing this call makes the guarded update
 * affect zero rows and the transaction aborts, rather than leaving a live token
 * hanging off a revoked grant to be resurrected by a later re-consent.
 */
export async function issueAccessToken(
  grant: { id: string; epoch: number },
  userId: string,
  scopes: string[],
  clientName: string
): Promise<IssuedTokens> {
  const access = generateMcpToken()

  await prisma.$transaction(async (tx) => {
    const { count } = await tx.oAuthGrant.updateMany({
      where: { id: grant.id, epoch: grant.epoch, revokedAt: null },
      data: { lastUsedAt: new Date() },
    })
    if (count !== 1) throw new GrantStateError()

    await tx.mcpAccessToken.updateMany({
      where: { grantId: grant.id, revokedAt: null },
      data: { revokedAt: new Date() },
    })
    await tx.mcpAccessToken.create({
      data: {
        userId,
        tokenHash: access.hash,
        prefix: access.prefix,
        name: clientName,
        scopes: accessScopes(scopes),
        source: "oauth",
        grantId: grant.id,
        expiresAt: new Date(Date.now() + ACCESS_TOKEN_TTL_MS),
      },
    })
  })

  return {
    access_token: access.raw,
    token_type: "Bearer",
    expires_in: Math.floor(ACCESS_TOKEN_TTL_MS / 1000),
    scope: scopes.join(" "),
  }
}

/**
 * Mints the next refresh token in a family, guarded on the grant's epoch for the
 * same reason as above.
 */
export async function issueRefreshToken(
  grant: { id: string; epoch: number },
  generation: number
): Promise<string> {
  const { token, hash } = generateRefreshToken()

  await prisma.$transaction(async (tx) => {
    const active = await tx.oAuthGrant.count({
      where: { id: grant.id, epoch: grant.epoch, revokedAt: null },
    })
    if (active !== 1) throw new GrantStateError()

    await tx.oAuthRefreshToken.create({
      data: {
        grantId: grant.id,
        tokenHash: hash,
        generation,
        expiresAt: new Date(Date.now() + REFRESH_TOKEN_TTL_MS),
      },
    })
  })

  return token
}

export type RefreshConsumption =
  | { kind: "unknown" }
  | { kind: "expired" }
  | { kind: "reuse"; grantId: string }
  | { kind: "ok"; grantId: string; generation: number }

/**
 * Redeems a refresh token exactly once.
 *
 * Every generation is retained, so a token from any point in the family's life
 * is still recognisable as reuse - which is the whole point. The guarded update
 * on `usedAt` means two concurrent presentations of the same token cannot both
 * win: the loser is reported as reuse, which is indistinguishable from theft and
 * treated the same way.
 */
export async function consumeRefreshToken(hash: string): Promise<RefreshConsumption> {
  const row = await prisma.oAuthRefreshToken.findUnique({
    where: { tokenHash: hash },
    select: { id: true, grantId: true, generation: true, usedAt: true, expiresAt: true },
  })

  if (!row) return { kind: "unknown" }
  if (row.usedAt) return { kind: "reuse", grantId: row.grantId }
  if (row.expiresAt <= new Date()) return { kind: "expired" }

  const { count } = await prisma.oAuthRefreshToken.updateMany({
    where: { id: row.id, usedAt: null },
    data: { usedAt: new Date() },
  })
  if (count !== 1) return { kind: "reuse", grantId: row.grantId }

  return { kind: "ok", grantId: row.grantId, generation: row.generation }
}

/** Kills a grant, every access token it issued, and its whole refresh family. */
export async function revokeGrant(grantId: string): Promise<void> {
  const now = new Date()
  await prisma.$transaction([
    prisma.mcpAccessToken.updateMany({
      where: { grantId, revokedAt: null },
      data: { revokedAt: now },
    }),
    prisma.oAuthRefreshToken.deleteMany({ where: { grantId } }),
    // Outstanding codes die too, so one minted before the revoke cannot be
    // redeemed against a later re-consent.
    prisma.oAuthAuthCode.deleteMany({ where: { grantId, consumedAt: null } }),
    prisma.oAuthGrant.update({
      where: { id: grantId },
      data: { revokedAt: now },
    }),
  ])
}