remora/web/server/oauth.mjs

import { createHash, createPublicKey, createVerify } from 'node:crypto'

import {
  equalSecret,
  expiredSessionCookie,
  randomValue,
  SESSION_RENEWAL_PENDING,
  sessionCookie,
  sessionIdentifier,
} from './sessions.mjs'

const CLOCK_SKEW_SECONDS = 30
const OIDC_TIMEOUT_MILLISECONDS = 15_000

export class SessionRenewalPendingError extends Error {}

export class OAuthClient {
  constructor(config, store, fetchImplementation = fetch) {
    this.config = config
    this.store = store
    this.fetch = fetchImplementation
    this.discovery = null
  }

  async route(request, response, url) {
    if (request.method === 'GET' && url.pathname === '/auth/login')
      return this.login(request, response)
    if (request.method === 'GET' && url.pathname === '/auth/callback')
      return this.callback(request, response, url)
    if (request.method === 'GET' && url.pathname === '/auth/session')
      return this.session(request, response)
    if (request.method === 'POST' && url.pathname === '/auth/logout')
      return this.logout(request, response)
    return false
  }

  async login(request, response) {
    const previous = sessionIdentifier(request)
    if (previous !== null) await this.store.destroy(previous)
    const discovery = await this.metadata()
    const verifier = randomValue(48)
    const state = randomValue()
    const nonce = randomValue()
    const identifier = await this.store.create({
      kind: 'authorization',
      state,
      nonce,
      verifier,
    })
    const authorization = new URL(discovery.authorization_endpoint)
    authorization.search = new URLSearchParams({
      client_id: this.config.oidcClientId,
      redirect_uri: `${this.config.publicOrigin}/auth/callback`,
      response_type: 'code',
      scope: this.config.oidcScope,
      state,
      nonce,
      code_challenge: createHash('sha256').update(verifier).digest('base64url'),
      code_challenge_method: 'S256',
    }).toString()
    response.writeHead(302, {
      Location: authorization.toString(),
      'Set-Cookie': sessionCookie(identifier, this.config),
      'Cache-Control': 'no-store',
    })
    response.end()
    return true
  }

  async callback(request, response, url) {
    const identifier = sessionIdentifier(request)
    const pending = identifier && (await this.store.read(identifier))
    const state = url.searchParams.get('state')
    const code = url.searchParams.get('code')
    if (
      pending?.kind !== 'authorization' ||
      code === null ||
      state === null ||
      !equalSecret(state, pending.state)
    )
      return this.authenticationFailure(response, identifier)

    try {
      const discovery = await this.metadata()
      const tokenResponse = await this.fetch(discovery.token_endpoint, {
        method: 'POST',
        headers: {
          Authorization: clientAuthorization(this.config),
          'Content-Type': 'application/x-www-form-urlencoded',
        },
        body: new URLSearchParams({
          grant_type: 'authorization_code',
          code,
          redirect_uri: `${this.config.publicOrigin}/auth/callback`,
          code_verifier: pending.verifier,
        }),
        signal: AbortSignal.timeout(OIDC_TIMEOUT_MILLISECONDS),
      })
      if (!tokenResponse.ok) throw new Error('token exchange failed')
      const tokens = await tokenResponse.json()
      const identity = await this.validateIdToken(
        tokens.id_token,
        pending.nonce,
      )
      const csrfToken = randomValue()
      const session = authenticatedSession(tokens, identity, csrfToken)
      const replacement = await this.store.replace(identifier, session)
      response.writeHead(302, {
        Location: this.config.publicOrigin,
        'Set-Cookie': sessionCookie(replacement, this.config),
        'Cache-Control': 'no-store',
      })
      response.end()
    } catch {
      await this.authenticationFailure(response, identifier)
    }
    return true
  }

  async session(request, response) {
    const loaded = await this.authenticated(request)
    if (loaded === null) {
      json(
        response,
        200,
        { kind: 'signed_out' },
        {
          'Set-Cookie': expiredSessionCookie(this.config),
        },
      )
      return true
    }
    json(response, 200, {
      kind: 'authenticated',
      identity: loaded.session.identity,
      csrfToken: loaded.session.csrfToken,
    })
    return true
  }

