feat(prisma-next): v3 bulk-encrypt middleware (plain-JSONB payloads)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: 2 people

packages/prisma-next/src/middleware/bulk-encrypt.ts

@@ -194,7 +194,14 @@ function collectTargets(
   return targets
 }
 
-function stampRoutingKeysFromAst(ast: AnyQueryAst | undefined): void {
+/**
+ * Walk a lowered `InsertAst` / `UpdateAst` and stamp `(table, column)`
+ * routing context onto every cipherstash envelope embedded in a
+ * `ParamRef`. Version-neutral: it touches only the shared envelope base
+ * and the AST — no wire/codec knowledge — so the v3 middleware
+ * (`../v3/bulk-encrypt-v3.ts`) reuses it rather than forking the walk.
+ */
+export function stampRoutingKeysFromAst(ast: AnyQueryAst | undefined): void {
   if (!ast) return
   switch (ast.kind) {
     case 'insert':

packages/prisma-next/src/v3/bulk-encrypt-v3.ts

@@ -0,0 +1,155 @@
+/**
+ * v3 bulk-encrypt middleware — the plain-JSONB twin of the v2
+ * `../middleware/bulk-encrypt.ts`. Same `beforeExecute` lifecycle and
+ * batching contract, with two deliberate differences:
+ *
+ *   - **Jurisdiction**: params are matched against the closed v3
+ *     codec-id set ({@link CIPHERSTASH_V3_CODEC_ID_SET}), never the v2
+ *     set — the two middlewares partition the param space by codec id,
+ *     so a plan whose params carry v2 codec ids is invisible here.
+ *   - **Wire**: the param slot is replaced with the plain JSONB text of
+ *     the EQL payload ({@link v3ToDriver} — v3 columns are PG domains
+ *     over `jsonb`), never the v2 `eql_v2_encrypted` composite literal.
+ *     This module MUST NOT import `encodeEqlV2EncryptedWire`.
+ *
+ * The `(table, column)` routing-key stamping is NOT forked: the AST walk
+ * is version-neutral (it touches only the shared envelope base) and is
+ * reused from the v2 module via the exported
+ * {@link stampRoutingKeysFromAst}.
+ *
+ * Everything else — one `sdk.bulkEncrypt` per `(table, column)` group,
+ * `setHandleCiphertext` stamping with plaintext retention, `ctx.signal`
+ * forwarding by identity plus the pre-flight abort check and the abort
+ * race, and the ciphertext-count diagnostic — mirrors the v2 middleware;
+ * see its module doc for the full lifecycle rationale.
+ */
+
+import type {
+  ParamRefHandle,
+  SqlParamRefMutator,
+} from '@prisma-next/sql-relational-core/middleware'
+import type { SqlMiddleware } from '@prisma-next/sql-runtime'
+import { ifDefined } from '@prisma-next/utils/defined'
+import {
+  checkCipherstashAborted,
+  raceCipherstashAbort,
+} from '../execution/abort'
+import {
+  EncryptedEnvelopeBase,
+  setHandleCiphertext,
+} from '../execution/envelope-base'
+import { markBulkEncryptMiddlewareRegistered } from '../execution/middleware-registry'
+import { type BulkEncryptTarget, groupByRoutingKey } from '../execution/routing'
+import type { CipherstashSdk } from '../execution/sdk'
+import { CIPHERSTASH_V3_CODEC_ID_SET } from '../extension-metadata/constants-v3'
+// Version-neutral AST routing-key stamping — reused, not the v2 encode path.
+import { stampRoutingKeysFromAst } from '../middleware/bulk-encrypt'
+import { v3ToDriver } from './wire-v3'
+
+/**
+ * Construct the v3 bulk-encrypt middleware. The returned middleware is
+ * stateless aside from the captured `sdk` reference; one instance per
+ * (v3) runtime extension is the expected pattern.
+ */
+export function bulkEncryptMiddlewareV3(sdk: CipherstashSdk): SqlMiddleware {
+  // Mark this sdk as wired up so the codec's `encode` can distinguish
+  // "happy path: middleware will run later and fill in the ciphertext"
+  // from "misconfig: this sdk has no middleware registered". See
+  // `../execution/middleware-registry.ts`.
+  markBulkEncryptMiddlewareRegistered(sdk)
+  return {
+    name: 'cipherstash.bulk-encrypt-v3',
+    familyId: 'sql',
+    async beforeExecute(plan, ctx, params) {
+      if (!params) {
+        return
+      }
+
+      stampRoutingKeysFromAst(plan.ast)
+
+      const targets = collectV3Targets(params)
+      if (targets.length === 0) {
+        return
+      }
+
+      const groups = groupByRoutingKey(targets)
+      for (const [groupKey, group] of groups) {
+        const first = group[0]
+        if (!first) continue
+        const routingKey = first.routingKey
+
+        checkCipherstashAborted(ctx.signal, 'bulk-encrypt')
+        const ciphertexts = await raceCipherstashAbort(
+          sdk.bulkEncrypt({
+            routingKey,
+            values: group.map((t) => t.plaintext),
+            ...ifDefined('signal', ctx.signal),
+          }),
+          ctx.signal,
+          'bulk-encrypt',
+        )
+
+        if (ciphertexts.length !== group.length) {
+          throw new Error(
+            `cipherstash v3 bulk-encrypt: SDK returned ${ciphertexts.length} ciphertexts ` +
+              `for routing key ${groupKey} but ${group.length} were requested.`,
+          )
+        }
+
+        // Replace each ParamRef's value with the plain JSONB text the
+        // pg driver can serialise directly (the envelope instance
+        // itself would fail at the driver boundary). The envelope's
+        // ciphertext slot is also stamped via `setHandleCiphertext`
+        // and its plaintext slot retained, so follow-on reuse of the
+        // same envelope skips the re-encrypt round-trip.
+        params.replaceValues(
+          group.map((t, i) => {
+            const ciphertext = ciphertexts[i]
+            setHandleCiphertext(t.envelope, ciphertext)
+            return {
+              ref: t.ref,
+              newValue: v3ToDriver(ciphertext),
+            }
+          }),
+        )
+      }
+    },
+  }
+}
+
+function collectV3Targets(
+  params: SqlParamRefMutator,
+): BulkEncryptTarget<ParamRefHandle<string | undefined>>[] {
+  const targets: BulkEncryptTarget<ParamRefHandle<string | undefined>>[] = []
+  for (const entry of params.entries()) {
+    if (
+      entry.codecId === undefined ||
+      !CIPHERSTASH_V3_CODEC_ID_SET.has(entry.codecId)
+    )
+      continue
+    const value = entry.value
+    if (!(value instanceof EncryptedEnvelopeBase)) continue
+    const handle = value.expose()
+    if (handle.plaintext === undefined) {
+      throw new Error(
+        'cipherstash v3 bulk-encrypt: encountered an envelope with no plaintext on the write path. ' +
+          'Use the relevant `Encrypted*.from(plaintext)` factory to construct write-side envelopes.',
+      )
+    }
+    if (handle.table === undefined || handle.column === undefined) {
+      throw new Error(
+        'cipherstash v3 bulk-encrypt: envelope reached the bulk-encrypt phase without a (table, column) ' +
+          "routing context. The middleware's AST walk only handles `InsertAst` and `UpdateAst`; " +
+          'cipherstash envelopes embedded in other plan shapes (e.g. raw SQL) must stamp routing ' +
+          'context explicitly via `setHandleRoutingKey` before execute.',
+      )
+    }
+    targets.push({
+      ref: entry.ref,
+      plaintext: handle.plaintext,
+      envelope: value,
+      routingKey: { table: handle.table, column: handle.column },
+    })
+  }
+  return targets
+}

packages/prisma-next/test/bulk-encrypt-middleware.helpers.ts

@@ -0,0 +1,144 @@
+/**
+ * Shared harness for the bulk-encrypt middleware suites (v2
+ * `bulk-encrypt-middleware.test.ts` and v3 `v3/bulk-encrypt-v3.test.ts`).
+ * Version-neutral by construction: plan builders take the codec id to
+ * stamp on each `ParamRef` (defaulting to the v2 string codec id so
+ * existing v2 call sites read unchanged), and the mock SDK echoes
+ * whatever ciphertext shape the caller's `encryptImpl` produces.
+ */
+
+import type { Contract, PlanMeta } from '@prisma-next/contract/types'
+import type { SqlStorage } from '@prisma-next/sql-contract/types'
+import {
+  type ColumnRef,
+  InsertAst,
+  ParamRef,
+  TableSource,
+  UpdateAst,
+} from '@prisma-next/sql-relational-core/ast'
+import type { SqlExecutionPlan } from '@prisma-next/sql-relational-core/plan'
+import type { SqlMiddlewareContext } from '@prisma-next/sql-runtime'
+import { vi } from 'vitest'
+import type {
+  CipherstashBulkDecryptArgs,
+  CipherstashBulkEncryptArgs,
+  CipherstashSdk,
+  CipherstashSingleDecryptArgs,
+} from '../src/execution/sdk'
+import { CIPHERSTASH_STRING_CODEC_ID } from '../src/extension-metadata/constants'
+
+export { createSqlParamRefMutator } from '@prisma-next/sql-relational-core/middleware'
+
+export const baseMeta: PlanMeta = {
+  target: 'postgres',
+  storageHash: 'sha256:test',
+  lane: 'dsl',
+}
+
+export function createCtx(
+  overrides?: Partial<SqlMiddlewareContext>,
+): SqlMiddlewareContext {
+  return {
+    contract: {} as Contract<SqlStorage>,
+    mode: 'strict' as const,
+    scope: 'runtime' as const,
+    now: () => Date.now(),
+    log: {
+      info: vi.fn(),
+      warn: vi.fn(),
+      error: vi.fn(),
+    },
+    contentHash: async () => 'mock-hash',
+    planExecutionId: 'test-plan-execution',
+    ...overrides,
+  }
+}
+
+export interface CounterSdk extends CipherstashSdk {
+  readonly bulkEncryptCalls: CipherstashBulkEncryptArgs[]
+  readonly bulkDecryptCalls: CipherstashBulkDecryptArgs[]
+  readonly singleDecryptCalls: CipherstashSingleDecryptArgs[]
+}
+
+export function makeCounterSdk(options?: {
+  encryptImpl?: (args: CipherstashBulkEncryptArgs) => ReadonlyArray<unknown>
+}): CounterSdk {
+  const bulkEncryptCalls: CipherstashBulkEncryptArgs[] = []
+  const bulkDecryptCalls: CipherstashBulkDecryptArgs[] = []
+  const singleDecryptCalls: CipherstashSingleDecryptArgs[] = []
+  const encryptImpl =
+    options?.encryptImpl ??
+    ((args: CipherstashBulkEncryptArgs) =>
+      args.values.map(
+        (plaintext) =>
+          `cipher:${args.routingKey.table}.${args.routingKey.column}:${plaintext}`,
+      ))
+  return {
+    bulkEncryptCalls,
+    bulkDecryptCalls,
+    singleDecryptCalls,
+    decrypt(args) {
+      singleDecryptCalls.push(args)
+      return Promise.resolve(`single:${String(args.ciphertext)}`)
+    },
+    bulkEncrypt(args) {
+      bulkEncryptCalls.push(args)
+      return Promise.resolve(encryptImpl(args))
+    },
+    bulkDecrypt(args) {
+      bulkDecryptCalls.push(args)
+      return Promise.resolve(
+        args.ciphertexts.map((c) => `bulk-decrypt:${String(c)}`),
+      )
+    },
+  }
+}
+
+export function buildInsertPlan(
+  table: string,
+  rows: ReadonlyArray<Record<string, unknown>>,
+  codecId: string = CIPHERSTASH_STRING_CODEC_ID,
+): SqlExecutionPlan {
+  const params: unknown[] = []
+  const astRows = rows.map((row) => {
+    const out: Record<string, ParamRef> = {}
+    for (const [column, value] of Object.entries(row)) {
+      const ref = ParamRef.of(value, {
+        codec: { codecId },
+      })
+      out[column] = ref
+      params.push(value)
+    }
+    return out
+  })
+  const ast = new InsertAst(TableSource.named(table), astRows)
+  return {
+    sql: `INSERT INTO "${table}" (...) VALUES (...)`,
+    params,
+    meta: { ...baseMeta },
+    ast,
+  } as SqlExecutionPlan
+}
+
+export function buildUpdatePlan(
+  table: string,
+  set: Record<string, unknown>,
+  codecId: string = CIPHERSTASH_STRING_CODEC_ID,
+): SqlExecutionPlan {
+  const params: unknown[] = []
+  const astSet: Record<string, ParamRef | ColumnRef> = {}
+  for (const [column, value] of Object.entries(set)) {
+    const ref = ParamRef.of(value, {
+      codec: { codecId },
+    })
+    astSet[column] = ref
+    params.push(value)
+  }
+  const ast = new UpdateAst(TableSource.named(table), astSet)
+  return {
+    sql: `UPDATE "${table}" SET ...`,
+    params,
+    meta: { ...baseMeta },
+    ast,
+  } as SqlExecutionPlan
+}