fix: duplicate subscription ids now properly error & close the connection (#63)

Author: reconbot

README.md

@@ -599,7 +599,6 @@ const instance = makeServer({
 
 </details>
 
-
 ## Caveats
 
 ### Ping/Pong
@@ -609,3 +608,7 @@ For whatever reason, AWS API Gateway does not support WebSocket protocol level p
 ### Socket idleness
 
 API Gateway considers an idle connection to be one where no messages have been sent on the socket for a fixed duration [(currently 10 minutes)](https://docs.aws.amazon.com/apigateway/latest/developerguide/limits.html#apigateway-execution-service-websocket-limits-table). The WebSocket spec has support for detecting idle connections (ping/pong) but API Gateway doesn't use it. This means, in the case where both parties are connected, and no message is sent on the socket for the defined duration (direction agnostic), API Gateway will close the socket. A fix for this is to set up immediate reconnection on the client side.
+
+### Socket Close Reasons
+
+API Gateway doesn't support custom reasons or codes for WebSockets being closed. So the codes and reason strings wont match `graphql-ws`.

lib/ddb/DDB.ts

@@ -2,8 +2,8 @@ import { DynamoDB } from 'aws-sdk'
 import { LoggerFunction, DDBType } from '../types'
 
 export interface DDBClient<T extends DDBType, TKey> {
-  get: (Key: TKey) => Promise<T|null>
-  put: (obj: T) => Promise<T>
+  get: (Key: TKey) => Promise<T | null>
+  put: (obj: T, putOptions?: Partial<DynamoDB.DocumentClient.PutItemInput>) => Promise<T>
   update: (Key: TKey, obj: Partial<T>) => Promise<T>
   delete: (Key: TKey) => Promise<T>
   query: (options: Omit<DynamoDB.DocumentClient.QueryInput, 'TableName' | 'Select'>) => AsyncGenerator<T, void, undefined>
@@ -35,13 +35,14 @@ export const DDB = <T extends DDBType, TKey>({
     }
   }
 
-  const put = async (Item: T): Promise<T> => {
+  const put = async (Item: T, putOptions?: Partial<DynamoDB.DocumentClient.PutItemInput>): Promise<T> => {
     log('put', { tableName: tableName, Item })
     try {
       const { Attributes } = await documentClient.put({
         TableName: tableName,
         Item,
         ReturnValues: 'ALL_OLD',
+        ...putOptions,
       }).promise()
       return Attributes as T
     } catch (e) {

lib/messages/subscribe.ts

@@ -18,14 +18,15 @@ export const subscribe: MessageHandler<SubscribeMessage> =
     try {
       await setupSubscription({ server, event, message })
     } catch (err) {
+      server.log('subscribe:error', { connectionId: event.requestContext.connectionId, error: err.message })
       await server.onError?.(err, { event, message })
       await deleteConnection(server)(event.requestContext)
     }
   }
 
 const setupSubscription: MessageHandler<SubscribeMessage> = async ({ server, event, message }) => {
   const connectionId = event.requestContext.connectionId
-  server.log('subscribe', { connectionId, query: message.payload.query })
+  server.log('subscribe', { connectionId, messageId: message.id, query: message.payload.query })
 
   const connection = await server.models.connection.get({ id: connectionId })
   if (!connection) {
@@ -83,7 +84,7 @@ const setupSubscription: MessageHandler<SubscribeMessage> = async ({ server, eve
 
   server.log('onSubscribe', { onSubscribe: !!onSubscribe })
   const onSubscribeErrors = await onSubscribe?.(root, args, context, info)
-  if (onSubscribeErrors){
+  if (onSubscribeErrors) {
     server.log('onSubscribe', { onSubscribeErrors })
     return postToConnection(server)({
       ...event.requestContext,
@@ -97,8 +98,9 @@ const setupSubscription: MessageHandler<SubscribeMessage> = async ({ server, eve
 
   const filterData = typeof filter === 'function' ? await filter(root, args, context, info) : filter
 
+  const subscriptionId = `${connection.id}|${message.id}`
   const subscription: Subscription = {
-    id: `${connection.id}|${message.id}`,
+    id: subscriptionId,
     topic,
     filter: filterData || {},
     subscriptionId: message.id,
@@ -113,7 +115,22 @@ const setupSubscription: MessageHandler<SubscribeMessage> = async ({ server, eve
     createdAt: Date.now(),
   }
   server.log('subscribe:putSubscription', subscription)
-  await server.models.subscription.put(subscription)
+  try{
+    await server.models.subscription.put(subscription, {
+      ConditionExpression: '#id <> :id',
+      ExpressionAttributeNames: {
+        '#id': 'id',
+      },
+      ExpressionAttributeValues: {
+        ':id': subscriptionId,
+      },
+    })
+  } catch (error) {
+    if (error.code === 'ConditionalCheckFailedException') {
+      throw new Error(`Subscriber for ${message.id} already exists`)
+    }
+    throw error
+  }
 
   server.log('onAfterSubscribe', { onAfterSubscribe: !!onAfterSubscribe })
   await onAfterSubscribe?.(root, args, context, info)

lib/test/execute-helper.ts

@@ -16,12 +16,13 @@ export const executeQuery = async function* (query: string, {
   url = URL,
   stayConnected = false,
   timeout = 20_000,
+  id = 1,
 }: {
   url?: string
   stayConnected?: boolean
   timeout?: number
+  id?: number
 } = {}): AsyncGenerator<unknown, void, unknown> {
-  let id = 1
   const ws = new WebSocket(url, 'graphql-transport-ws')
 
   const incomingMessages = deferGenerator()
@@ -61,7 +62,7 @@ export const executeQuery = async function* (query: string, {
   }
 
   await send({
-    id: `${id++}`,
+    id: `${id}`,
     type: 'subscribe',
     payload: { query },
   })
@@ -84,10 +85,11 @@ export const executeQuery = async function* (query: string, {
 
 export const executeToComplete = async function (query: string, {
   url = URL,
+  id = 1,
 }: {
   url?: string
+  id?: number
 } = {}): Promise<() => Promise<void>> {
-  let id = 1
   const ws = new WebSocket(url, 'graphql-transport-ws')
 
   const incomingMessages = deferGenerator()
@@ -118,16 +120,14 @@ export const executeToComplete = async function (query: string, {
     throw new Error(`Bad ack ${messageToString(connectionAck)}`)
   }
 
-  const subId = id++
-
   await send({
-    id: `${subId}`,
+    id: `${id}`,
     type: 'subscribe',
     payload: { query },
   })
 
   return () => send({
-    id: `${subId}`,
+    id,
     type: 'complete',
   })
 }
@@ -136,10 +136,11 @@ export const executeToComplete = async function (query: string, {
 
 export const executeToDisconnect = async function (query: string, {
   url = URL,
+  id = 1,
 }: {
   url?: string
+  id?: number
 } = {}): Promise<() => void> {
-  let id = 1
   const ws = new WebSocket(url, 'graphql-transport-ws')
 
   const incomingMessages = deferGenerator()
@@ -170,13 +171,91 @@ export const executeToDisconnect = async function (query: string, {
     throw new Error(`Bad ack ${messageToString(connectionAck)}`)
   }
 
-  const subId = id++
-
   await send({
-    id: `${subId}`,
+    id: `${id}`,
     type: 'subscribe',
     payload: { query },
   })
 
   return () => ws.close()
 }
+
+export const executeDoubleQuery = async function* (query: string, {
+  url = URL,
+  stayConnected = false,
+  timeout = 20_000,
+  id = 1,
+}: {
+  url?: string
+  stayConnected?: boolean
+  timeout?: number
+  id?: number
+} = {}): AsyncGenerator<unknown, void, unknown> {
+  const ws = new WebSocket(url, 'graphql-transport-ws')
+
+  const incomingMessages = deferGenerator()
+
+  ws.on('message', data => {
+    const message = JSON.parse(data.toString())
+    incomingMessages.queueValue(message)
+    if (message.type === 'error' || message.type === 'complete') {
+      incomingMessages.queueReturn()
+    }
+  })
+
+  ws.on('error', error => {
+    incomingMessages.queueValue( { type: 'websocketError', value: error.message })
+    incomingMessages.queueReturn()
+  })
+  ws.on('close', (code, reason) => {
+    incomingMessages.queueValue({ type: 'close', code, reason: reason.toString() })
+    incomingMessages.queueReturn()
+  })
+
+  let timer: NodeJS.Timeout|null = null
+  if (timeout) {
+    timer = setTimeout(() => {
+      incomingMessages.queueValue({ type: 'timeout', timeout })
+      incomingMessages.queueReturn()
+    }, timeout)
+  }
+
+  const send = (data: any) => new Promise<void>(resolve => ws.send(JSON.stringify(data), () => resolve()))
+
+  await new Promise(resolve => ws.on('open', resolve))
+  await send({ type: 'connection_init' })
+  const connectionAck: any = (await incomingMessages.generator.next()).value
+  if (connectionAck.type !== 'connection_ack') {
+    throw new Error(`Bad ack ${messageToString(connectionAck)}`)
+  }
+
+  await send({
+    id: `${id}`,
+    type: 'subscribe',
+    payload: { query },
+  })
+
+  const firstMessage = await incomingMessages.generator.next()
+  if (firstMessage.done) {
+    return
+  }
+  yield firstMessage.value
+
+  await send({
+    id: `${id}`,
+    type: 'subscribe',
+    payload: { query },
+  })
+
+
+  for await (const message of incomingMessages.generator) {
+    yield message
+  }
+
+  if (!stayConnected){
+    ws.close()
+  }
+  if (timer) {
+    clearTimeout(timer)
+  }
+}

lib/test/graphql-ws-schema.ts

@@ -14,6 +14,7 @@ const typeDefs = `
     greetings: String
     onSubscribeError: String
     onResolveError: String
+    oneEvent: String
   }
 `
 
@@ -43,6 +44,14 @@ const resolvers = {
         throw new Error('resolver error')
       },
     },
+    oneEvent:{
+      subscribe: async function*(){
+        yield { oneEvent:  'lets start!' }
+        // eslint-disable-next-line @typescript-eslint/no-empty-function
+        await new Promise(() => {})
+      },
+    },
+
   },
 }
 

lib/test/integration-events-test.ts

@@ -2,7 +2,7 @@
 import { assert } from 'chai'
 import { start as sandBoxStart, end as sandBoxStop } from '@architect/sandbox'
 import { collect, map } from 'streaming-iterables'
-import { executeQuery, executeToComplete, executeToDisconnect } from './execute-helper'
+import { executeDoubleQuery, executeQuery, executeToComplete, executeToDisconnect } from './execute-helper'
 import { startGqlWSServer } from './graphql-ws-schema'
 
 describe('Events', () => {
@@ -64,6 +64,18 @@ describe('Events', () => {
       assert.deepEqual(lambdaResult, gqlWSResult)
       await stop()
     })
+
+    it('errors when duplicating subscription ids', async () => {
+      const { url, stop } = await startGqlWSServer()
+
+      const lambdaError = await collect(executeDoubleQuery('subscription { oneEvent }', { id: 1 }))
+      const gqlWSError = await collect(executeDoubleQuery('subscription { oneEvent }', { url, id: 1 }))
+      assert.deepEqual(lambdaError[0], gqlWSError[0])
+      // This would be exactly equal but apigateway doesn't support close reasons *eye roll*
+      assert.containSubset(lambdaError[1], { type: 'close' })
+      assert.containSubset(gqlWSError[1], { type: 'close' })
+      await stop()
+    })
   })
 
   describe('Filter Events', () => {

mocks/arc-basic-events/lib/graphql.js

@@ -34,6 +34,7 @@ const typeDefs = `
     onCompleteTestClientDisconnect: String
     onResolveError: String
     onCompleteServerComplete: String
+    oneEvent: String
   }
 `
 
@@ -139,6 +140,17 @@ const resolvers = {
         return payload.message
       },
     },
+    oneEvent:{
+      subscribe: subscribe('oneEvent', {
+        async onAfterSubscribe(_, __, { publish }) {
+          await publish({ topic: 'oneEvent', payload: { message: 'lets start!' } })
+        },
+      }),
+      resolve: ({ payload }) => {
+        return payload.message
+      },
+    },
+
   },
 }