  async logout(request, response) {
    const loaded = await this.storedSession(request)
    if (
      loaded === null ||
      !validUnsafeRequest(request, this.config, loaded.session.csrfToken)
    ) {
      json(response, 403, { error: 'invalid_csrf' })
      return true
    }
    await this.store.destroy(loaded.identifier)
    const discovery = await this.metadata().catch(() => null)
    let location = this.config.publicOrigin
    if (discovery?.end_session_endpoint) {
      const logout = new URL(discovery.end_session_endpoint)
      logout.searchParams.set(
        'post_logout_redirect_uri',
        this.config.publicOrigin,
      )
      logout.searchParams.set('id_token_hint', loaded.session.idToken)
      location = logout.toString()
    }
    json(
      response,
      200,
      { redirect: location },
      { 'Set-Cookie': expiredSessionCookie(this.config) },
    )
    return true
  }

  async authenticated(request, refresh = true) {
    const loaded = await this.storedSession(request)
    if (loaded === null) return null
    const { identifier } = loaded
    let { session } = loaded
    if (session.accessExpiresAt <= now() + CLOCK_SKEW_SECONDS) {
      if (!refresh) {
        if (session.accessExpiresAt <= now()) {
          session = await this.store.renew(
            identifier,
            (current) =>
              validAuthenticatedSession(current) &&
              current.accessExpiresAt <= now(),
            async () => null,
          )
          if (session === SESSION_RENEWAL_PENDING)
            throw new SessionRenewalPendingError()
          if (!validAuthenticatedSession(session)) return null
        }
      } else {
        const renewed = await this.store.renew(
          identifier,
          (current) =>
            validAuthenticatedSession(current) &&
            current.accessExpiresAt <= now() + CLOCK_SKEW_SECONDS,
          (current) => this.refresh(current).catch(() => null),
        )
        if (renewed === SESSION_RENEWAL_PENDING) {
          if (session.accessExpiresAt <= now())
            throw new SessionRenewalPendingError()
        } else {
          session = renewed
        }
        if (!validAuthenticatedSession(session)) {
          if (session?.kind === 'authenticated')
            await this.store.destroy(identifier)
          return null
        }
      }
    }
    return { identifier, session }
  }

  async storedSession(request) {
    const identifier = sessionIdentifier(request)
    if (identifier === null) return null
    const session = await this.store.read(identifier)
    if (!validAuthenticatedSession(session)) {
      if (session?.kind === 'authenticated')
        await this.store.destroy(identifier)
      return null
    }
    return { identifier, session }
  }

  async refresh(session) {
    if (typeof session.refreshToken !== 'string') return null
    const discovery = await this.metadata()
    const response = await this.fetch(discovery.token_endpoint, {
      method: 'POST',
      headers: {
        Authorization: clientAuthorization(this.config),
        'Content-Type': 'application/x-www-form-urlencoded',
      },
      body: new URLSearchParams({
        grant_type: 'refresh_token',
        refresh_token: session.refreshToken,
      }),
      signal: AbortSignal.timeout(OIDC_TIMEOUT_MILLISECONDS),
    })
    if (!response.ok) return null
    const tokens = await response.json()
    const idToken =
      typeof tokens.id_token === 'string' ? tokens.id_token : session.idToken
    const identity =
      typeof tokens.id_token === 'string'
        ? await this.validateIdToken(tokens.id_token, undefined)
        : session.identity
    if (identity !== session.identity) return null
    return authenticatedSession(
      {
        ...tokens,
        id_token: idToken,
        refresh_token:
          typeof tokens.refresh_token === 'string'
            ? tokens.refresh_token
            : session.refreshToken,
      },
      identity,
      session.csrfToken,
    )
  }

  async metadata() {
    if (this.discovery !== null) return this.discovery
    const response = await this.fetch(
      `${this.config.oidcAuthority}/.well-known/openid-configuration`,
      { signal: AbortSignal.timeout(OIDC_TIMEOUT_MILLISECONDS) },
    )
    if (!response.ok) throw new Error('OIDC discovery failed')
    const metadata = await response.json()
    for (const field of [
      'issuer',
      'authorization_endpoint',
      'token_endpoint',
      'jwks_uri',
    ])
      if (typeof metadata[field] !== 'string')
        throw new Error(`OIDC discovery omitted ${field}`)
    if (metadata.issuer !== this.config.oidcAuthority)
      throw new Error('OIDC discovery issuer mismatch')
    this.discovery = metadata
    return metadata
  }

