PotentialStorage/src/kv-storage/kv-storage-http-server.ts

import { AuthError } from '../utils/errors.js'
import { ExtendedJson } from '../utils/extended-json.js'
import { KeyRange } from '../utils/key-range.js'
import { AUTH_HEADERS, RequestAuth, type RequestAuthOptions, type SignatureVerifier } from '../utils/request-auth.js'
import type { BaseKVStorage, KVChange, ScanOptions } from './base-kv-storage.js'

export interface KVStorageHttpServerOptions {
  /**
   * `{ verifiers }`: requests must be signed (see `RequestAuth`), and a caller may only touch keys
   * under its own public key (`<publicKey>` and `<publicKey>/...`).
   * `'none'`: no authentication, everything is open; for local development and tests only.
   */
  auth: 'none' | ({ verifiers: SignatureVerifier[] } & RequestAuthOptions)
  /** Path prefix to strip, e.g. `/kv`. Default `''`. */
  basePath?: string
  /** Interval of SSE keep-alive comments. Default 15 000 ms. */
  keepAliveMs?: number
  /**
   * Origins allowed to call the server from a browser on another origin (CORS), e.g.
   * `['https://app.example.com']`, or `'*'` for any. Default: none (same origin only).
   */
  allowOrigins?: string[] | '*'
}

type RequestBody = { keys: string[]; changes: KVChange[]; paths: string[] } & ScanOptions

/**
 * Serves any KV storage over HTTP for `KVStorageHttp`. Framework-agnostic: `handle(Request)`
 * returns a `Response`. It executes exactly the operations clients send; it doesn't know about
 * versions or conflicts. Every route is `POST` with an Extended JSON body:
 *
 * - `/readMany { keys }`, `/write { changes }`, `/scan { gte, lt, limit, reverse, values }` (NDJSON).
 * - `/events { paths }`: Server-Sent Events relaying each change to those keys (and their
 *   subtrees) as it happens, including the subscriber's own writes. A `ready` event follows the
 *   connection. Nothing is replayed after a reconnect; clients catch up by reading.
 *
 * With authentication, each request is checked first (401), then every key, scan range and
 * subscription path must be the caller's public key or under it (403).
 *
 * Only changes made through the wrapped storage instance are relayed; several server processes
 * over one database each relay only their own writes.
 */
export class KVStorageHttpServer {
  private readonly auth: RequestAuth | undefined
  private readonly basePath: string
  private readonly keepAliveMs: number
  private readonly allowOrigins: string[] | '*'
  private readonly subscribers = new Set<{ ranges: KeyRange[]; send: (frame: string) => void }>()

  constructor(
    private readonly storage: BaseKVStorage,
    options: KVStorageHttpServerOptions,
  ) {
    this.auth = options.auth === 'none' ? undefined : new RequestAuth(options.auth.verifiers, options.auth)
    this.basePath = options.basePath ?? ''
    this.keepAliveMs = options.keepAliveMs ?? 15_000
    this.allowOrigins = options.allowOrigins ?? []
    storage.on('change', ({ changes }) => {
      for (const subscriber of this.subscribers) {
        const matching = changes.filter(({ key }) => subscriber.ranges.some((range) => range.contains(key)))
        if (matching.length) subscriber.send(`event: change\ndata: ${ExtendedJson.stringify(matching)}\n\n`)
      }
    })
  }

  async handle(request: Request): Promise<Response> {
    const response = await this.route(request)
    const origin = request.headers.get('origin')
    if (!origin || !(this.allowOrigins === '*' || this.allowOrigins.includes(origin))) return response
    response.headers.set('access-control-allow-origin', origin)
    response.headers.append('vary', 'origin')
    return response
  }

