From e81dc7b23a13839433fd3db3159065bd5131cdda Mon Sep 17 00:00:00 2001
From: bglgwyng <bglgwyng@gmail.com>
Date: Wed, 9 Aug 2023 20:37:31 +0900
Subject: [PATCH] Parse persited query on subscription.

---
 lib/routes.js                  | 44 ++++++++++++++++++++--------------
 lib/subscription-connection.js |  6 +++--
 lib/subscription.js            | 10 ++++----
 3 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/lib/routes.js b/lib/routes.js
index a6ab508b..34b48072 100644
--- a/lib/routes.js
+++ b/lib/routes.js
@@ -214,7 +214,7 @@ module.exports = async function (app, opts) {
     notSupportedError
   } = persistedQueryProvider || {}
 
-  async function executeQuery (query, variables, operationName, request, reply) {
+  async function executeQuery(query, variables, operationName, request, reply) {
     // Validate a query is present
     if (!query) {
       return new MER_ERR_GQL_PERSISTED_QUERY_NOT_FOUND('Unknown query')
@@ -237,10 +237,8 @@ module.exports = async function (app, opts) {
     return executeQuery(query, variables, operationName, request, reply)
   }
 
-  async function executePersistedQuery (body, request, reply) {
+  async function getPersistedQuery (body) {
     let { query } = body
-    const { operationName, variables } = body
-
     // Verify if a query matches the persisted format
     const persisted = isPersistedQuery(body)
     if (persisted) {
@@ -266,23 +264,32 @@ module.exports = async function (app, opts) {
       }
     }
 
-    // Execute the query
-    const result = await executeQuery(query, variables, operationName, request, reply)
-
-    // Only save queries which are not yet persisted
-    if (!persisted && query) {
-      // If provided the getHashForQuery, saveQuery settings we save this query
-      const hash = getHashForQuery && getHashForQuery(query)
-      if (hash) {
-        try {
-          await saveQuery(hash, query)
-        } catch (err) {
-          request.log.warn({ err, hash, query }, 'Failed to persist query')
+    return {
+      query,
+      persisted,
+      async saveQuery () {
+        const hash = getHashForQuery && getHashForQuery(query)
+        if (hash) {
+          try {
+            await saveQuery(hash, query)
+          } catch (err) {
+            // FIMXE: memory leak vulnerability
+            // request.log.warn({ err, hash, query }, 'Failed to persist query')
+          }
         }
       }
     }
+  }
+
+  async function executePersistedQuery (body, request, reply) {
+    const { operationName, variables } = body
+
+    const { query, saveQuery } = await getPersistedQuery(body)
+
+    const result = await executeQuery(query, variables, operationName, request, reply)
+
+    await saveQuery()
 
-    // Return the result
     return result
   }
 
@@ -326,7 +333,8 @@ module.exports = async function (app, opts) {
       subscriptionContextFn,
       keepAlive,
       fullWsTransport,
-      errorFormatter
+      errorFormatter,
+      getPersistedQuery,
     })
   } else {
     app.route(getOptions)
diff --git a/lib/subscription-connection.js b/lib/subscription-connection.js
index 65d5e1da..71d8f70e 100644
--- a/lib/subscription-connection.js
+++ b/lib/subscription-connection.js
@@ -20,7 +20,8 @@ module.exports = class SubscriptionConnection {
     resolveContext,
     keepAlive,
     fullWsTransport,
-    errorFormatter
+    errorFormatter,
+    getPersistedQuery
   }) {
     this.fastify = fastify
     this.socket = socket
@@ -37,6 +38,7 @@ module.exports = class SubscriptionConnection {
     this.fullWsTransport = fullWsTransport
     this.errorFormatter = errorFormatter
     this.headers = {}
+    this.getPersistedQuery = getPersistedQuery
 
     this.protocolMessageTypes = getProtocolByName(socket.protocol)
     this.socket.on('error', this.handleConnectionClose.bind(this))
@@ -208,7 +210,7 @@ module.exports = class SubscriptionConnection {
       }
     }
 
-    const document = typeof query !== 'string' ? query : parse(query)
+    const document = typeof query !== 'string' ? query : parse((await this.getPersistedQuery(payload)).query)
 
     if (!document) {
       throw new Error('Must provide document.')
diff --git a/lib/subscription.js b/lib/subscription.js
index c35cb590..2e4d6d16 100644
--- a/lib/subscription.js
+++ b/lib/subscription.js
@@ -6,7 +6,7 @@ const { kHooks } = require('./symbols')
 const SubscriptionConnection = require('./subscription-connection')
 const { getProtocolByName } = require('./subscription-protocol')
 
-function createConnectionHandler ({ subscriber, fastify, onConnect, onDisconnect, entityResolversFactory, subscriptionContextFn, keepAlive, fullWsTransport, errorFormatter }) {
+function createConnectionHandler ({ subscriber, fastify, onConnect, onDisconnect, entityResolversFactory, subscriptionContextFn, keepAlive, fullWsTransport, errorFormatter, getPersistedQuery }) {
   return async (connection, request) => {
     const { socket } = connection
 
@@ -48,7 +48,8 @@ function createConnectionHandler ({ subscriber, fastify, onConnect, onDisconnect
       resolveContext,
       keepAlive,
       fullWsTransport,
-      errorFormatter
+      errorFormatter,
+      getPersistedQuery
     })
 
     /* istanbul ignore next */
@@ -62,7 +63,7 @@ function createConnectionHandler ({ subscriber, fastify, onConnect, onDisconnect
 }
 
 module.exports = async function (fastify, opts) {
-  const { getOptions, subscriber, verifyClient, onConnect, onDisconnect, entityResolversFactory, subscriptionContextFn, keepAlive, fullWsTransport, errorFormatter } = opts
+  const { getOptions, subscriber, verifyClient, onConnect, onDisconnect, entityResolversFactory, subscriptionContextFn, keepAlive, fullWsTransport, errorFormatter, getPersistedQuery } = opts
 
   // If `fastify.websocketServer` exists, it means `@fastify/websocket` already registered.
   // Without this check, @fastify/websocket will be registered multiple times and raises FST_ERR_DEC_ALREADY_PRESENT.
@@ -86,7 +87,8 @@ module.exports = async function (fastify, opts) {
       subscriptionContextFn,
       keepAlive,
       fullWsTransport,
-      errorFormatter
+      errorFormatter,
+      getPersistedQuery
     })
   })
 }