Websocket links management in graphql gateway (related to https://github.com/apollographql/graphql-tools/issues/864)

server.ts

import { ApolloServer } from 'apollo-server-express';
import { ContextBuilderPayload } from './Types';

const server = new ApolloServer({
	  context: (payload: ContextBuilderPayload) => {
      let context: CollaborationContext = {};
      if (!isIntrospection(payload)) {
        if (isOperation(payload)) {
          context = { /* operation context */};
        } else {
          const wsContext = payload.connection.context;
          const token = wsContext.authorization || wsContext.idtoken;
          context = { idtoken: token };
        }
      }
      return context;
    },
		subscriptions: {
			onConnect: (connectionParams, websocket, context) => {
				console.info(`Connecting new websocket...`);
				return connectionParams;
			},
			onDisconnect: async (websocket, context) => {
				const params = await context.initPromise;
				const { authorization, idtoken } = params;
				console.info(`Disconnecting websocket...`);
				disconnectWebsocket(idtoken || authorization);
			}
		}
	});

const isOperation = (e: ContextBuilderPayload): e is OperationContextBuilderPayload {
	return !!(e as OperationContextBuilderPayload).req;
}

const isIntrospection = (payload: ContextBuilderPayload) => {
	return isOperation(payload) && payload.req.body.operationName === 'IntrospectionQuery';
};

SubscriptionManagement.ts

import { Operation, NextLink } from 'apollo-link';
import { SubscriptionClient } from 'subscriptions-transport-ws';
import WebSocket from 'ws';
import { CollaborationContext } from './Types';

export const linksMap: { [key: string]: { [key: string]: SubscriptionClient } } = {};

export const createCustomWsLink = (wsUri: string) => {
	console.info(`Creating custom WS Link for ${wsUri}`);
	
	return (operation: Operation, forward: NextLink) => {
		const context = operation.getContext();
		const connectionParams: CollaborationContext = context.graphqlContext;
		const { idtoken, user } = connectionParams;
		const userWsLinks = (linksMap[idtoken] = linksMap[idtoken] || {});
		let client = userWsLinks[wsUri];
		if (!client) {
			console.info(`Creating ws link for user ${user.name} & token ${idtoken.substring(0, 10)}`);
			client = new SubscriptionClient(
				wsUri,
				{
					connectionParams,
					reconnect: true
				},
				WebSocket
			);
			userWsLinks[wsUri] = client;
		} else {
			console.debug(`Reusing ws link for user ${user.name} & token ${idtoken.substring(0, 10)}`);
		}
		console.debug(`Executing ${operation.operationName}`);
		return client.request(operation);
	};
};

export const disconnectWebsocket = (idtoken: string) => {
	if (linksMap[idtoken]) {
		info(`Disconnecting all websockets clients with token: ${idtoken.substring(0, 10)}`);
		Object.values(linksMap[idtoken]).forEach(client => client.close());
		delete linksMap[idtoken];
	} else {
		console.debug(`Ws link not found for token ${idtoken}`);	
	}
	console.debug(`linksMap now: ${Object.values(linksMap)}`);
};

Types.ts

export type ContextBuilderPayload =
	| SubscriptionContextBuilderPayload
	| OperationContextBuilderPayload;

export interface SubscriptionContextBuilderPayload {
	payload: {
		extensions: any;
		operationName: string;
		query: string;
		variables: any;
	};
	connection: {
		context: any;
		query: string;
		variables: any;
	};
}
export interface OperationContextBuilderPayload {
	req: Request;
	res: Response;
}