Repository: mcollina/fastify-gql
Pull request #271 (2020-08-31 12:22:49)

feat: GraphQL over WebSocket protocol extensions and 'connectionInit' extension
This PR is an implementation of the solution proposed in #268 that should solve the `connection_init` handshake issue.

PacoDu: Waiting for #270, otherwise I think this PR is ready.
PacoDu: > Good work! Can you please add a full end-to-end test where this is fully tested?

This is already included in test/gateway/subscription.js (https://github.com/mcollina/fastify-gql/pull/271/files#diff-b9cdf3fa0a5c08b4c70395f7f82a9930R388). Is there anything missing ?
mcollina: I think we should have a test with two websocket clients to validate the end use case.
PacoDu: I think this is ready now 🚀 !

Patch:
diff --git a/README.md b/README.md
index 7b4ec7a..2f11a22 100644
--- a/README.md
+++ b/README.md
@@ -875,6 +875,8 @@ __fastify-gql__ supports the following options:
   [GraphQL IDE](https://www.npmjs.com/package/graphql-playground-react) on `/playground` if `'playground'`
   and if `routes` is `true`. Leave empty or `false` to disable.
   _only applies if `onlyPersisted` option is not `true`_
+* `playgroundSettings` Object. that allow you to configure GraphQL Playground with [playground 
+   options](https://github.com/prisma-labs/graphql-playground#usage). it works if the graphiql is set to `'playground'`.
 * `jit`: Integer. The minimum number of execution a query needs to be
   executed before being jit'ed.
 * `routes`: boolean. Serves the Default: `true`. A graphql endpoint is
diff --git a/bench.sh b/bench.sh
index 998918a..9202e4e 100755
--- a/bench.sh
+++ b/bench.sh
@@ -1,3 +1,3 @@
 #! /bin/bash
 
-./node_modules/.bin/autocannon -c 100 -d 5 -p 10 --on-port '/graphql?query={add(x:2,y:2)}' -- node example.js
+./node_modules/.bin/autocannon -c 100 -d 5 -p 10 --on-port '/graphql?query={add(x:2,y:2)}' -- node examples/basic.js
diff --git a/docs/protocol-extension.md b/docs/protocol-extension.md
new file mode 100644
index 0000000..d053ded
--- /dev/null
+++ b/docs/protocol-extension.md
@@ -0,0 +1,36 @@
+# GraphQL over WebSocket Protocol extension
+
+The GraphQL over WebSocket Protocol used by `fastify-gql` follows apollo's protocol defined (here)[https://github.com/apollographql/subscriptions-transport-ws/blob/master/PROTOCOL.md]. In order to achieve gateway-to-service communication and handle `connection_init` per client, an extension is used on the protocol for server-to-server communication.
+
+## Message structure extension
+
+This is the extended structure of each message:
+```
+export interface OperationMessage {
+  payload?: any;
+  id?: string;
+  type: string;
+
+  extensions?: Array<OperationExtension>;
+}
+
+export interface OperationExtension {
+  type: string;
+  payload?: any;
+}
+```
+
+## Server -> Server
+### `connectionInit` extension
+
+Gateway uses this extension to share the `connection_init` payload with a service when the connection is already established between gateway and services.
+
+```
+export interface ConnectionInitExtension extends OperationExtension {
+  type: string;
+  payload?: Object;
+}
+```
+
+* `type: String` : 'connectionInit'
+* `payload: Object` : optional parameters that the client specifies in connectionParams
diff --git a/index.js b/index.js
index 89831da..4a57e63 100644
--- a/index.js
+++ b/index.js
@@ -199,6 +199,7 @@ const plugin = fp(async function (app, opts) {
       errorHandler: opts.errorHandler,
       errorFormatter: opts.errorFormatter,
       ide: optsIde,
+      ideSettings: opts.playgroundSettings,
       prefix: opts.prefix,
       path: opts.path,
       context: opts.context,
diff --git a/lib/gateway/make-resolver.js b/lib/gateway/make-resolver.js
index 66e4223..4449568 100644
--- a/lib/gateway/make-resolver.js
+++ b/lib/gateway/make-resolver.js
@@ -404,7 +404,7 @@ function makeResolver ({ service, createOperation, transformData, isQuery, isRef
     }
 
     if (isSubscription) {
-      const subscriptionId = service.createSubscription(query, variables, pubsub.publish.bind(pubsub))
+      const subscriptionId = service.createSubscription(query, variables, pubsub.publish.bind(pubsub), context._connectionInit)
       return pubsub.subscribe(`${service.name}_${subscriptionId}`)
     }
 
diff --git a/lib/routes.js b/lib/routes.js
index a1fa099..f652819 100644
--- a/lib/routes.js
+++ b/lib/routes.js
@@ -315,6 +315,7 @@ module.exports = async function (app, opts) {
           GraphQLPlayground.init(document.getElementById('root'), {
             subscriptionEndpoint: '${opts.prefix || ''}${graphqlPath}',
             endpoint: '${opts.prefix || ''}${graphqlPath}',
+            settings: ${JSON.stringify(opts.ideSettings)},
           });
         });`)
       })
diff --git a/lib/subscription-client.js b/lib/subscription-client.js
index 6558279..19a09a4 100644
--- a/lib/subscription-client.js
+++ b/lib/subscription-client.js
@@ -8,6 +8,7 @@ const {
   GQL_CONNECTION_KEEP_ALIVE,
   GQL_START,
   GQL_DATA,
+  GQL_ERROR,
   GQL_COMPLETE,
   GQL_STOP,
   GRAPHQL_WS
@@ -90,11 +91,12 @@ class SubscriptionClient {
 
       if (tryReconnect) {
         for (const operationId of this.operations.keys()) {
-          const { options, handler } = this.operations.get(operationId)
+          const { options, handler, extensions } = this.operations.get(operationId)
 
           this.operations.set(operationId, {
             options,
             handler,
+            extensions,
             started: false
           })
         }
@@ -146,12 +148,13 @@ class SubscriptionClient {
     }
   }
 
-  sendMessage (operationId, type, payload = {}) {
+  sendMessage (operationId, type, payload = {}, extensions) {
     this.socket.send(
       JSON.stringify({
         id: operationId,
         type,
-        payload
+        payload,
+        extensions
       })
     )
   }
@@ -196,6 +199,14 @@ class SubscriptionClient {
           operation.handler(data.payload.data)
         }
         break
+      case GQL_ERROR:
+        /* istanbul ignore else */
+        if (operation) {
+          operation.handler(null)
+          this.operations.delete(operationId)
+          this.sendMessage(operationId, GQL_ERROR, data.payload)
+        }
+        break
       case GQL_COMPLETE:
         /* istanbul ignore else */
         if (operation) {
@@ -220,17 +231,17 @@ class SubscriptionClient {
   }
 
   startOperation (operationId) {
-    const { started, options, handler } = this.operations.get(operationId)
+    const { started, options, handler, extensions } = this.operations.get(operationId)
     if (!started) {
       if (!this.ready) {
         throw new Error('Connection is not ready')
       }
-      this.operations.set(operationId, { started: true, options, handler })
-      this.sendMessage(operationId, GQL_START, options)
+      this.operations.set(operationId, { started: true, options, handler, extensions })
+      this.sendMessage(operationId, GQL_START, options, extensions)
     }
   }
 
-  createSubscription (query, variables, publish) {
+  createSubscription (query, variables, publish, connectionInit) {
     const subscriptionString = JSON.stringify({ query, variables })
     let operationId = this.subscriptionQueryMap[subscriptionString]
 
@@ -241,7 +252,7 @@ class SubscriptionClient {
 
     operationId = String(++this.operationId)
 
-    this.operations.set(operationId, {
+    const operation = {
       started: false,
       options: { query, variables },
       handler: async (data) => {
@@ -250,7 +261,16 @@ class SubscriptionClient {
           payload: data
         })
       }
-    })
+    }
+
+    if (connectionInit) {
+      operation.extensions = [{
+        type: 'connectionInit',
+        payload: connectionInit
+      }]
+    }
+
+    this.operations.set(operationId, operation)
     this.startOperation(operationId)
     this.operationsCount[operationId] = 1
 
diff --git a/lib/subscription-connection.js b/lib/subscription-connection.js
index a24c9ea..2202d3a 100644
--- a/lib/subscription-connection.js
+++ b/lib/subscription-connection.js
@@ -23,7 +23,7 @@ module.exports = class SubscriptionConnection {
     fastify,
     lruGatewayResolvers,
     entityResolvers,
-    context,
+    context = {},
     onConnect,
     resolveContext
   }) {
@@ -94,7 +94,10 @@ module.exports = class SubscriptionConnection {
   async handleConnectionInit (data) {
     if (typeof this.resolveContext === 'function') {
       try {
-        this.context = await this.resolveContext()
+        this.context = {
+          ...this.context,
+          ...(await this.resolveContext())
+        }
       } catch (err) {
         this.fastify.log.error(err)
         // 1011: Internal Error
@@ -124,15 +127,38 @@ module.exports = class SubscriptionConnection {
       }
     }
 
+    this.context._connectionInit = data.payload
+
     this.sendMessage(GQL_CONNECTION_ACK)
     this.isReady = true
   }
 
   async handleGQLStart (data) {
     // Starting a GraphQL subscription
-    const { id, payload } = data
+    const { id, payload, extensions } = data
     const { query, variables, operationName } = payload
 
+    let context = { ...this.context }
+
+    if (extensions && extensions instanceof Array) {
+      for (const extension of extensions) {
+        switch (extension.type) {
+          case 'connectionInit': {
+            const authorize = await this.handleConnectionInitExtension(extension)
+            if (typeof authorize === 'object') {
+              context = {
+                ...context,
+                ...authorize
+              }
+            }
+            break
+          }
+          default:
+            throw new Error(`Unknown extension ${extension.type}`)
+        }
+      }
+    }
+
     const sc = new SubscriptionContext({
       fastify: this.fastify,
       pubsub: this.subscriber
@@ -146,7 +172,7 @@ module.exports = class SubscriptionConnection {
       document,
       {}, // rootValue
       {
-        ...this.context,
+        ...context,
         pubsub: sc,
         lruGatewayResolvers: this.lruGatewayResolvers,
         reply: {
@@ -201,6 +227,26 @@ module.exports = class SubscriptionConnection {
     }
   }
 
+  async handleConnectionInitExtension (extension) {
+    if (typeof this.onConnect === 'function') {
+      let authorize = false
+      try {
+        authorize = await this.onConnect({
+          type: 'connection_init',
+          payload: extension.payload
+        })
+      } catch (e) {
+        this.fastify.log.error(e)
+      }
+      if (!authorize) {
+        throw new Error('Forbidden')
+      }
+      return authorize
+    }
+
+    return true
+  }
+
   close () {
     this.handleConnectionClose()
   }


Test patch:
diff --git a/tap-snapshots/test-routes.js-TAP.test.js b/tap-snapshots/test-routes.js-TAP.test.js
index dea5d8c..f943bbc 100644
--- a/tap-snapshots/test-routes.js-TAP.test.js
+++ b/tap-snapshots/test-routes.js-TAP.test.js
@@ -70,11 +70,22 @@ exports['test/routes.js TAP if ide is graphiql, serve config.js with the correct
 window.GRAPHQL_ENDPOINT = '/app/graphql'
 `
 
+exports['test/routes.js TAP if ide is playground, and playgroundSettings is set, serve init.js with playground editor options > must match snapshot 1'] = `
+window.addEventListener('load', function(event) {
+          GraphQLPlayground.init(document.getElementById('root'), {
+            subscriptionEndpoint: '/graphql',
+            endpoint: '/graphql',
+            settings: {"editor.theme":"light","editor.fontSize":17},
+          });
+        });
+`
+
 exports['test/routes.js TAP if ide is playground, serve init.js with the correct endpoint > must match snapshot 1'] = `
 window.addEventListener('load', function(event) {
           GraphQLPlayground.init(document.getElementById('root'), {
             subscriptionEndpoint: '/app/graphql',
             endpoint: '/app/graphql',
+            settings: undefined,
           });
         });
 `
diff --git a/test/gateway/connectionInit-extension.js b/test/gateway/connectionInit-extension.js
new file mode 100644
index 0000000..c9e28b9
--- /dev/null
+++ b/test/gateway/connectionInit-extension.js
@@ -0,0 +1,362 @@
+'use strict'
+
+const { test } = require('tap')
+const Fastify = require('fastify')
+const WebSocket = require('ws')
+const GQL = require('../..')
+
+test('connectionInit extension e2e testing', t => {
+  t.plan(9)
+  const userContext = { name: 'test-user' }
+
+  function onConnect (data) {
+    const { payload } = data
+    if (typeof payload.headers === 'object') {
+      // 3 different 'dummy' methods for authentication
+      if (payload.headers.from === 'gateway' && payload.headers.allowGateway) {
+        return { user: userContext }
+      } else if (payload.headers.authorize) {
+        return { user: userContext }
+      } else if (payload.headers.foo) {
+        return { user: userContext }
+      }
+    }
+
+    return false
+  }
+
+  const userService = Fastify()
+  const notificationService = Fastify()
+  const gateway = Fastify()
+  t.tearDown(() => {
+    userService.close()
+    notificationService.close()
+    gateway.close()
+  })
+
+  let userId = 1
+  const users = [
+    {
+      id: 1,
+      name: 'toto'
+    }
+  ]
+  let notificationId = 1
+  const notifications = [
+    {
+      id: 1,
+      message: 'test!'
+    }
+  ]
+
+  userService.register(GQL, {
+    schema: `
+      type User {
+        id: ID!
+        name: String
+      }
+
+      extend type Query {
+        users: [User]
+      }
+
+      extend type Mutation {
+        addUser(name: String!): User
+      }
+
+      extend type Subscription {
+        userAdded: User
+      }
+    `,
+    resolvers: {
+      Query: {
+        users: () => users
+      },
+      Mutation: {
+        addUser: (root, { name }, { pubsub }) => {
+          const user = {
+            id: ++userId,
+            name
+          }
+          pubsub.publish({
+            topic: 'USER_ADDED',
+            payload: {
+              userAdded: user
+            }
+          })
+        }
+      },
+      Subscription: {
+        userAdded: {
+          subscribe: (root, args, { pubsub, user }) => {
+            t.equal(user, userContext)
+            return pubsub.subscribe('USER_ADDED')
+          }
+        }
+      }
+    },
+    federationMetadata: true,
+    subscription: { onConnect }
+  })
+
+  notificationService.register(GQL, {
+    schema: `
+      type Notification {
+        id: ID!
+        message: String
+      }
+
+      extend type Query {
+        notifications: [Notification]
+      }
+
+      extend type Mutation {
+        addNotification(message: String!): Notification
+      }
+
+      extend type Subscription {
+        notificationAdded: Notification
+      }
+    `,
+    resolvers: {
+      Query: {
+        notifications: () => notifications
+      },
+      Mutation: {
+        addNotification: (root, { message }, { pubsub }) => {
+          const notification = {
+            id: ++notificationId,
+            message
+          }
+          pubsub.publish({
+            topic: 'NOTIFICATION_ADDED',
+            payload: {
+              notificationAdded: notification
+            }
+          })
+        }
+      },
+      Subscription: {
+        notificationAdded: {
+          subscribe: (root, args, { pubsub, user }) => {
+            t.equal(user, userContext)
+            return pubsub.subscribe('NOTIFICATION_ADDED')
+          }
+        }
+      }
+    },
+    federationMetadata: true,
+    subscription: { onConnect }
+  })
+
+  Promise.all([
+    userService.ready(),
+    notificationService.ready(),
+    userService.listen(0),
+    notificationService.listen(0)
+  ]).then(() => {
+    gateway.register(GQL, {
+      subscription: true,
+      gateway: {
+        services: [
+          {
+            name: 'user',
+            url: `http://localhost:${(userService.server.address()).port}/graphql`,
+            wsUrl: `ws://localhost:${(userService.server.address()).port}/graphql`,
+            wsConnectionParams: {
+              connectionInitPayload () {
+                return {
+                  headers: {
+                    allowGateway: true,
+                    from: 'gateway'
+                  }
+                }
+              }
+            }
+          },
+          {
+            name: 'notification',
+            url: `http://localhost:${(notificationService.server.address()).port}/graphql`,
+            wsUrl: `ws://localhost:${(notificationService.server.address()).port}/graphql`,
+            wsConnectionParams: {
+              connectionInitPayload () {
+                return {
+                  headers: {
+                    allowGateway: true,
+                    from: 'gateway'
+                  }
+                }
+              }
+            }
+          }
+        ]
+      }
+    })
+
+    gateway.listen(0, err => {
+      t.error(err)
+      async function addUser () {
+        await gateway.inject({
+          method: 'POST',
+          url: '/graphql',
+          body: {
+            query: `
+              mutation {
+                addUser(name: "titi") {
+                  id
+                }
+              }
+            `
+          }
+        })
+      }
+
+      async function addNotification () {
+        await gateway.inject({
+          method: 'POST',
+          url: '/graphql',
+          body: {
+            query: `
+              mutation {
+                addNotification(message: "test") {
+                  id
+                }
+              }
+            `
+          }
+        })
+      }
+
+      const ws1 = new WebSocket(`ws://localhost:${(gateway.server.address()).port}/graphql`, 'graphql-ws')
+      const ws2 = new WebSocket(`ws://localhost:${(gateway.server.address()).port}/graphql`, 'graphql-ws')
+      const client1 = WebSocket.createWebSocketStream(ws1, { encoding: 'utf8', objectMode: true })
+      const client2 = WebSocket.createWebSocketStream(ws2, { encoding: 'utf8', objectMode: true })
+      client1.setEncoding('utf8')
+      client2.setEncoding('utf8')
+      t.tearDown(() => {
+        client1.destroy()
+        client2.destroy()
+      })
+
+      let ready1 = false
+      let ready2 = false
+      let done1 = false
+      let done2 = false
+      async function connectionAckCallback (clientNb) {
+        if (clientNb === 1) {
+          ready1 = true
+        } else if (clientNb === 2) {
+          ready2 = true
+        }
+
+        if (ready1 && ready2) {
+          await Promise.all([
+            addUser(),
+            addNotification()
+          ])
+        }
+      }
+      function terminateCallback (clientNb) {
+        if (clientNb === 1) {
+          done1 = true
+        } else if (clientNb === 2) {
+          done2 = true
+        }
+
+        if (done1 && done2) {
+          t.end()
+        }
+      }
+
+      function makeDataHandler (clientNb) {
+        return async function dataHandler (chunk) {
+          const data = JSON.parse(chunk)
+          switch (data.type) {
+            case 'connection_ack':
+              t.pass('should ack')
+              await connectionAckCallback(clientNb)
+              break
+            case 'data':
+              switch (data.id) {
+                case 1:
+                  t.deepEqual(data.payload.data, { notificationAdded: { id: '2', message: 'test' } })
+                  terminateCallback(clientNb)
+                  break
+                case 2:
+                  t.deepEqual(data.payload.data, { userAdded: { id: '2', name: 'titi' } })
+                  break
+              }
+              break
+            case 'connection_error':
+              t.fail('should not send connection_error')
+              break
+            case 'error':
+              t.fail('should not send error')
+              break
+            case 'complete':
+              t.fail('should not send complete')
+              break
+            default:
+              t.fail('unknown response')
+          }
+        }
+      }
+
+      client1.write(JSON.stringify({
+        type: 'connection_init',
+        payload: {
+          headers: {
+            authorize: true
+          }
+        }
+      }))
+      client2.write(JSON.stringify({
+        type: 'connection_init',
+        payload: {
+          headers: {
+            foo: 'bar'
+          }
+        }
+      }))
+
+      const messages = [
+        {
+          id: 1,
+          type: 'start',
+          payload: {
+            query: `
+            subscription {
+              notificationAdded {
+                id
+                message
+              }
+            }
+          `
+          }
+        },
+        {
+          id: 2,
+          type: 'start',
+          payload: {
+            query: `
+            subscription {
+              userAdded {
+                id
+                name
+              }
+            }
+          `
+          }
+        }
+      ]
+
+      for (const message of messages) {
+        client1.write(JSON.stringify(message))
+        client2.write(JSON.stringify(message))
+      }
+
+      client1.on('data', makeDataHandler(1))
+      client2.on('data', makeDataHandler(2))
+    })
+  })
+})
diff --git a/test/gateway/subscription.js b/test/gateway/subscription.js
index cc4da55..0371601 100644
--- a/test/gateway/subscription.js
+++ b/test/gateway/subscription.js
@@ -384,3 +384,108 @@ test('gateway wsConnectionParams function is passed to SubscriptionClient', t =>
     await gateway.ready()
   })
 })
+
+test('gateway forwards the connectin_init payload to the federated service on gql_start using the connectionInit extension', t => {
+  t.plan(3)
+  function onConnect (data) {
+    if (data && data.payload && Object.entries(data.payload).length) {
+      t.deepEqual(data.payload, connectionInitPayload)
+    }
+
+    return true
+  }
+
+  const connectionInitPayload = {
+    hello: 'world'
+  }
+  const testService = Fastify()
+  t.tearDown(() => {
+    testService.close()
+  })
+
+  testService.register(GQL, {
+    schema: `
+      type Notification {
+        id: ID!
+        message: String
+      }
+
+      type Query {
+        notifications: [Notification]
+      }
+
+      type Subscription {
+        notificationAdded: Notification
+      }
+    `,
+    resolvers: {
+      Query: {
+        notifications: () => []
+      },
+      Subscription: {
+        notificationAdded: {
+          subscribe: (root, args, { pubsub, topic, hello }) => {
+            t.end()
+          }
+        }
+      }
+    },
+    federationMetadata: true,
+    subscription: { onConnect }
+  })
+
+  testService.listen(0, async err => {
+    t.error(err)
+
+    const testServicePort = testService.server.address().port
+
+    const gateway = Fastify()
+    t.tearDown(() => { gateway.close() })
+    gateway.register(GQL, {
+      subscription: true,
+      gateway: {
+        services: [{
+          name: 'test',
+          url: `http://localhost:${testServicePort}/graphql`,
+          wsUrl: `ws://localhost:${testServicePort}/graphql`
+        }]
+      }
+    })
+
+    gateway.listen(0, async err => {
+      t.error(err)
+      const ws = new WebSocket(`ws://localhost:${(gateway.server.address()).port}/graphql`, 'graphql-ws')
+      const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8', objectMode: true })
+      t.tearDown(() => {
+        client.destroy()
+      })
+      client.setEncoding('utf8')
+
+      client.write(JSON.stringify({
+        type: 'connection_init',
+        payload: connectionInitPayload
+      }))
+
+      client.on('data', chunk => {
+        const data = JSON.parse(chunk)
+        if (data.type === 'connection_ack') {
+          client.write(JSON.stringify({
+            id: 1,
+            type: 'start',
+            payload: {
+              query: `
+                subscription {
+                  notificationAdded {
+                    id
+                    message
+                  }
+                }
+              `
+            }
+          }))
+          client.destroy()
+        }
+      })
+    })
+  })
+})
diff --git a/test/routes.js b/test/routes.js
index f23eef9..81b6bd6 100644
--- a/test/routes.js
+++ b/test/routes.js
@@ -1733,6 +1733,31 @@ test('if ide is graphiql, serve config.js with the correct endpoint', async (t)
   t.matchSnapshot(res.body)
 })
 
+test('if ide is playground, and playgroundSettings is set, serve init.js with playground editor options ', async (t) => {
+  const app = Fastify()
+  const schema = `
+    type Query {
+      add(x: Int, y: Int): Int
+    }
+  `
+  app.register(GQL, {
+    ide: 'playground',
+    playgroundSettings: {
+      'editor.theme': 'light',
+      'editor.fontSize': 17
+    },
+    schema
+  })
+
+  const res = await app.inject({
+    method: 'GET',
+    url: '/playground/init.js'
+  })
+  t.strictEqual(res.statusCode, 200)
+  t.strictEqual(res.headers['content-type'], 'application/javascript')
+  t.matchSnapshot(res.body)
+})
+
 test('if operationName is null, it should work fine', async (t) => {
   const app = Fastify()
   const schema = `
diff --git a/test/subscription-connection.js b/test/subscription-connection.js
index e84d423..86b0df2 100644
--- a/test/subscription-connection.js
+++ b/test/subscription-connection.js
@@ -261,7 +261,135 @@ test('subscription connection extends context with onConnect return value', asyn
     }
   })
 
-  await sc.handleConnectionInit()
+  await sc.handleConnectionInit({})
   t.is(sc.context.data, true)
   t.is(sc.context.a, 1)
 })
+
+test('subscription connection send GQL_ERROR message if connectionInit extension is defined and onConnect returns a falsy value', async (t) => {
+  const sc = new SubscriptionConnection({
+    on () { },
+    close () { },
+    send (message) {
+      t.deepEquals(JSON.parse(message), {
+        id: 1,
+        type: 'error',
+        payload: 'Forbidden'
+      })
+    }
+  }, {
+    onConnect: function () {
+      throw new Error('Not allowed')
+    },
+    fastify: {
+      log: {
+        error: () => {}
+      }
+    }
+  })
+
+  sc.isReady = true
+  await sc.handleMessage(JSON.stringify({
+    id: 1,
+    type: 'start',
+    payload: {},
+    extensions: [
+      { type: 'connectionInit' }
+    ]
+  }))
+})
+
+test('subscription connection does not create subscription if connectionInit extension is defined and onConnect returns a falsy value', async (t) => {
+  const sc = new SubscriptionConnection({
+    on () { },
+    close () { },
+    send (message) {}
+  }, {
+    onConnect: function () {
+      throw new Error('Not allowed')
+    },
+    fastify: {
+      log: {
+        error: () => {}
+      }
+    }
+  })
+
+  sc.isReady = true
+  await sc.handleMessage(JSON.stringify({
+    id: 1,
+    type: 'start',
+    payload: {},
+    extensions: [
+      { type: 'connectionInit' }
+    ]
+  }))
+
+  t.notOk(sc.subscriptionContexts.get(0))
+})
+
+test('subscription connection send GQL_ERROR on unknown extension', async (t) => {
+  const sc = new SubscriptionConnection({
+    on () { },
+    close () { },
+    send (message) {
+      t.deepEquals(JSON.parse(message), {
+        id: 1,
+        type: 'error',
+        payload: 'Unknown extension unknown'
+      })
+    }
+  }, { })
+
+  sc.isReady = true
+  await sc.handleMessage(JSON.stringify({
+    id: 1,
+    type: 'start',
+    payload: {},
+    extensions: [
+      { type: 'unknown' }
+    ]
+  }))
+
+  t.notOk(sc.subscriptionContexts.get(0))
+})
+
+test('subscription connection handleConnectionInitExtension returns the onConnect return value', async (t) => {
+  const onConnectResult = {
+    hello: 'world'
+  }
+  const sc = new SubscriptionConnection({
+    on () { },
+    close () { },
+    send (message) { }
+  }, {
+    onConnect: function () {
+      return onConnectResult
+    },
+    fastify: {
+      log: {
+        error: () => { }
+      }
+    }
+  })
+
+  sc.isReady = true
+  const res = await sc.handleConnectionInitExtension({ type: 'connectionInit' })
+
+  t.deepEqual(res, onConnectResult)
+})
+
+test('subscription connection externds the context with the connection_init payload', async (t) => {
+  const connectionInitPayload = {
+    hello: 'world'
+  }
+  const sc = new SubscriptionConnection({
+    on () { },
+    close () { },
+    send (message) { }
+  }, {})
+
+  await sc.handleConnectionInit({ type: 'connection_init', payload: connectionInitPayload })
+
+  t.deepEqual(sc.context._connectionInit, connectionInitPayload)
+})
diff --git a/test/subscription.js b/test/subscription.js
index f6743a1..c3be6f7 100644
--- a/test/subscription.js
+++ b/test/subscription.js
@@ -936,3 +936,164 @@ test('subscription server exposes pubsub', t => {
     })
   })
 })
+
+test('subscription context is extended with onConnect return value if connectionInit extension is defined in gql_start message', t => {
+  const app = Fastify()
+  t.tearDown(() => app.close())
+
+  const schema = `
+    type Notification {
+      id: ID!
+      message: String
+    }
+
+    type Query {
+      notifications: [Notification]
+    }
+
+    type Subscription {
+      notificationAdded: Notification
+    }
+  `
+
+  const resolvers = {
+    Query: {
+      notifications: () => []
+    },
+    Subscription: {
+      notificationAdded: {
+        subscribe: (root, args, { pubsub, topic, hello }) => {
+          t.equal(hello, 'world')
+          t.end()
+          pubsub.subscribe(topic)
+        }
+      }
+    }
+  }
+
+  app.register(GQL, {
+    schema,
+    resolvers,
+    subscription: {
+      onConnect: () => ({ hello: 'world' })
+    }
+  })
+
+  app.listen(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'
+    }))
+
+    client.on('data', chunk => {
+      const data = JSON.parse(chunk)
+
+      if (data.type === 'connection_ack') {
+        client.write(JSON.stringify({
+          id: 1,
+          type: 'start',
+          payload: {
+            query: `
+          subscription {
+            notificationAdded {
+              id
+              message
+            }
+          }
+        `
+          },
+          extensions: [
+            { type: 'connectionInit' }
+          ]
+        }))
+
+        client.end()
+      }
+    })
+  })
+})
+
+test('subscription works properly if onConnect is not defined and connectionInit extension is defined in gql_start message', t => {
+  const app = Fastify()
+  t.tearDown(() => app.close())
+
+  const schema = `
+    type Notification {
+      id: ID!
+      message: String
+    }
+
+    type Query {
+      notifications: [Notification]
+    }
+
+    type Subscription {
+      notificationAdded: Notification
+    }
+  `
+
+  const resolvers = {
+    Query: {
+      notifications: () => []
+    },
+    Subscription: {
+      notificationAdded: {
+        subscribe: (root, args, { pubsub, topic, hello }) => {
+          t.end()
+          pubsub.subscribe(topic)
+        }
+      }
+    }
+  }
+
+  app.register(GQL, {
+    schema,
+    resolvers,
+    subscription: true
+  })
+
+  app.listen(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'
+    }))
+
+    client.on('data', chunk => {
+      const data = JSON.parse(chunk)
+
+      if (data.type === 'connection_ack') {
+        client.write(JSON.stringify({
+          id: 1,
+          type: 'start',
+          payload: {
+            query: `
+          subscription {
+            notificationAdded {
+              id
+              message
+            }
+          }
+        `
+          },
+          extensions: [
+            { type: 'connectionInit' }
+          ]
+        }))
+
+        client.end()
+      }
+    })
+  })
+})