  async validateIdToken(token, nonce) {
    const [headerPart, payloadPart, signaturePart] = String(token).split('.')
    if (!headerPart || !payloadPart || !signaturePart)
      throw new Error('invalid ID token')
    const header = JSON.parse(Buffer.from(headerPart, 'base64url').toString())
    const claims = JSON.parse(Buffer.from(payloadPart, 'base64url').toString())
    if (header.alg !== 'RS256' || typeof header.kid !== 'string')
      throw new Error('unsupported ID token signature')
    const discovery = await this.metadata()
    const response = await this.fetch(discovery.jwks_uri, {
      signal: AbortSignal.timeout(OIDC_TIMEOUT_MILLISECONDS),
    })
    if (!response.ok) throw new Error('JWKS fetch failed')
    const jwks = await response.json()
    const key = jwks.keys?.find((candidate) => candidate.kid === header.kid)
    if (key === undefined) throw new Error('ID token key missing')
    const verifier = createVerify('RSA-SHA256')
    verifier.update(`${headerPart}.${payloadPart}`)
    verifier.end()
    if (
      !verifier.verify(
        createPublicKey({ key, format: 'jwk' }),
        Buffer.from(signaturePart, 'base64url'),
      )
    )
      throw new Error('invalid ID token signature')
    const audience = Array.isArray(claims.aud) ? claims.aud : [claims.aud]
    if (
      claims.iss !== this.config.oidcAuthority ||
      !audience.includes(this.config.oidcClientId) ||
      (audience.length > 1 && claims.azp !== this.config.oidcClientId) ||
      (nonce !== undefined && claims.nonce !== nonce) ||
      typeof claims.sub !== 'string' ||
      typeof claims.iat !== 'number' ||
      claims.iat > now() + CLOCK_SKEW_SECONDS ||
      typeof claims.exp !== 'number' ||
      claims.exp <= now() - CLOCK_SKEW_SECONDS
    )
      throw new Error('invalid ID token claims')
    return `${claims.iss}\u0000${claims.sub}`
  }

  async authenticationFailure(response, identifier) {
    if (identifier) await this.store.destroy(identifier)
    response.writeHead(302, {
      Location: `${this.config.publicOrigin}/?authentication=failed`,
      'Set-Cookie': expiredSessionCookie(this.config),
      'Cache-Control': 'no-store',
    })
    response.end()
    return true
  }
}

export function validUnsafeRequest(request, config, csrfToken) {
  return (
    validBrowserContext(request, config) &&
    equalSecret(request.headers['x-remora-csrf'], csrfToken)
  )
}

export function validBrowserContext(request, config) {
  const fetchSite = request.headers['sec-fetch-site']
  return (
    request.headers.origin === config.publicOrigin &&
    (fetchSite === undefined ||
      fetchSite === 'same-origin' ||
      fetchSite === 'none')
  )
}

function authenticatedSession(tokens, identity, csrfToken) {
  if (
    typeof tokens.access_token !== 'string' ||
    typeof tokens.id_token !== 'string' ||
    typeof tokens.expires_in !== 'number' ||
    !Number.isFinite(tokens.expires_in) ||
    tokens.expires_in <= 0
  )
    throw new Error('token response incomplete')
  return {
    kind: 'authenticated',
    identity,
    csrfToken,
    accessToken: tokens.access_token,
    refreshToken:
      typeof tokens.refresh_token === 'string' ? tokens.refresh_token : null,
    idToken: tokens.id_token,
    accessExpiresAt: now() + tokens.expires_in,
  }
}

function validAuthenticatedSession(session) {
  return (
    session?.kind === 'authenticated' &&
    typeof session.identity === 'string' &&
    session.identity !== '' &&
    typeof session.csrfToken === 'string' &&
    session.csrfToken !== '' &&
    typeof session.accessToken === 'string' &&
    session.accessToken !== '' &&
    (session.refreshToken === null ||
      (typeof session.refreshToken === 'string' &&
        session.refreshToken !== '')) &&
    typeof session.idToken === 'string' &&
    session.idToken !== '' &&
    Number.isFinite(session.accessExpiresAt) &&
    session.accessExpiresAt > 0
  )
}

function clientAuthorization(config) {
  const identifier = formComponent(config.oidcClientId)
  const secret = formComponent(config.oidcClientSecret)
  return `Basic ${Buffer.from(`${identifier}:${secret}`).toString('base64')}`
}

function formComponent(value) {
  return new URLSearchParams({ value }).toString().slice('value='.length)
}

function now() {
  return Math.floor(Date.now() / 1000)
}

function json(response, status, value, headers = {}) {
  const body = JSON.stringify(value)
  response.writeHead(status, {
    ...headers,
    'Cache-Control': 'no-store',
    'Content-Type': 'application/json',
    'Content-Length': Buffer.byteLength(body),
  })
  response.end(body)
}