aws-appsync-subscription-link 2.2.5 → 2.3.0

package/CHANGELOG.md

@@ -3,6 +3,40 @@
 All notable changes to this project will be documented in this file.
 See [Conventional Commits](https://conventionalcommits.org) for commit guidelines.
 
+<a name="2.3.0"></a>
+# [2.3.0](https://github.com/awslabs/aws-mobile-appsync-sdk-js/compare/aws-appsync-subscription-link@2.2.7...aws-appsync-subscription-link@2.3.0) (2022-05-02)
+
+
+### Features
+
+* add custom domain support for Apollo V2 ([#698](https://github.com/awslabs/aws-mobile-appsync-sdk-js/issues/698)) ([42db16d](https://github.com/awslabs/aws-mobile-appsync-sdk-js/commit/42db16d))
+
+
+
+
+<a name="2.2.7"></a>
+## [2.2.7](https://github.com/awslabs/aws-mobile-appsync-sdk-js/compare/aws-appsync-subscription-link@2.2.5...aws-appsync-subscription-link@2.2.7) (2022-03-04)
+
+
+### Bug Fixes
+
+* Port over Amplify fix for subscription race conditions for ApolloV2 ([#509](https://github.com/awslabs/aws-mobile-appsync-sdk-js/issues/509)) ([#705](https://github.com/awslabs/aws-mobile-appsync-sdk-js/issues/705)) ([f5022ba](https://github.com/awslabs/aws-mobile-appsync-sdk-js/commit/f5022ba))
+* **aws-appsync-subscription-link:** graphql header refactor to fix IAM-based auth for Apollo V2 packages ([#679](https://github.com/awslabs/aws-mobile-appsync-sdk-js/issues/679)) ([09fc430](https://github.com/awslabs/aws-mobile-appsync-sdk-js/commit/09fc430))
+
+
+
+
+<a name="2.2.6"></a>
+## [2.2.6](https://github.com/awslabs/aws-mobile-appsync-sdk-js/compare/aws-appsync-subscription-link@2.2.5...aws-appsync-subscription-link@2.2.6) (2021-10-18)
+
+
+### Bug Fixes
+
+* **aws-appsync-subscription-link:** graphql header refactor to fix IAM-based auth for Apollo V2 packages ([#679](https://github.com/awslabs/aws-mobile-appsync-sdk-js/issues/679)) ([09fc430](https://github.com/awslabs/aws-mobile-appsync-sdk-js/commit/09fc430))
+
+
+
+
 <a name="2.2.5"></a>
 ## [2.2.5](https://github.com/awslabs/aws-mobile-appsync-sdk-js/compare/aws-appsync-subscription-link@2.2.4...aws-appsync-subscription-link@2.2.5) (2021-09-24)
 

package/__tests__/link/realtime-subscription-handshake-link-test.ts

@@ -50,8 +50,22 @@ describe("RealTime subscription link", () => {
                 type: AUTH_TYPE.API_KEY,
                 apiKey: 'xxxxx'
             },
-            region: 'us-east-1',
-            url: 'https://xxxxx.appsync-api.amazonaws.com/graphql'
+            region: 'us-west-2',
+            url: 'https://firsttesturl12345678901234.appsync-api.us-west-2.amazonaws.com/graphql'
+        });
+
+        expect(link).toBeInstanceOf(AppSyncRealTimeSubscriptionHandshakeLink);
+    });
+
+    test("Can instantiate link with custom domain", () => {
+        expect.assertions(1);
+        const link = new AppSyncRealTimeSubscriptionHandshakeLink({
+            auth: {
+                type: AUTH_TYPE.API_KEY,
+                apiKey: 'xxxxx'
+            },
+            region: 'us-west-2',
+            url: 'https://test1.testcustomdomain.com/graphql'
         });
 
         expect(link).toBeInstanceOf(AppSyncRealTimeSubscriptionHandshakeLink);
@@ -63,7 +77,7 @@ describe("RealTime subscription link", () => {
             return "2019-11-13T18:47:04.733Z";
         }));
         AppSyncRealTimeSubscriptionHandshakeLink.createWebSocket = jest.fn((url, protocol) => {
-            expect(url).toBe('wss://xxxxx.appsync-realtime-api.amazonaws.com/graphql?header=eyJob3N0IjoieHh4eHguYXBwc3luYy1hcGkuYW1hem9uYXdzLmNvbSIsIngtYW16LWRhdGUiOiIyMDE5MTExM1QxODQ3MDRaIiwieC1hcGkta2V5IjoieHh4eHgifQ==&payload=e30=');
+            expect(url).toBe('wss://apikeytesturl1234567890123.appsync-realtime-api.us-west-2.amazonaws.com/graphql?header=eyJob3N0IjoiYXBpa2V5dGVzdHVybDEyMzQ1Njc4OTAxMjMuYXBwc3luYy1hcGkudXMtd2VzdC0yLmFtYXpvbmF3cy5jb20iLCJ4LWFtei1kYXRlIjoiMjAxOTExMTNUMTg0NzA0WiIsIngtYXBpLWtleSI6Inh4eHh4In0=&payload=e30=');
             expect(protocol).toBe('graphql-ws');
             done();
             return new myWebSocket();
@@ -73,13 +87,50 @@ describe("RealTime subscription link", () => {
                 type: AUTH_TYPE.API_KEY,
                 apiKey: 'xxxxx'
             },
-            region: 'us-east-1',
-            url: 'https://xxxxx.appsync-api.amazonaws.com/graphql'
+            region: 'us-west-2',
+            url: 'https://apikeytesturl1234567890123.appsync-api.us-west-2.amazonaws.com/graphql'
         });
 
         execute(link, { query }).subscribe({
             error: (err) => {
-                console.log({ err });
+                console.log(JSON.stringify(err));
+                fail;
+            },
+            next: (data) => {
+                console.log({ data });
+                done();
+            },
+            complete: () => {
+                console.log('done with this');
+                done();
+            }
+
+        });
+    });
+
+    test("Initialize WebSocket correctly for API KEY with custom domain", (done) => {
+        expect.assertions(2);
+        jest.spyOn(Date.prototype, 'toISOString').mockImplementation(jest.fn(() => {
+            return "2019-11-13T18:47:04.733Z";
+        }));
+        AppSyncRealTimeSubscriptionHandshakeLink.createWebSocket = jest.fn((url, protocol) => {
+            expect(url).toBe('wss://apikeytest.testcustomdomain.com/graphql/realtime?header=eyJob3N0IjoiYXBpa2V5dGVzdC50ZXN0Y3VzdG9tZG9tYWluLmNvbSIsIngtYW16LWRhdGUiOiIyMDE5MTExM1QxODQ3MDRaIiwieC1hcGkta2V5IjoieHh4eHgifQ==&payload=e30=');
+            expect(protocol).toBe('graphql-ws');
+            done();
+            return new myWebSocket();
+        });
+        const link = new AppSyncRealTimeSubscriptionHandshakeLink({
+            auth: {
+                type: AUTH_TYPE.API_KEY,
+                apiKey: 'xxxxx'
+            },
+            region: 'us-west-2',
+            url: 'https://apikeytest.testcustomdomain.com/graphql'
+        });
+
+        execute(link, { query }).subscribe({
+            error: (err) => {
+                console.log(JSON.stringify(err));
                 fail;
             },
             next: (data) => {
@@ -100,7 +151,7 @@ describe("RealTime subscription link", () => {
             return "2019-11-13T18:47:04.733Z";
         }));
         AppSyncRealTimeSubscriptionHandshakeLink.createWebSocket = jest.fn((url, protocol) => {
-            expect(url).toBe('wss://xxxxx.appsync-realtime-api.amazonaws.com/graphql?header=eyJBdXRob3JpemF0aW9uIjoidG9rZW4iLCJob3N0IjoieHh4eHguYXBwc3luYy1hcGkuYW1hem9uYXdzLmNvbSJ9&payload=e30=');
+            expect(url).toBe('wss://cognitouserpooltesturl1234.appsync-realtime-api.us-west-2.amazonaws.com/graphql?header=eyJBdXRob3JpemF0aW9uIjoidG9rZW4iLCJob3N0IjoiY29nbml0b3VzZXJwb29sdGVzdHVybDEyMzQuYXBwc3luYy1hcGkudXMtd2VzdC0yLmFtYXpvbmF3cy5jb20ifQ==&payload=e30=');
             expect(protocol).toBe('graphql-ws');
             done();
             return new myWebSocket();
@@ -110,13 +161,50 @@ describe("RealTime subscription link", () => {
                 type: AUTH_TYPE.AMAZON_COGNITO_USER_POOLS,
                 jwtToken: 'token'
             },
-            region: 'us-east-1',
-            url: 'https://xxxxx.appsync-api.amazonaws.com/graphql'
+            region: 'us-west-2',
+            url: 'https://cognitouserpooltesturl1234.appsync-api.us-west-2.amazonaws.com/graphql'
         });
 
         execute(link, { query }).subscribe({
             error: (err) => {
-                console.log({ err });
+                console.log(JSON.stringify(err));
+                fail;
+            },
+            next: (data) => {
+                console.log({ data });
+                done();
+            },
+            complete: () => {
+                console.log('done with this');
+                done();
+            }
+
+        });
+    });
+
+    test("Initialize WebSocket correctly for COGNITO USER POOLS with custom domain", (done) => {
+        expect.assertions(2);
+        jest.spyOn(Date.prototype, 'toISOString').mockImplementation(jest.fn(() => {
+            return "2019-11-13T18:47:04.733Z";
+        }));
+        AppSyncRealTimeSubscriptionHandshakeLink.createWebSocket = jest.fn((url, protocol) => {
+            expect(url).toBe('wss://cognitouserpools.testcustomdomain.com/graphql/realtime?header=eyJBdXRob3JpemF0aW9uIjoidG9rZW4iLCJob3N0IjoiY29nbml0b3VzZXJwb29scy50ZXN0Y3VzdG9tZG9tYWluLmNvbSJ9&payload=e30=');
+            expect(protocol).toBe('graphql-ws');
+            done();
+            return new myWebSocket();
+        });
+        const link = new AppSyncRealTimeSubscriptionHandshakeLink({
+            auth: {
+                type: AUTH_TYPE.AMAZON_COGNITO_USER_POOLS,
+                jwtToken: 'token'
+            },
+            region: 'us-west-2',
+            url: 'https://cognitouserpools.testcustomdomain.com/graphql'
+        });
+
+        execute(link, { query }).subscribe({
+            error: (err) => {
+                console.log(JSON.stringify(err));
                 fail;
             },
             next: (data) => {
@@ -137,7 +225,44 @@ describe("RealTime subscription link", () => {
             return "2019-11-13T18:47:04.733Z";
         }));
         AppSyncRealTimeSubscriptionHandshakeLink.createWebSocket = jest.fn((url, protocol) => {
-            expect(url).toBe('wss://xxxxx.appsync-realtime-api.amazonaws.com/graphql?header=eyJBdXRob3JpemF0aW9uIjoidG9rZW4iLCJob3N0IjoieHh4eHguYXBwc3luYy1hcGkuYW1hem9uYXdzLmNvbSJ9&payload=e30=');
+            expect(url).toBe('wss://openidconnecttesturl123456.appsync-realtime-api.us-west-2.amazonaws.com/graphql?header=eyJBdXRob3JpemF0aW9uIjoidG9rZW4iLCJob3N0Ijoib3BlbmlkY29ubmVjdHRlc3R1cmwxMjM0NTYuYXBwc3luYy1hcGkudXMtd2VzdC0yLmFtYXpvbmF3cy5jb20ifQ==&payload=e30=');
+            expect(protocol).toBe('graphql-ws');
+            done();
+            return new myWebSocket();
+        });
+        const link = new AppSyncRealTimeSubscriptionHandshakeLink({
+            auth: {
+                type: AUTH_TYPE.OPENID_CONNECT,
+                jwtToken: 'token'
+            },
+            region: 'us-west-2',
+            url: 'https://openidconnecttesturl123456.appsync-api.us-west-2.amazonaws.com/graphql'
+        });
+
+        execute(link, { query }).subscribe({
+            error: (err) => {
+                console.log(JSON.stringify(err));
+                fail;
+            },
+            next: (data) => {
+                console.log({ data });
+                done();
+            },
+            complete: () => {
+                console.log('done with this');
+                done();
+            }
+
+        });
+    });
+
+    test("Initialize WebSocket correctly for OPENID_CONNECT with custom domain", (done) => {
+        expect.assertions(2);
+        jest.spyOn(Date.prototype, 'toISOString').mockImplementation(jest.fn(() => {
+            return "2019-11-13T18:47:04.733Z";
+        }));
+        AppSyncRealTimeSubscriptionHandshakeLink.createWebSocket = jest.fn((url, protocol) => {
+            expect(url).toBe('wss://openidconnecttesturl.testcustomdomain.com/graphql/realtime?header=eyJBdXRob3JpemF0aW9uIjoidG9rZW4iLCJob3N0Ijoib3BlbmlkY29ubmVjdHRlc3R1cmwudGVzdGN1c3RvbWRvbWFpbi5jb20ifQ==&payload=e30=');
             expect(protocol).toBe('graphql-ws');
             done();
             return new myWebSocket();
@@ -147,13 +272,13 @@ describe("RealTime subscription link", () => {
                 type: AUTH_TYPE.OPENID_CONNECT,
                 jwtToken: 'token'
             },
-            region: 'us-east-1',
-            url: 'https://xxxxx.appsync-api.amazonaws.com/graphql'
+            region: 'us-west-2',
+            url: 'https://openidconnecttesturl.testcustomdomain.com/graphql'
         });
 
         execute(link, { query }).subscribe({
             error: (err) => {
-                console.log({ err });
+                console.log(JSON.stringify(err));
                 fail;
             },
             next: (data) => {
@@ -174,7 +299,41 @@ describe("RealTime subscription link", () => {
             return "2019-11-13T18:47:04.733Z";
         }));
         AppSyncRealTimeSubscriptionHandshakeLink.createWebSocket = jest.fn((url, protocol) => {
-            expect(url).toBe('wss://xxxxx.appsync-realtime-api.amazonaws.com/graphql?header=eyJBdXRob3JpemF0aW9uIjoidG9rZW4iLCJob3N0IjoieHh4eHguYXBwc3luYy1hcGkuYW1hem9uYXdzLmNvbSJ9&payload=e30=');
+            expect(url).toBe('wss://awslambdatesturl1234567890.appsync-realtime-api.us-west-2.amazonaws.com/graphql?header=eyJBdXRob3JpemF0aW9uIjoidG9rZW4iLCJob3N0IjoiYXdzbGFtYmRhdGVzdHVybDEyMzQ1Njc4OTAuYXBwc3luYy1hcGkudXMtd2VzdC0yLmFtYXpvbmF3cy5jb20ifQ==&payload=e30=');
+            expect(protocol).toBe('graphql-ws');
+            done();
+            return new myWebSocket();
+        });
+        const link = new AppSyncRealTimeSubscriptionHandshakeLink({
+            auth: {
+                type: AUTH_TYPE.AWS_LAMBDA,
+                token: 'token'
+            },
+            region: 'us-west-2',
+            url: 'https://awslambdatesturl1234567890.appsync-api.us-west-2.amazonaws.com/graphql'
+        });
+
+        execute(link, { query }).subscribe({
+            error: (err) => {
+                fail;
+            },
+            next: (data) => {
+                done();
+            },
+            complete: () => {
+                done();
+            }
+
+        });
+    })
+
+    test('Initialize WebSocket correctly for AWS_LAMBDA with custom domain', (done) => {
+        expect.assertions(2);
+        jest.spyOn(Date.prototype, 'toISOString').mockImplementation(jest.fn(() => {
+            return "2019-11-13T18:47:04.733Z";
+        }));
+        AppSyncRealTimeSubscriptionHandshakeLink.createWebSocket = jest.fn((url, protocol) => {
+            expect(url).toBe('wss://awslambdatesturl.testcustomdomain.com/graphql/realtime?header=eyJBdXRob3JpemF0aW9uIjoidG9rZW4iLCJob3N0IjoiYXdzbGFtYmRhdGVzdHVybC50ZXN0Y3VzdG9tZG9tYWluLmNvbSJ9&payload=e30=');
             expect(protocol).toBe('graphql-ws');
             done();
             return new myWebSocket();
@@ -184,8 +343,8 @@ describe("RealTime subscription link", () => {
                 type: AUTH_TYPE.AWS_LAMBDA,
                 token: 'token'
             },
-            region: 'us-east-1',
-            url: 'https://xxxxx.appsync-api.amazonaws.com/graphql'
+            region: 'us-west-2',
+            url: 'https://awslambdatesturl.testcustomdomain.com/graphql'
         });
 
         execute(link, { query }).subscribe({

package/lib/non-terminating-link.d.ts

@@ -9,7 +9,5 @@ export declare class NonTerminatingLink extends ApolloLink {
     constructor(contextKey: string, { link }: {
         link: ApolloLink;
     });
-    request(operation: any, forward?: NextLink): import("apollo-link").Observable<import("apollo-link").FetchResult<{
-        [key: string]: any;
-    }, Record<string, any>, Record<string, any>>>;
+    request(operation: any, forward?: NextLink): ReturnType<ApolloLink['request']>;
 }

package/lib/realtime-subscription-handshake-link.d.ts

@@ -1,8 +1,8 @@
 /*!
- * Copyright 2017-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * Copyright 2017-2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-import { ApolloLink, Observable, Operation } from "apollo-link";
+import { ApolloLink, Observable, Operation, FetchResult } from "apollo-link";
 import { UrlInfo } from "./types";
 export declare const CONTROL_EVENTS_KEY = "@@controlEvents";
 export declare class AppSyncRealTimeSubscriptionHandshakeLink extends ApolloLink {
@@ -16,11 +16,14 @@ export declare class AppSyncRealTimeSubscriptionHandshakeLink extends ApolloLink
     private subscriptionObserverMap;
     private promiseArray;
     constructor({ url: theUrl, region: theRegion, auth: theAuth }: UrlInfo);
-    request(operation: Operation): Observable<any>;
+    private isCustomDomain;
+    request(operation: Operation): Observable<FetchResult<{
+        [key: string]: any;
+    }, Record<string, any>, Record<string, any>>>;
     private _verifySubscriptionAlreadyStarted;
     private _sendUnsubscriptionMessage;
     private _removeSubscriptionObserver;
-    private _closeSocketWhenFlushed;
+    private _closeSocketIfRequired;
     private _startSubscriptionWithAWSAppSyncRealTime;
     private _initializeWebSocketConnection;
     private _awsRealTimeHeaderBasedAuth;
@@ -33,5 +36,4 @@ export declare class AppSyncRealTimeSubscriptionHandshakeLink extends ApolloLink
     private _errorDisconnect;
     private _timeoutStartSubscriptionAck;
     static createWebSocket(awsRealTimeUrl: string, protocol: string): WebSocket;
-    private static _discoverAppSyncRealTimeEndpoint;
 }

package/lib/realtime-subscription-handshake-link.js

@@ -61,7 +61,7 @@ var __generator = (this && this.__generator) || function (thisArg, body) {
 };
 Object.defineProperty(exports, "__esModule", { value: true });
 /*!
- * Copyright 2017-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * Copyright 2017-2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 var apollo_link_1 = require("apollo-link");
@@ -77,9 +77,9 @@ exports.CONTROL_EVENTS_KEY = "@@controlEvents";
 var NON_RETRYABLE_CODES = [400, 401, 403];
 var SERVICE = "appsync";
 var APPSYNC_REALTIME_HEADERS = {
-    accept: 'application/json, text/javascript',
-    'content-encoding': 'amz-1.0',
-    'content-type': 'application/json; charset=UTF-8'
+    accept: "application/json, text/javascript",
+    "content-encoding": "amz-1.0",
+    "content-type": "application/json; charset=UTF-8",
 };
 /**
  * Time in milliseconds to wait for GQL_CONNECTION_INIT message
@@ -93,6 +93,8 @@ var START_ACK_TIMEOUT = 15000;
  * Default Time in milliseconds to wait for GQL_CONNECTION_KEEP_ALIVE message
  */
 var DEFAULT_KEEP_ALIVE_TIMEOUT = 5 * 60 * 1000;
+var standardDomainPattern = /^https:\/\/\w{26}\.appsync\-api\.\w{2}(?:(?:\-\w{2,})+)\-\d\.amazonaws.com\/graphql$/i;
+var customDomainPath = '/realtime';
 var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super) {
     __extends(AppSyncRealTimeSubscriptionHandshakeLink, _super);
     function AppSyncRealTimeSubscriptionHandshakeLink(_a) {
@@ -107,6 +109,10 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
         _this.auth = theAuth;
         return _this;
     }
+    // Check if url matches standard domain pattern
+    AppSyncRealTimeSubscriptionHandshakeLink.prototype.isCustomDomain = function (url) {
+        return url.match(standardDomainPattern) === null;
+    };
     AppSyncRealTimeSubscriptionHandshakeLink.prototype.request = function (operation) {
         var _a;
         var _this = this;
@@ -129,22 +135,23 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                     _this.auth.type === aws_appsync_auth_link_1.AUTH_TYPE.OPENID_CONNECT
                     ? _this.auth.jwtToken
                     : null;
-                token = _this.auth.type === aws_appsync_auth_link_1.AUTH_TYPE.AWS_LAMBDA ? _this.auth.token : token;
+                token =
+                    _this.auth.type === aws_appsync_auth_link_1.AUTH_TYPE.AWS_LAMBDA ? _this.auth.token : token;
                 var options = {
                     appSyncGraphqlEndpoint: _this.url,
                     authenticationType: _this.auth.type,
                     query: graphql_1.print(query),
                     region: _this.region,
-                    graphql_headers: function () { return (headers); },
+                    graphql_headers: function () { return headers; },
                     variables: variables,
                     apiKey: _this.auth.type === aws_appsync_auth_link_1.AUTH_TYPE.API_KEY ? _this.auth.apiKey : "",
                     credentials: _this.auth.type === aws_appsync_auth_link_1.AUTH_TYPE.AWS_IAM ? _this.auth.credentials : null,
-                    token: token
+                    token: token,
                 };
                 _this._startSubscriptionWithAWSAppSyncRealTime({
                     options: options,
                     observer: observer,
-                    subscriptionId: subscriptionId_1
+                    subscriptionId: subscriptionId_1,
                 });
                 return function () { return __awaiter(_this, void 0, void 0, function () {
                     var subscriptionState;
@@ -190,7 +197,7 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                                 variables: variables,
                                 query: query,
                                 subscriptionReadyCallback: res,
-                                subscriptionFailedCallback: rej
+                                subscriptionFailedCallback: rej,
                             });
                         })];
                 }
@@ -206,7 +213,7 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                 // Preparing unsubscribe message to stop receiving messages for that subscription
                 var unsubscribeMessage = {
                     id: subscriptionId,
-                    type: types_1.MESSAGE_TYPES.GQL_STOP
+                    type: types_1.MESSAGE_TYPES.GQL_STOP,
                 };
                 var stringToAWSRealTime = JSON.stringify(unsubscribeMessage);
                 this.awsRealTimeSocket.send(stringToAWSRealTime);
@@ -220,22 +227,25 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
     };
     AppSyncRealTimeSubscriptionHandshakeLink.prototype._removeSubscriptionObserver = function (subscriptionId) {
         this.subscriptionObserverMap.delete(subscriptionId);
-        if (this.subscriptionObserverMap.size === 0) {
-            // Socket could be sending data to unsubscribe so is required to wait until is flushed
-            this._closeSocketWhenFlushed();
-        }
+        // Verifying for 1000ms after removing subscription in case there are new subscriptions on mount / unmount
+        setTimeout(this._closeSocketIfRequired.bind(this), 1000);
     };
-    AppSyncRealTimeSubscriptionHandshakeLink.prototype._closeSocketWhenFlushed = function () {
-        logger("closing WebSocket...");
-        clearTimeout(this.keepAliveTimeoutId);
+    AppSyncRealTimeSubscriptionHandshakeLink.prototype._closeSocketIfRequired = function () {
+        if (this.subscriptionObserverMap.size > 0) {
+            // There are active subscriptions on the WebSocket
+            return;
+        }
         if (!this.awsRealTimeSocket) {
             this.socketStatus = types_1.SOCKET_STATUS.CLOSED;
             return;
         }
         if (this.awsRealTimeSocket.bufferedAmount > 0) {
-            setTimeout(this._closeSocketWhenFlushed.bind(this), 1000);
+            // There is still data on the WebSocket
+            setTimeout(this._closeSocketIfRequired.bind(this), 1000);
         }
         else {
+            logger("closing WebSocket...");
+            clearTimeout(this.keepAliveTimeoutId);
             var tempSocket = this.awsRealTimeSocket;
             tempSocket.close(1000);
             this.awsRealTimeSocket = null;
@@ -245,17 +255,17 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
     AppSyncRealTimeSubscriptionHandshakeLink.prototype._startSubscriptionWithAWSAppSyncRealTime = function (_a) {
         var options = _a.options, observer = _a.observer, subscriptionId = _a.subscriptionId;
         return __awaiter(this, void 0, void 0, function () {
-            var appSyncGraphqlEndpoint, authenticationType, query, variables, apiKey, region, _b, graphql_headers, credentials, token, subscriptionState, data, dataString, headerObj, _c, _d, subscriptionMessage, stringToAWSRealTime, err_1, _e, message, subscriptionFailedCallback_1, _f, subscriptionFailedCallback, subscriptionReadyCallback;
-            var _g;
+            var appSyncGraphqlEndpoint, authenticationType, query, variables, apiKey, region, _b, graphql_headers, credentials, token, subscriptionState, data, dataString, headerObj, _c, subscriptionMessage, stringToAWSRealTime, err_1, _d, message, subscriptionFailedCallback_1, _e, subscriptionFailedCallback, subscriptionReadyCallback;
+            var _f;
             var _this = this;
-            return __generator(this, function (_h) {
-                switch (_h.label) {
+            return __generator(this, function (_g) {
+                switch (_g.label) {
                     case 0:
                         appSyncGraphqlEndpoint = options.appSyncGraphqlEndpoint, authenticationType = options.authenticationType, query = options.query, variables = options.variables, apiKey = options.apiKey, region = options.region, _b = options.graphql_headers, graphql_headers = _b === void 0 ? function () { return ({}); } : _b, credentials = options.credentials, token = options.token;
                         subscriptionState = types_1.SUBSCRIPTION_STATUS.PENDING;
                         data = {
                             query: query,
-                            variables: variables
+                            variables: variables,
                         };
                         // Having a subscription id map will make it simple to forward messages received
                         this.subscriptionObserverMap.set(subscriptionId, {
@@ -275,45 +285,43 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                                 canonicalUri: "",
                                 region: region,
                                 credentials: credentials,
-                                token: token
+                                token: token,
+                                graphql_headers: graphql_headers,
                             })];
                     case 1:
-                        _d = [__assign.apply(void 0, _c.concat([(_h.sent())]))];
-                        return [4 /*yield*/, graphql_headers()];
-                    case 2:
-                        headerObj = __assign.apply(void 0, [__assign.apply(void 0, _d.concat([(_h.sent())])), (_g = {}, _g[aws_appsync_auth_link_1.USER_AGENT_HEADER] = aws_appsync_auth_link_1.USER_AGENT, _g)]);
+                        headerObj = __assign.apply(void 0, [__assign.apply(void 0, _c.concat([(_g.sent())])), (_f = {}, _f[aws_appsync_auth_link_1.USER_AGENT_HEADER] = aws_appsync_auth_link_1.USER_AGENT, _f)]);
                         subscriptionMessage = {
                             id: subscriptionId,
                             payload: {
                                 data: dataString,
                                 extensions: {
-                                    authorization: __assign({}, headerObj)
-                                }
+                                    authorization: __assign({}, headerObj),
+                                },
                             },
-                            type: types_1.MESSAGE_TYPES.GQL_START
+                            type: types_1.MESSAGE_TYPES.GQL_START,
                         };
                         stringToAWSRealTime = JSON.stringify(subscriptionMessage);
-                        _h.label = 3;
-                    case 3:
-                        _h.trys.push([3, 5, , 6]);
+                        _g.label = 2;
+                    case 2:
+                        _g.trys.push([2, 4, , 5]);
                         return [4 /*yield*/, this._initializeWebSocketConnection({
                                 apiKey: apiKey,
                                 appSyncGraphqlEndpoint: appSyncGraphqlEndpoint,
                                 authenticationType: authenticationType,
                                 region: region,
                                 credentials: credentials,
-                                token: token
+                                token: token,
                             })];
+                    case 3:
+                        _g.sent();
+                        return [3 /*break*/, 5];
                     case 4:
-                        _h.sent();
-                        return [3 /*break*/, 6];
-                    case 5:
-                        err_1 = _h.sent();
-                        _e = err_1.message, message = _e === void 0 ? "" : _e;
+                        err_1 = _g.sent();
+                        _d = err_1.message, message = _d === void 0 ? "" : _d;
                         observer.error({
                             errors: [
-                                __assign({}, new graphql_1.GraphQLError("Connection failed: " + message))
-                            ]
+                                __assign({}, new graphql_1.GraphQLError("Connection failed: " + message)),
+                            ],
                         });
                         observer.complete();
                         subscriptionFailedCallback_1 = (this.subscriptionObserverMap.get(subscriptionId) || {}).subscriptionFailedCallback;
@@ -322,8 +330,8 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                             subscriptionFailedCallback_1();
                         }
                         return [2 /*return*/];
-                    case 6:
-                        _f = this.subscriptionObserverMap.get(subscriptionId), subscriptionFailedCallback = _f.subscriptionFailedCallback, subscriptionReadyCallback = _f.subscriptionReadyCallback;
+                    case 5:
+                        _e = this.subscriptionObserverMap.get(subscriptionId), subscriptionFailedCallback = _e.subscriptionFailedCallback, subscriptionReadyCallback = _e.subscriptionReadyCallback;
                         // This must be done before sending the message in order to be listening immediately
                         this.subscriptionObserverMap.set(subscriptionId, {
                             observer: observer,
@@ -334,7 +342,7 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                             subscriptionFailedCallback: subscriptionFailedCallback,
                             startAckTimeoutId: setTimeout(function () {
                                 _this._timeoutStartSubscriptionAck.call(_this, subscriptionId);
-                            }, START_ACK_TIMEOUT)
+                            }, START_ACK_TIMEOUT),
                         });
                         if (this.awsRealTimeSocket) {
                             this.awsRealTimeSocket.send(stringToAWSRealTime);
@@ -351,7 +359,7 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
             return;
         }
         return new Promise(function (res, rej) { return __awaiter(_this, void 0, void 0, function () {
-            var discoverableEndpoint, payloadString, headerString, _a, _b, headerQs, payloadQs, awsRealTimeUrl, err_2;
+            var payloadString, headerString, _a, _b, headerQs, payloadQs, discoverableEndpoint, awsRealTimeUrl, err_2;
             return __generator(this, function (_c) {
                 switch (_c.label) {
                     case 0:
@@ -361,7 +369,6 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                     case 1:
                         _c.trys.push([1, 4, , 5]);
                         this.socketStatus = types_1.SOCKET_STATUS.CONNECTING;
-                        discoverableEndpoint = AppSyncRealTimeSubscriptionHandshakeLink._discoverAppSyncRealTimeEndpoint(this.url);
                         payloadString = "{}";
                         _b = (_a = JSON).stringify;
                         return [4 /*yield*/, this._awsRealTimeHeaderBasedAuth({
@@ -372,12 +379,23 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                                 appSyncGraphqlEndpoint: appSyncGraphqlEndpoint,
                                 region: region,
                                 credentials: credentials,
-                                token: token
+                                token: token,
+                                graphql_headers: function () { },
                             })];
                     case 2:
                         headerString = _b.apply(_a, [_c.sent()]);
                         headerQs = Buffer.from(headerString).toString("base64");
                         payloadQs = Buffer.from(payloadString).toString("base64");
+                        discoverableEndpoint = appSyncGraphqlEndpoint;
+                        if (this.isCustomDomain(discoverableEndpoint)) {
+                            discoverableEndpoint = discoverableEndpoint.concat(customDomainPath);
+                        }
+                        else {
+                            discoverableEndpoint = discoverableEndpoint.replace('appsync-api', 'appsync-realtime-api').replace('gogi-beta', 'grt-beta');
+                        }
+                        discoverableEndpoint = discoverableEndpoint
+                            .replace("https://", "wss://")
+                            .replace('http://', 'ws://');
                         awsRealTimeUrl = discoverableEndpoint + "?header=" + headerQs + "&payload=" + payloadQs;
                         return [4 /*yield*/, this._initializeRetryableHandshake({ awsRealTimeUrl: awsRealTimeUrl })];
                     case 3:
@@ -410,7 +428,7 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
         }); });
     };
     AppSyncRealTimeSubscriptionHandshakeLink.prototype._awsRealTimeHeaderBasedAuth = function (_a) {
-        var authenticationType = _a.authenticationType, payload = _a.payload, canonicalUri = _a.canonicalUri, appSyncGraphqlEndpoint = _a.appSyncGraphqlEndpoint, apiKey = _a.apiKey, region = _a.region, credentials = _a.credentials, token = _a.token;
+        var authenticationType = _a.authenticationType, payload = _a.payload, canonicalUri = _a.canonicalUri, appSyncGraphqlEndpoint = _a.appSyncGraphqlEndpoint, apiKey = _a.apiKey, region = _a.region, credentials = _a.credentials, token = _a.token, graphql_headers = _a.graphql_headers;
         return __awaiter(this, void 0, void 0, function () {
             var headerHandler, handler, host, result;
             return __generator(this, function (_b) {
@@ -421,7 +439,7 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                             AWS_IAM: this._awsRealTimeIAMHeader.bind(this),
                             OPENID_CONNECT: this._awsRealTimeAuthorizationHeader.bind(this),
                             AMAZON_COGNITO_USER_POOLS: this._awsRealTimeAuthorizationHeader.bind(this),
-                            AWS_LAMBDA: this._awsRealTimeAuthorizationHeader.bind(this)
+                            AWS_LAMBDA: this._awsRealTimeAuthorizationHeader.bind(this),
                         };
                         handler = headerHandler[authenticationType];
                         if (typeof handler !== "function") {
@@ -437,7 +455,8 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                                 region: region,
                                 host: host,
                                 credentials: credentials,
-                                token: token
+                                token: token,
+                                graphql_headers: graphql_headers,
                             })];
                     case 1:
                         result = _b.sent();
@@ -447,41 +466,43 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
         });
     };
     AppSyncRealTimeSubscriptionHandshakeLink.prototype._awsRealTimeAuthorizationHeader = function (_a) {
-        var host = _a.host, token = _a.token;
+        var host = _a.host, token = _a.token, graphql_headers = _a.graphql_headers;
         return __awaiter(this, void 0, void 0, function () {
-            var _b, _c;
-            return __generator(this, function (_d) {
-                switch (_d.label) {
+            var _b, _c, _d;
+            return __generator(this, function (_e) {
+                switch (_e.label) {
                     case 0:
                         _b = {};
                         if (!(typeof token === "function")) return [3 /*break*/, 2];
                         return [4 /*yield*/, token.call(undefined)];
                     case 1:
-                        _c = _d.sent();
+                        _c = _e.sent();
                         return [3 /*break*/, 4];
                     case 2: return [4 /*yield*/, token];
                     case 3:
-                        _c = _d.sent();
-                        _d.label = 4;
-                    case 4: return [2 /*return*/, (_b.Authorization = _c,
-                            _b.host = host,
-                            _b)];
+                        _c = _e.sent();
+                        _e.label = 4;
+                    case 4:
+                        _d = [(_b.Authorization = _c, _b.host = host, _b)];
+                        return [4 /*yield*/, graphql_headers()];
+                    case 5: return [2 /*return*/, __assign.apply(void 0, _d.concat([(_e.sent())]))];
                 }
             });
         });
     };
     AppSyncRealTimeSubscriptionHandshakeLink.prototype._awsRealTimeApiKeyHeader = function (_a) {
-        var apiKey = _a.apiKey, host = _a.host;
+        var apiKey = _a.apiKey, host = _a.host, graphql_headers = _a.graphql_headers;
         return __awaiter(this, void 0, void 0, function () {
-            var dt, dtStr;
-            return __generator(this, function (_b) {
-                dt = new Date();
-                dtStr = dt.toISOString().replace(/[:\-]|\.\d{3}/g, "");
-                return [2 /*return*/, {
-                        host: host,
-                        "x-amz-date": dtStr,
-                        "x-api-key": apiKey
-                    }];
+            var dt, dtStr, _b;
+            return __generator(this, function (_c) {
+                switch (_c.label) {
+                    case 0:
+                        dt = new Date();
+                        dtStr = dt.toISOString().replace(/[:\-]|\.\d{3}/g, "");
+                        _b = [{ host: host, "x-amz-date": dtStr, "x-api-key": apiKey }];
+                        return [4 /*yield*/, graphql_headers()];
+                    case 1: return [2 /*return*/, __assign.apply(void 0, _b.concat([(_c.sent())]))];
+                }
             });
         });
     };
@@ -494,7 +515,7 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                     case 0:
                         endpointInfo = {
                             region: region,
-                            service: SERVICE
+                            service: SERVICE,
                         };
                         creds = typeof credentials === "function"
                             ? credentials.call()
@@ -514,13 +535,13 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                         formattedCredentials = {
                             access_key: accessKeyId,
                             secret_key: secretAccessKey,
-                            session_token: sessionToken
+                            session_token: sessionToken,
                         };
                         request = {
                             url: "" + appSyncGraphqlEndpoint + canonicalUri,
                             body: payload,
                             method: "POST",
-                            headers: __assign({}, APPSYNC_REALTIME_HEADERS)
+                            headers: __assign({}, APPSYNC_REALTIME_HEADERS),
                         };
                         signed_params = aws_appsync_auth_link_1.Signer.sign(request, formattedCredentials, endpointInfo);
                         return [2 /*return*/, signed_params.headers];
@@ -536,7 +557,7 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                     case 0:
                         logger("Initializaling retryable Handshake");
                         return [4 /*yield*/, retry_1.jitteredExponentialRetry(this._initializeHandshake.bind(this), [
-                                { awsRealTimeUrl: awsRealTimeUrl }
+                                { awsRealTimeUrl: awsRealTimeUrl },
                             ])];
                     case 1:
                         _b.sent();
@@ -592,7 +613,8 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                                         if (type === types_1.MESSAGE_TYPES.GQL_CONNECTION_ACK) {
                                             ackOk = true;
                                             _this.keepAliveTimeout = connectionTimeoutMs;
-                                            _this.awsRealTimeSocket.onmessage = _this._handleIncomingSubscriptionMessage.bind(_this);
+                                            _this.awsRealTimeSocket.onmessage =
+                                                _this._handleIncomingSubscriptionMessage.bind(_this);
                                             _this.awsRealTimeSocket.onerror = function (err) {
                                                 logger(err);
                                                 _this._errorDisconnect(types_1.CONTROL_MSG.CONNECTION_CLOSED);
@@ -610,7 +632,7 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                                         }
                                     };
                                     var gqlInit = {
-                                        type: types_1.MESSAGE_TYPES.GQL_CONNECTION_INIT
+                                        type: types_1.MESSAGE_TYPES.GQL_CONNECTION_INIT,
                                     };
                                     _this.awsRealTimeSocket.send(JSON.stringify(gqlInit));
                                     function checkAckOk() {
@@ -667,8 +689,8 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                 observer.next({
                     data: payload,
                     extensions: {
-                        controlMsgType: "CONNECTED"
-                    }
+                        controlMsgType: "CONNECTED",
+                    },
                 });
             }
             else {
@@ -682,7 +704,7 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                 startAckTimeoutId: null,
                 subscriptionState: subscriptionState,
                 subscriptionReadyCallback: subscriptionReadyCallback,
-                subscriptionFailedCallback: subscriptionFailedCallback
+                subscriptionFailedCallback: subscriptionFailedCallback,
             });
             return;
         }
@@ -700,12 +722,12 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                 startAckTimeoutId: startAckTimeoutId,
                 subscriptionReadyCallback: subscriptionReadyCallback,
                 subscriptionFailedCallback: subscriptionFailedCallback,
-                subscriptionState: subscriptionState
+                subscriptionState: subscriptionState,
             });
             observer.error({
                 errors: [
-                    __assign({}, new graphql_1.GraphQLError("Connection failed: " + JSON.stringify(payload)))
-                ]
+                    __assign({}, new graphql_1.GraphQLError("Connection failed: " + JSON.stringify(payload))),
+                ],
             });
             clearTimeout(startAckTimeoutId);
             observer.complete();
@@ -739,13 +761,13 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
             observer: observer,
             query: query,
             variables: variables,
-            subscriptionState: types_1.SUBSCRIPTION_STATUS.FAILED
+            subscriptionState: types_1.SUBSCRIPTION_STATUS.FAILED,
         });
         if (observer && !observer.closed) {
             observer.error({
                 errors: [
-                    __assign({}, new graphql_1.GraphQLError("Subscription timeout " + JSON.stringify({ query: query, variables: variables })))
-                ]
+                    __assign({}, new graphql_1.GraphQLError("Subscription timeout " + JSON.stringify({ query: query, variables: variables }))),
+                ],
             });
             // Cleanup will be automatically executed
             observer.complete();
@@ -755,13 +777,6 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
     AppSyncRealTimeSubscriptionHandshakeLink.createWebSocket = function (awsRealTimeUrl, protocol) {
         return new WebSocket(awsRealTimeUrl, protocol);
     };
-    AppSyncRealTimeSubscriptionHandshakeLink._discoverAppSyncRealTimeEndpoint = function (url) {
-        return url
-            .replace("https://", "wss://")
-            .replace('http://', 'ws://')
-            .replace("appsync-api", "appsync-realtime-api")
-            .replace("gogi-beta", "grt-beta");
-    };
     return AppSyncRealTimeSubscriptionHandshakeLink;
 }(apollo_link_1.ApolloLink));
 exports.AppSyncRealTimeSubscriptionHandshakeLink = AppSyncRealTimeSubscriptionHandshakeLink;

package/lib/subscription-handshake-link.d.ts

@@ -15,7 +15,7 @@ export declare class SubscriptionHandshakeLink extends ApolloLink {
     private topicObservers;
     private clientObservers;
     constructor(subsInfoContextKey: any);
-    request(operation: Operation): Observable<any>;
+    request(operation: Operation): Observable<unknown>;
     connectNewClients(connectionInfo: MqttConnectionInfo[], observer: ZenObservable.Observer<FetchResult>, operation: Operation): Promise<any[]>;
     connectNewClient(connectionInfo: MqttConnectionInfo, observer: ZenObservable.Observer<FetchResult>, selectionNames: string[]): Promise<any>;
     subscribeToTopics<T>(client: any, topics: string[], observer: ZenObservable.Observer<T>): Promise<unknown[]>;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "aws-appsync-subscription-link",
-  "version": "2.2.5",
+  "version": "2.3.0",
   "main": "lib/index.js",
   "license": "Apache-2.0",
   "description": "AWS Mobile AppSync SDK for JavaScript",
@@ -23,12 +23,13 @@
     "apollo-link-context": "1.0.11",
     "apollo-link-http": "1.5.8",
     "apollo-link-retry": "2.2.7",
-    "aws-appsync-auth-link": "^2.0.6",
+    "aws-appsync-auth-link": "^2.0.8",
     "debug": "2.6.9",
     "url": "^0.11.0"
   },
   "devDependencies": {
-    "@redux-offline/redux-offline": "2.5.2-native.0"
+    "@redux-offline/redux-offline": "2.5.2-native.0",
+    "apollo-client": "2.x"
   },
   "peerDependencies": {
     "apollo-client": "2.x"