@nifrajs/graphql 4.0.2 → 5.1.0

package/src/ws.ts

@@ -16,13 +16,14 @@
 import type { NifraWebSocket, WebSocketContext, WebSocketHandler } from "@nifrajs/core/ws"
 import {
   type DocumentNode,
+  type ExecutionArgs,
   type ExecutionResult,
   execute,
   GraphQLError,
   type GraphQLSchema,
   parse,
-  specifiedRules,
   subscribe,
+  type ValidationRule,
   validate,
 } from "graphql"
 import {
@@ -38,7 +39,9 @@ import {
   documentMetrics,
   type GraphqlLimits,
   graphqlLimits,
+  overlapLimitError,
   requestLimitError,
+  validationRules,
 } from "./limits.ts"
 
 /** Per-connection state carried on `ws.data` for the graphql-ws bridge. */
@@ -49,6 +52,16 @@ interface GraphqlWsConnection {
   /** Tear the graphql-ws connection down (graphql-ws's `opened()` return value). */
   closed?: (code: number, reason: string) => void | Promise<void>
   readonly subscriptions: Set<string>
+  /** One abort controller per running operation id, aborted by the client's complete or a close. */
+  operations?: Map<string, AbortController>
+  /** Set once graphql-ws sends `connection_ack`, i.e. `connection_init` was admitted. */
+  acknowledged?: boolean
+}
+
+/** The `extra` graphql-ws hands its hooks for one socket. */
+interface GraphqlWsExtra {
+  readonly request: Request
+  readonly operations: Map<string, AbortController>
 }
 
 export interface GraphqlWsOptions<
@@ -58,6 +71,11 @@ export interface GraphqlWsOptions<
   readonly schema: GraphQLSchema
   /** Build the resolver `contextValue` from the upgrade request. Shared with the HTTP transport. */
   readonly context?: GraphqlContextBuilder<Context>
+  /**
+   * Admit a socket from its upgrade request, before {@link onConnect}; return `false` (or throw) to
+   * reject the connection. `mountGraphql({ subscriptions: true })` passes its HTTP `authorize` here.
+   */
+  readonly authorize?: (request: Request) => boolean | Promise<boolean>
   /**
    * Called once per connection with the client's `connection_init` payload; return `false` (or throw)
    * to reject the connection. The seam for authenticating a socket from an init token.
@@ -75,9 +93,19 @@ export interface GraphqlWsOptions<
   readonly maxAliases?: number
   readonly maxComplexity?: number
   readonly maxOperations?: number
+  readonly maxFieldOverlaps?: number
   /** Maximum concurrent subscription operation ids per connection. Default 20. */
   readonly maxSubscriptions?: number
-  /** Maximum time for initial execute/subscribe setup. Default 10 seconds. */
+  /**
+   * Answer introspection (`__schema`, `__type`). Default `true`. Set `false` to refuse such an
+   * operation with an `error` frame; `__typename` still resolves.
+   */
+  readonly introspection?: boolean
+  /**
+   * Maximum time for initial execute/subscribe setup. Default 10 seconds. Each operation's abort
+   * signal, which a resolver reads through `info.getAbortSignal()`, is aborted at this deadline, when
+   * the client completes the operation, and when the socket closes.
+   */
   readonly executionTimeoutMs?: number
 }
 
