Repository: mercurius-js/mercurius
Pull request #896 (2022-10-24 12:39:39)

Use the real Request object in subscriptions
Signed-off-by: Matteo Collina <hello@matteocollina.com>

This PR does a few things:

1. uses the real `Request` object in the subscription `Context`. Reply is still "faked" as there is no `Reply` for websockets.
2. copies over all the data incoming as part of `connection_init` packet to `req.headers`. If there is a `headers` object that's used instead (as it was before). This change was needed because GraphiQL treats all the payload as headers.


mcollina: Docs added.

Patch:
diff --git a/docs/subscriptions.md b/docs/subscriptions.md
index b043b2c..7d53eba 100644
--- a/docs/subscriptions.md
+++ b/docs/subscriptions.md
@@ -5,6 +5,7 @@
 - [Subscriptions](#subscriptions)
   - [Subscription support (simple)](#subscription-support-simple)
   - [Subscription filters](#subscription-filters)
+  - [Subscription Context](#subscription-context)
   - [Build a custom GraphQL context object for subscriptions](#build-a-custom-graphql-context-object-for-subscriptions)
   - [Subscription support (with redis)](#subscription-support-with-redis)
   - [Subscriptions with custom PubSub](#subscriptions-with-custom-pubsub)
@@ -153,6 +154,18 @@ app.register(mercurius, {
 })
 ```
 
+### Subscription Context
+
+The context for the `Subscription` includes:
+
+* `app`, the Fastify application
+* `reply`, an object that pretend to be a `Reply` object from Fastify without its decorators.
+* `reply.request`, the real request object from Fastify
+
+During the connection initialization phase, all content of proerty `payload` of the `connection_init` packet
+is automatically copied inside `request.headers`. In case an `headers` property is specified within `payload`,
+that's used instead.
+
 ### Build a custom GraphQL context object for subscriptions
 
 ```js
diff --git a/lib/subscription-connection.js b/lib/subscription-connection.js
index 4835222..8bd7d30 100644
--- a/lib/subscription-connection.js
+++ b/lib/subscription-connection.js
@@ -159,8 +159,12 @@ module.exports = class SubscriptionConnection {
 
     this.context._connectionInit = data.payload
 
-    if (data.payload && data.payload.headers) {
-      this.headers = data.payload.headers
+    if (data.payload) {
+      if (data.payload.headers) {
+        this.headers = data.payload.headers
+      } else {
+        this.headers = data.payload
+      }
     }
 
     this.sendMessage(this.protocolMessageTypes.GQL_CONNECTION_ACK)
@@ -257,6 +261,13 @@ module.exports = class SubscriptionConnection {
       subscriptionLoaders = this.fastify[kSubscriptionFactory]
     }
 
+    for (const [key, value] of Object.entries(this.headers)) {
+      // as a security mechanism, do not override existing headers
+      if (context.request.headers[key] === undefined) {
+        context.request.headers[key] = value
+      }
+    }
+
     const subIter = await subscribe({
       schema,
       document,
@@ -270,7 +281,7 @@ module.exports = class SubscriptionConnection {
         lruGatewayResolvers: this.lruGatewayResolvers,
         reply: {
           entityResolversFactory: this.entityResolversFactory,
-          request: { headers: this.headers },
+          request: context.request,
           [kLoaders]:
             subscriptionLoaders && subscriptionLoaders.create(context),
           [kEntityResolvers]: this.entityResolvers
diff --git a/lib/subscription.js b/lib/subscription.js
index 710c0dc..f75acb6 100644
--- a/lib/subscription.js
+++ b/lib/subscription.js
@@ -22,7 +22,8 @@ function createConnectionHandler ({ subscriber, fastify, onConnect, onDisconnect
 
     let context = {
       app: fastify,
-      pubsub: subscriber
+      pubsub: subscriber,
+      request
     }
 
     if (context.app.graphql && context.app.graphql[kHooks]) {


Test patch:
diff --git a/test/subscription.js b/test/subscription.js
index 143e1ff..4abaa4d 100644
--- a/test/subscription.js
+++ b/test/subscription.js
@@ -1,4 +1,6 @@
 'use strict'
+// we test that __proto__ is actually ignored
+/* eslint-disable no-proto */
 
 const { test, t } = require('tap')
 const Fastify = require('fastify')
@@ -204,7 +206,9 @@ test('subscription server sends update to subscriptions', t => {
     },
     Subscription: {
       notificationAdded: {
-        subscribe: (root, args, { pubsub }) => pubsub.subscribe('NOTIFICATION_ADDED')
+        subscribe: (root, args, ctx) => {
+          return ctx.pubsub.subscribe('NOTIFICATION_ADDED')
+        }
       }
     }
   }
@@ -2102,3 +2106,377 @@ test('subscription passes context to its loaders', t => {
     })
   })
 })
+
+test('request and reply objects in subscription context', t => {
+  const app = Fastify()
+  t.teardown(() => app.close())
+
+  const sendTestQuery = () => {
+    app.inject({
+      method: 'POST',
+      url: '/graphql',
+      body: {
+        query: `
+          query {
+            notifications {
+              id
+              message
+            }
+          }
+        `
+      }
+    }, () => {
+      sendTestMutation()
+    })
+  }
+
+  const sendTestMutation = () => {
+    app.inject({
+      method: 'POST',
+      url: '/graphql',
+      body: {
+        query: `
+          mutation {
+            addNotification(message: "Hello World") {
+              id
+            }
+          }
+        `
+      }
+    }, () => {})
+  }
+
+  const emitter = mq()
+  const schema = `
+    type Notification {
+      id: ID!
+      message: String
+    }
+
+    type Query {
+      notifications: [Notification]
+    }
+
+    type Mutation {
+      addNotification(message: String): Notification
+    }
+
+    type Subscription {
+      notificationAdded: Notification
+    }
+  `
+
+  app.decorateRequest('foo', function () { return 'bar' })
+
+  let idCount = 1
+  const notifications = [{
+    id: idCount,
+    message: 'Notification message'
+  }]
+
+  const resolvers = {
+    Query: {
+      notifications: () => notifications
+    },
+    Mutation: {
+      addNotification: async (_, { message }) => {
+        const id = idCount++
+        const notification = {
+          id,
+          message
+        }
+        notifications.push(notification)
+        await emitter.emit({
+          topic: 'NOTIFICATION_ADDED',
+          payload: {
+            notificationAdded: notification
+          }
+        })
+
+        return notification
+      }
+    },
+    Subscription: {
+      notificationAdded: {
+        subscribe: (root, args, ctx) => {
+          t.equal(ctx.reply.request.foo(), 'bar')
+          t.equal(ctx.reply.request.headers.authorization, 'Bearer foobar')
+          return ctx.pubsub.subscribe('NOTIFICATION_ADDED')
+        }
+      }
+    }
+  }
+
+  app.register(GQL, {
+    schema,
+    resolvers,
+    subscription: {
+      emitter
+    }
+  })
+
+  app.listen({ port: 0 }, err => {
+    t.error(err)
+
+    const ws = new WebSocket('ws://localhost:' + (app.server.address()).port + '/graphql', 'graphql-ws')
+    const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8', objectMode: true })
+    t.teardown(client.destroy.bind(client))
+    client.setEncoding('utf8')
+
+    client.write(JSON.stringify({
+      type: 'connection_init',
+      payload: {
+        headers: {
+          authorization: 'Bearer foobar'
+        }
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 1,
+      type: 'start',
+      payload: {
+        query: `
+          subscription {
+            notificationAdded {
+              id
+              message
+            }
+          }
+        `
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 2,
+      type: 'start',
+      payload: {
+        query: `
+          subscription {
+            notificationAdded {
+              id
+              message
+            }
+          }
+        `
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 2,
+      type: 'stop'
+    }))
+
+    client.on('data', chunk => {
+      const data = JSON.parse(chunk)
+
+      if (data.id === 1 && data.type === 'data') {
+        t.equal(chunk, JSON.stringify({
+          type: 'data',
+          id: 1,
+          payload: {
+            data: {
+              notificationAdded: {
+                id: '1',
+                message: 'Hello World'
+              }
+            }
+          }
+        }))
+
+        client.end()
+        t.end()
+      } else if (data.id === 2 && data.type === 'complete') {
+        sendTestQuery()
+      }
+    })
+  })
+})
+
+test('request and reply objects in subscription context - no headers wrapper', t => {
+  const app = Fastify()
+  t.teardown(() => app.close())
+
+  const sendTestQuery = () => {
+    app.inject({
+      method: 'POST',
+      url: '/graphql',
+      body: {
+        query: `
+          query {
+            notifications {
+              id
+              message
+            }
+          }
+        `
+      }
+    }, () => {
+      sendTestMutation()
+    })
+  }
+
+  const sendTestMutation = () => {
+    app.inject({
+      method: 'POST',
+      url: '/graphql',
+      body: {
+        query: `
+          mutation {
+            addNotification(message: "Hello World") {
+              id
+            }
+          }
+        `
+      }
+    }, () => {})
+  }
+
+  const emitter = mq()
+  const schema = `
+    type Notification {
+      id: ID!
+      message: String
+    }
+
+    type Query {
+      notifications: [Notification]
+    }
+
+    type Mutation {
+      addNotification(message: String): Notification
+    }
+
+    type Subscription {
+      notificationAdded: Notification
+    }
+  `
+
+  app.decorateRequest('foo', function () { return 'bar' })
+
+  let idCount = 1
+  const notifications = [{
+    id: idCount,
+    message: 'Notification message'
+  }]
+
+  const resolvers = {
+    Query: {
+      notifications: () => notifications
+    },
+    Mutation: {
+      addNotification: async (_, { message }) => {
+        const id = idCount++
+        const notification = {
+          id,
+          message
+        }
+        notifications.push(notification)
+        await emitter.emit({
+          topic: 'NOTIFICATION_ADDED',
+          payload: {
+            notificationAdded: notification
+          }
+        })
+
+        return notification
+      }
+    },
+    Subscription: {
+      notificationAdded: {
+        subscribe: (root, args, ctx) => {
+          t.equal(ctx.reply.request.foo(), 'bar')
+          t.equal(ctx.reply.request.headers.authorization, 'Bearer foobar')
+          t.equal(ctx.reply.request.headers.constructor, Object)
+          t.equal(ctx.reply.request.headers.__proto__, {}.__proto__)
+          return ctx.pubsub.subscribe('NOTIFICATION_ADDED')
+        }
+      }
+    }
+  }
+
+  app.register(GQL, {
+    schema,
+    resolvers,
+    subscription: {
+      emitter
+    }
+  })
+
+  app.listen({ port: 0 }, err => {
+    t.error(err)
+
+    const ws = new WebSocket('ws://localhost:' + (app.server.address()).port + '/graphql', 'graphql-ws')
+    const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8', objectMode: true })
+    t.teardown(client.destroy.bind(client))
+    client.setEncoding('utf8')
+
+    client.write(JSON.stringify({
+      type: 'connection_init',
+      payload: {
+        authorization: 'Bearer foobar',
+        constructor: 'aaa',
+        __proto__: 'bbb'
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 1,
+      type: 'start',
+      payload: {
+        query: `
+          subscription {
+            notificationAdded {
+              id
+              message
+            }
+          }
+        `
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 2,
+      type: 'start',
+      payload: {
+        query: `
+          subscription {
+            notificationAdded {
+              id
+              message
+            }
+          }
+        `
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 2,
+      type: 'stop'
+    }))
+
+    client.on('data', chunk => {
+      const data = JSON.parse(chunk)
+
+      if (data.id === 1 && data.type === 'data') {
+        t.equal(chunk, JSON.stringify({
+          type: 'data',
+          id: 1,
+          payload: {
+            data: {
+              notificationAdded: {
+                id: '1',
+                message: 'Hello World'
+              }
+            }
+          }
+        }))
+
+        client.end()
+        t.end()
+      } else if (data.id === 2 && data.type === 'complete') {
+        sendTestQuery()
+      }
+    })
+  })
+})