  private async route(request: Request): Promise<Response> {
    const path = new URL(request.url).pathname
    if (request.method === 'OPTIONS') {
      // CORS preflight (signed requests carry custom headers). `handle` adds the allowed origin.
      const headers = ['content-type', 'cache-control', 'last-event-id', ...Object.values(AUTH_HEADERS)].join(', ')
      return new Response(null, {
        status: 204,
        headers: { 'access-control-allow-methods': 'POST', 'access-control-allow-headers': headers, 'access-control-max-age': '600' },
      })
    }
    if (request.method !== 'POST' || !path.startsWith(this.basePath)) return KVStorageHttpServer.error(404, 'Not found')
    try {
      const text = await request.text()
      const owner = this.auth ? await this.auth.verify(request.headers, path, text) : undefined
      const body = ExtendedJson.parse(text) as RequestBody
      switch (path.slice(this.basePath.length)) {
        case '/readMany': {
          this.authorize(owner, body.keys)
          const values = await this.storage.readMany(body.keys)
          const missing = values.flatMap((v, i) => (v === undefined ? [i] : []))
          return new Response(ExtendedJson.stringify({ values: values.map((v) => v ?? null), missing }), { headers: { 'content-type': 'application/json' } })
        }
        case '/write':
          this.authorize(
            owner,
            body.changes.map(({ key }) => key),
          )
          await this.storage.writeMany(body.changes)
          return new Response(null, { status: 204 })
        case '/scan':
          if (owner !== undefined && !new KeyRange(owner).covers(body.gte, body.lt)) throw new AuthError(403, 'Scan outside your keys')
          return this.scan(body)
        case '/events':
          if (!Array.isArray(body.paths) || body.paths.length === 0) return KVStorageHttpServer.error(400, 'Subscribe to at least one path')
          this.authorize(owner, body.paths)
          return this.events(body.paths)
        default:
          return KVStorageHttpServer.error(404, 'Not found')
      }
    } catch (error) {
      if (error instanceof AuthError) return KVStorageHttpServer.error(error.status, error.message)
      return KVStorageHttpServer.error(500, error instanceof Error ? error.message : String(error))
    }
  }

  /** Every key must be the owner's key or under it (no-op without authentication). */
  private authorize(owner: string | undefined, keys: string[]): void {
    if (owner === undefined) return
    const allowed = new KeyRange(owner)
    for (const key of keys) {
      if (typeof key !== 'string' || !allowed.contains(key)) throw new AuthError(403, `No access to "${key}"`)
    }
  }

  private events(paths: string[]): Response {
    const encoder = new TextEncoder()
    let subscriber: { ranges: KeyRange[]; send: (frame: string) => void } | undefined
    let keepAlive: ReturnType<typeof setInterval> | undefined
    const stream = new ReadableStream<Uint8Array>({
      start: (controller) => {
        subscriber = { ranges: paths.map((path) => new KeyRange(path)), send: (frame) => controller.enqueue(encoder.encode(frame)) }
        this.subscribers.add(subscriber)
        subscriber.send('event: ready\ndata: null\n\n')
        keepAlive = setInterval(() => subscriber!.send(': keep-alive\n\n'), this.keepAliveMs)
        ;(keepAlive as { unref?: () => void }).unref?.()
      },
      cancel: () => {
        if (subscriber) this.subscribers.delete(subscriber)
        clearInterval(keepAlive)
      },
    })
    return new Response(stream, { headers: { 'content-type': 'text/event-stream', 'cache-control': 'no-cache' } })
  }

  private scan(options: ScanOptions): Response {
    const iterator = this.storage.scan(options)[Symbol.asyncIterator]()
    const encoder = new TextEncoder()
    const stream = new ReadableStream<Uint8Array>({
      async pull(controller) {
        // Send in small batches so short reads (skip-scan seeks) don't make the server run far ahead.
        let chunk = ''
        for (let i = 0; i < 32; i++) {
          const next = await iterator.next()
          if (next.done) {
            if (chunk) controller.enqueue(encoder.encode(chunk))
            return controller.close()
          }
          chunk += ExtendedJson.stringify(options.values ? next.value : [next.value[0]]) + '\n'
        }
        controller.enqueue(encoder.encode(chunk))
      },
      async cancel() {
        await iterator.return?.()
      },
    })
    return new Response(stream, { headers: { 'content-type': 'application/x-ndjson' } })
  }

  private static error(status: number, message: string): Response {
    return Response.json({ error: message }, { status })
  }
}