@@ -106,6 +134,7 @@ function subscriptionFrameError(
   message: string,
   schema: GraphQLSchema,
   limits: GraphqlLimits,
+  rules: readonly ValidationRule[],
 ): { readonly id: unknown; readonly message: string } | undefined {
   let frame: unknown
   try {
@@ -133,11 +162,13 @@ function subscriptionFrameError(
   } catch (error) {
     return { id, message: error instanceof GraphQLError ? error.message : "Invalid GraphQL query." }
   }
-  const validationErrors = validate(schema, document, specifiedRules)
-  if (validationErrors.length > 0)
-    return { id, message: validationErrors[0]?.message ?? "Invalid GraphQL query." }
-  const limitError = documentLimitError(documentMetrics(document), limits)
-  return limitError === undefined ? undefined : { id, message: limitError }
+  const limitError =
+    documentLimitError(documentMetrics(document), limits) ?? overlapLimitError(document, limits)
+  if (limitError !== undefined) return { id, message: limitError }
+  const validationErrors = validate(schema, document, rules)
+  return validationErrors.length > 0
+    ? { id, message: validationErrors[0]?.message ?? "Invalid GraphQL query." }
+    : undefined
 }
 
 function maskWsError(error: unknown): unknown {
@@ -155,12 +186,14 @@ function maskWsError(error: unknown): unknown {
  * so the id stops counting against `maxSubscriptions`: a client is not expected to complete an
  * operation the server already ended.
  */
-function outboundFrame(message: string, subscriptions: Set<string>): string {
+function outboundFrame(message: string, connection: GraphqlWsConnection): string {
   try {
     const parsed = recordOf(JSON.parse(message))
     if (parsed === undefined) return message
+    if (parsed.type === "connection_ack") connection.acknowledged = true
     if ((parsed.type === "complete" || parsed.type === "error") && typeof parsed.id === "string") {
-      subscriptions.delete(parsed.id)
+      connection.subscriptions.delete(parsed.id)
+      connection.operations?.delete(parsed.id)
     }
     if (parsed.type === "next") {
       const payload = recordOf(parsed.payload)
@@ -180,26 +213,68 @@ function outboundFrame(message: string, subscriptions: Set<string>): string {
   return message
 }
 
+/** Close a result that arrived after its operation timed out, if it is a stream nobody will read. */
+function disposeLate(result: unknown): void {
+  if (typeof result !== "object" || result === null || !(Symbol.asyncIterator in result)) return
+  const iterator = recordOf(result)
+  if (iterator === undefined || typeof iterator.return !== "function") return
+  Promise.resolve(iterator.return.call(result)).catch(() => {})
+}
+
+const wrappedSubscribeFields = new WeakSet<object>()
+
+/**
+ * graphql drops a subscribe resolver's result that settles after the operation's signal aborted, so
+ * a source the resolver opened anyway would stay open. Each root subscription field's resolver is
+ * wrapped once to return such a late source; with no signal, or before an abort, it is unchanged.
+ */
+function returnLateSources(schema: GraphQLSchema): void {
+  const fields = schema.getSubscriptionType()?.getFields() ?? {}
+  for (const field of Object.values(fields)) {
+    const resolver = field.subscribe
+    if (resolver === undefined || wrappedSubscribeFields.has(field)) continue
+    wrappedSubscribeFields.add(field)
+    field.subscribe = (source, args, context, info) => {
+      const result = resolver(source, args, context, info)
+      if (!(result instanceof Promise)) return result
+      return result.then((value) => {
+        if (info.getAbortSignal()?.aborted === true) disposeLate(value)
+        return value
+      })
+    }
+  }
+}
+
 /**
  * A timeout is that operation's error result. graphql-ws does not catch a rejected execute or
- * subscribe, so a rejection would end the whole socket instead of the one operation.
+ * subscribe, so a rejection would end the whole socket instead of the one operation. A timeout aborts
+ * the operation's `abortSignal`, which resolvers read through `info.getAbortSignal()`, and work that
+ * settles after it is not abandoned open: a late subscription stream is returned.
  */
 async function withWsTimeout<T>(
-  work: T | Promise<T>,
+  start: (abortSignal: AbortSignal) => T | Promise<T>,
+  outer: AbortSignal | null | undefined,
   timeoutMs: number,
 ): Promise<T | ExecutionResult> {
-  if (!(work instanceof Promise) || timeoutMs === 0) return work
+  const controller = new AbortController()
+  const work = start(outer ? AbortSignal.any([outer, controller.signal]) : controller.signal)
+  if (!(work instanceof Promise)) return work
   let timer: ReturnType<typeof setTimeout> | undefined
+  let timedOut = false
   const timeout = new Promise<ExecutionResult>((resolve) => {
-    timer = setTimeout(
-      () => resolve({ errors: [new GraphQLError("Execution timed out.")] }),
-      timeoutMs,
-    )
+    timer = setTimeout(() => {
+      timedOut = true
+      const error = new GraphQLError("Execution timed out.")
+      // Resolved first: the abort can reject the work, and the race must settle on this result.
+      resolve({ errors: [error] })
+      controller.abort(error)
+    }, timeoutMs)
   })
   try {
     return await Promise.race([work, timeout])
   } finally {
     if (timer !== undefined) clearTimeout(timer)
+    if (timedOut) work.then(disposeLate, () => {})
   }
 }
 
@@ -221,15 +296,29 @@ export function graphqlWebSocket<Context extends Record<string, unknown> = Recor
   options: GraphqlWsOptions<Context>,
 ): WebSocketHandler<GraphqlWsConnection> {
   const limits = graphqlWsLimits(options)
+  const rules = validationRules(options.introspection)
   const maxSubscriptions = options.maxSubscriptions ?? 20
   assertWsNumber(maxSubscriptions, "maxSubscriptions", 1)
   const executionTimeoutMs = options.executionTimeoutMs ?? 10_000
   assertWsNumber(executionTimeoutMs, "executionTimeoutMs", 0)
   const executeWithTimeout: typeof execute = (args) =>
-    withWsTimeout(execute(args), executionTimeoutMs)
+    executionTimeoutMs === 0
+      ? execute(args)
+      : withWsTimeout(
+          (abortSignal) => execute({ ...args, abortSignal }),
+          args.abortSignal,
+          executionTimeoutMs,
+        )
   const subscribeWithTimeout: typeof subscribe = (args) =>
-    withWsTimeout(subscribe(args), executionTimeoutMs)
-  const gqlServer = makeServer<Record<string, unknown>, { request: Request }>({
+    executionTimeoutMs === 0
+      ? subscribe(args)
+      : withWsTimeout(
+          (abortSignal) => subscribe({ ...args, abortSignal }),
+          args.abortSignal,
+          executionTimeoutMs,
+        )
+  returnLateSources(options.schema)
+  const gqlServer = makeServer<Record<string, unknown>, GraphqlWsExtra>({
     // Keep the raw escape hatch for auxiliary graphql-ws hooks, but make the package's safety
     // controls authoritative: callers cannot accidentally replace the bounded execute/subscribe
     // wrappers or the context/authentication seams with an unbounded implementation.
@@ -237,22 +326,41 @@ export function graphqlWebSocket<Context extends Record<string, unknown> = Recor
     schema: options.schema,
     execute: executeWithTimeout,
     subscribe: subscribeWithTimeout,
-    onConnect: options.onConnect
-      ? async (ctx: GraphqlWsContext<Record<string, unknown>, { request: Request }>) => {
-          try {
-            return await options.onConnect?.(
-              ctx.connectionParams as ConnectionInitMessage["payload"],
-              ctx.extra.request,
-            )
-          } catch {
-            return false
+    onConnect:
+      options.onConnect || options.authorize
+        ? async (ctx: GraphqlWsContext<Record<string, unknown>, GraphqlWsExtra>) => {
+            try {
+              if (options.authorize !== undefined && !(await options.authorize(ctx.extra.request)))
+                return false
+              return (
+                (await options.onConnect?.(
+                  ctx.connectionParams as ConnectionInitMessage["payload"],
+                  ctx.extra.request,
+                )) ?? true
+              )
+            } catch {
+              return false
+            }
           }
-        }
-      : undefined,
-    context: options.context
-      ? async (ctx: GraphqlWsContext<Record<string, unknown>, { request: Request }>) =>
-          await buildContext(options.context, { request: ctx.extra.request })
-      : undefined,
+        : undefined,
+    // graphql-ws builds the context once per operation, just before executing it, so the operation's
+    // abort signal is attached to its execution arguments here.
+    context: async (
+      ctx: GraphqlWsContext<Record<string, unknown>, GraphqlWsExtra>,
+      id: string,
+      _payload: unknown,
+      args: ExecutionArgs,
+    ) => {
+      const controller = new AbortController()
+      ctx.extra.operations.set(id, controller)
+      args.abortSignal =
+        args.abortSignal === undefined || args.abortSignal === null
+          ? controller.signal
+          : AbortSignal.any([args.abortSignal, controller.signal])
+      return options.context === undefined
+        ? undefined
+        : await buildContext(options.context, { request: ctx.extra.request })
+    },
   })
 
   return {
@@ -264,11 +372,13 @@ export function graphqlWebSocket<Context extends Record<string, unknown> = Recor
       const protocol =
         (ws.raw as { protocol?: string } | undefined)?.protocol ?? GRAPHQL_TRANSPORT_WS_PROTOCOL
       let onFrame: ((message: string) => void | Promise<void>) | undefined
+      const operations = new Map<string, AbortController>()
+      ws.data.operations = operations
       const closed = gqlServer.opened(
         {
           protocol,
           send: async (data: string) => {
-            ws.send(outboundFrame(data, ws.data.subscriptions))
+            ws.send(outboundFrame(data, ws.data))
           },
           close: (code: number, reason: string) => {
             ws.close(code, reason)
@@ -277,7 +387,7 @@ export function graphqlWebSocket<Context extends Record<string, unknown> = Recor
             onFrame = cb
           },
         },
-        { request: ws.data.request },
+        { request: ws.data.request, operations },
       )
       ws.data.closed = closed
       ws.data.dispatch = async (message: string) => {
@@ -297,7 +407,13 @@ export function graphqlWebSocket<Context extends Record<string, unknown> = Recor
       // this socket - including the client's `Complete` that stops the subscription. graphql-ws drives
       // its own sends and error frames from inside that loop; we only feed it inbound frames.
       const message = asText(data)
-      const rejected = subscriptionFrameError(message, options.schema, limits)
+      // Until connection_init is admitted, graphql-ws answers a subscribe frame by closing with 4401
+      // without reading the query. Parsing or validating it here would hand an unauthenticated socket
+      // schema suggestions and synchronous validation work.
+      const rejected =
+        ws.data.acknowledged === true
+          ? subscriptionFrameError(message, options.schema, limits, rules)
+          : undefined
       if (rejected !== undefined) {
         ws.send(
           JSON.stringify({
@@ -324,6 +440,8 @@ export function graphqlWebSocket<Context extends Record<string, unknown> = Recor
           ws.data.subscriptions.add(frame.id)
         } else if (frame?.type === "complete" && typeof frame.id === "string") {
           ws.data.subscriptions.delete(frame.id)
+          ws.data.operations?.get(frame.id)?.abort(new GraphQLError("Operation completed."))
+          ws.data.operations?.delete(frame.id)
         }
       } catch {
         // graphql-ws owns malformed protocol-frame handling; pass it through unchanged.
@@ -333,6 +451,9 @@ export function graphqlWebSocket<Context extends Record<string, unknown> = Recor
 
     close(ws: NifraWebSocket<GraphqlWsConnection>, code: number, reason: string): void {
       ws.data.subscriptions.clear()
+      for (const controller of ws.data.operations?.values() ?? [])
+        controller.abort(new GraphQLError("Socket closed."))
+      ws.data.operations?.clear()
       // The socket is gone: a failing iterator return() or onClose hook has no one left to answer.
       void Promise.resolve(ws.data.closed?.(code, reason)).catch(() => {})
     },