aws-appsync-subscription-link 3.0.8 → 3.0.11

package/CHANGELOG.md

@@ -3,6 +3,38 @@
 All notable changes to this project will be documented in this file.
 See [Conventional Commits](https://conventionalcommits.org) for commit guidelines.
 
+<a name="3.0.11"></a>
+## [3.0.11](https://github.com/awslabs/aws-mobile-appsync-sdk-js/compare/aws-appsync-subscription-link@3.0.10...aws-appsync-subscription-link@3.0.11) (2022-05-02)
+
+
+
+
+**Note:** Version bump only for package aws-appsync-subscription-link
+
+<a name="3.0.10"></a>
+## [3.0.10](https://github.com/awslabs/aws-mobile-appsync-sdk-js/compare/aws-appsync-subscription-link@3.0.9...aws-appsync-subscription-link@3.0.10) (2022-03-04)
+
+
+### Bug Fixes
+
+* Port over Amplify fix for subscription race conditions ([#509](https://github.com/awslabs/aws-mobile-appsync-sdk-js/issues/509)) ([#704](https://github.com/awslabs/aws-mobile-appsync-sdk-js/issues/704)) ([92b50c4](https://github.com/awslabs/aws-mobile-appsync-sdk-js/commit/92b50c4))
+
+
+
+
+<a name="3.0.9"></a>
+## [3.0.9](https://github.com/awslabs/aws-mobile-appsync-sdk-js/compare/aws-appsync-subscription-link@3.0.8...aws-appsync-subscription-link@3.0.9) (2021-09-24)
+
+
+### Bug Fixes
+
+* add type fix to subscription link to fix 'prepare' script ([#669](https://github.com/awslabs/aws-mobile-appsync-sdk-js/issues/669)) ([c5202dc](https://github.com/awslabs/aws-mobile-appsync-sdk-js/commit/c5202dc))
+* **aws-appsync-subscription-link:** graphql header refactor to fix IAM-based auth ([#619](https://github.com/awslabs/aws-mobile-appsync-sdk-js/issues/619)) ([#671](https://github.com/awslabs/aws-mobile-appsync-sdk-js/issues/671)) ([dbf4959](https://github.com/awslabs/aws-mobile-appsync-sdk-js/commit/dbf4959))
+* remove aws-sdk V2; import aws-sdk V3 packages ([#637](https://github.com/awslabs/aws-mobile-appsync-sdk-js/issues/637)); refactor auth-link, signer, and types ([0996740](https://github.com/awslabs/aws-mobile-appsync-sdk-js/commit/0996740))
+
+
+
+
 <a name="3.0.8"></a>
 ## [3.0.8](https://github.com/awslabs/aws-mobile-appsync-sdk-js/compare/aws-appsync-subscription-link@3.0.7...aws-appsync-subscription-link@3.0.8) (2021-07-28)
 

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

@@ -6,7 +6,8 @@ import { AppSyncRealTimeSubscriptionHandshakeLink } from '../../src/realtime-sub
 const query = gql`subscription { someSubscription { aField } }`
 
 class myWebSocket implements WebSocket {
-    binaryType: BinaryType; bufferedAmount: number;
+    binaryType: BinaryType; 
+    bufferedAmount: number;
     extensions: string;
     onclose: (this: WebSocket, ev: CloseEvent) => any;
     onerror: (this: WebSocket, ev: Event) => any;
@@ -49,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);
@@ -62,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();
@@ -72,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) => {
@@ -99,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();
@@ -109,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) => {
@@ -136,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();
@@ -146,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) => {
@@ -173,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();
@@ -183,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/index.js

@@ -28,7 +28,7 @@ function createSubscriptionHandshakeLink(infoOrUrl, theResultsFetcherLink) {
                 link: new core_1.ApolloLink(function (operation, _forward) {
                     return new core_1.Observable(function (observer) {
                         var _a;
-                        var _b = operation.variables, _c = subscription_handshake_link_1.CONTROL_EVENTS_KEY, controlEvents = _b[_c], variables = __rest(_b, [typeof _c === "symbol" ? _c : _c + ""]);
+                        var _b = operation, _c = _b.variables, _d = subscription_handshake_link_1.CONTROL_EVENTS_KEY, controlEvents = _c[_d], variables = __rest(_c, [typeof _d === "symbol" ? _d : _d + ""]);
                         if (typeof controlEvents !== "undefined") {
                             operation.variables = variables;
                         }

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

@@ -2,12 +2,14 @@
  * Copyright 2017-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-import { ApolloLink, NextLink } from '@apollo/client/core';
+import { ApolloLink } from '@apollo/client/core';
+import type { NextLink, FetchResult } from '@apollo/client/core';
+import type { Observable } from 'zen-observable-ts';
 export declare class NonTerminatingLink extends ApolloLink {
     private contextKey;
     private link;
     constructor(contextKey: string, { link }: {
         link: ApolloLink;
     });
-    request(operation: any, forward?: NextLink): any;
+    request(operation: any, forward?: NextLink): Observable<FetchResult>;
 }

package/lib/non-terminating-link.js

@@ -77,7 +77,7 @@ var NonTerminatingLink = /** @class */ (function (_super) {
     }
     NonTerminatingLink.prototype.request = function (operation, forward) {
         var _this = this;
-        return context_1.setContext(function (_request, prevContext) { return __awaiter(_this, void 0, void 0, function () {
+        return (context_1.setContext(function (_request, prevContext) { return __awaiter(_this, void 0, void 0, function () {
             var result;
             var _a;
             var _this = this;
@@ -94,7 +94,7 @@ var NonTerminatingLink = /** @class */ (function (_super) {
                         return [2 /*return*/, __assign(__assign({}, prevContext), (_a = {}, _a[this.contextKey] = result, _a))];
                 }
             });
-        }); }).request(operation, forward);
+        }); })).request(operation, forward);
     };
     return NonTerminatingLink;
 }(core_1.ApolloLink));

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, Operation } from "@apollo/client/core";
+import { ApolloLink, Observable, Operation, FetchResult } from "@apollo/client/core";
 import { UrlInfo } from "./types";
 export declare const CONTROL_EVENTS_KEY = "@@controlEvents";
 export declare class AppSyncRealTimeSubscriptionHandshakeLink extends ApolloLink {
@@ -16,11 +16,12 @@ export declare class AppSyncRealTimeSubscriptionHandshakeLink extends ApolloLink
     private subscriptionObserverMap;
     private promiseArray;
     constructor({ url: theUrl, region: theRegion, auth: theAuth }: UrlInfo);
-    request(operation: Operation): any;
+    private isCustomDomain;
+    request(operation: Operation): Observable<FetchResult<Record<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 +34,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 core_1 = require("@apollo/client/core");
@@ -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;
@@ -220,22 +226,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,11 +254,11 @@ 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;
@@ -275,13 +284,11 @@ 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: {
@@ -293,9 +300,9 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                             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,
@@ -304,12 +311,12 @@ var AppSyncRealTimeSubscriptionHandshakeLink = /** @class */ (function (_super)
                                 credentials: credentials,
                                 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))
@@ -322,8 +329,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,
@@ -351,7 +358,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 +368,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 +378,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 +427,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) {
@@ -437,7 +454,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 +465,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())]))];
+                }
             });
         });
     };
@@ -755,13 +775,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;
 }(core_1.ApolloLink));
 exports.AppSyncRealTimeSubscriptionHandshakeLink = AppSyncRealTimeSubscriptionHandshakeLink;

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

@@ -2,7 +2,7 @@
  * Copyright 2017-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-import { ApolloLink, Operation, FetchResult } from "@apollo/client/core";
+import { ApolloLink, Observable, Operation, FetchResult } from "@apollo/client/core";
 declare type MqttConnectionInfo = {
     client: string;
     url: string;
@@ -14,7 +14,7 @@ export declare class SubscriptionHandshakeLink extends ApolloLink {
     private topicObservers;
     private clientObservers;
     constructor(subsInfoContextKey: any);
-    request(operation: Operation): 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": "3.0.8",
+  "version": "3.0.11",
   "main": "lib/index.js",
   "license": "Apache-2.0",
   "description": "AWS Mobile AppSync SDK for JavaScript",
@@ -19,7 +19,7 @@
     "test-watch": "jest --watch"
   },
   "dependencies": {
-    "aws-appsync-auth-link": "^3.0.6",
+    "aws-appsync-auth-link": "^3.0.7",
     "debug": "2.6.9",
     "url": "^0.11.0"
   },