From 874cdbbccbc2bf8f4ee18abe220e87dc52e6a8db Mon Sep 17 00:00:00 2001
From: Maneesh Tewani <maneesht@users.noreply.github.com>
Date: Thu, 5 May 2022 17:27:03 -0700
Subject: [PATCH] Added GMPID to websocket connection (#6232)

* Added applicationid to websocket connection
---
 .changeset/beige-turkeys-occur.md             |  5 ++
 .../src/realtime/WebSocketConnection.ts       | 22 ++++-----
 .../database/test/websocketconnection.test.ts | 48 +++++++++++++++++++
 3 files changed, 63 insertions(+), 12 deletions(-)
 create mode 100644 .changeset/beige-turkeys-occur.md
 create mode 100644 packages/database/test/websocketconnection.test.ts

diff --git a/.changeset/beige-turkeys-occur.md b/.changeset/beige-turkeys-occur.md
new file mode 100644
index 00000000000..7a020d99188
--- /dev/null
+++ b/.changeset/beige-turkeys-occur.md
@@ -0,0 +1,5 @@
+---
+"@firebase/database": patch
+---
+
+Added GMPID to websocket connection.
diff --git a/packages/database/src/realtime/WebSocketConnection.ts b/packages/database/src/realtime/WebSocketConnection.ts
index abeb0f70474..d23dd2afe86 100644
--- a/packages/database/src/realtime/WebSocketConnection.ts
+++ b/packages/database/src/realtime/WebSocketConnection.ts
@@ -25,6 +25,7 @@ import { logWrapper, splitStringBySize } from '../core/util/util';
 import { SDK_VERSION } from '../core/version';
 
 import {
+  APPLICATION_ID_PARAM,
   APP_CHECK_TOKEN_PARAM,
   FORGE_DOMAIN_RE,
   FORGE_REF,
@@ -99,7 +100,8 @@ export class WebSocketConnection implements Transport {
       repoInfo,
       transportSessionId,
       lastSessionId,
-      appCheckToken
+      appCheckToken,
+      applicationId
     );
     this.nodeAdmin = repoInfo.nodeAdmin;
   }
@@ -115,7 +117,8 @@ export class WebSocketConnection implements Transport {
     repoInfo: RepoInfo,
     transportSessionId?: string,
     lastSessionId?: string,
-    appCheckToken?: string
+    appCheckToken?: string,
+    applicationId?: string
   ): string {
     const urlParams: { [k: string]: string } = {};
     urlParams[VERSION_PARAM] = PROTOCOL_VERSION;
@@ -137,6 +140,9 @@ export class WebSocketConnection implements Transport {
     if (appCheckToken) {
       urlParams[APP_CHECK_TOKEN_PARAM] = appCheckToken;
     }
+    if (applicationId) {
+      urlParams[APPLICATION_ID_PARAM] = applicationId;
+    }
 
     return repoInfoConnectionURL(repoInfo, WEBSOCKET, urlParams);
   }
@@ -156,6 +162,7 @@ export class WebSocketConnection implements Transport {
     PersistentStorage.set('previous_websocket_failure', true);
 
     try {
+      let options: { [k: string]: object };
       if (isNodeSdk()) {
         const device = this.nodeAdmin ? 'AdminNode' : 'Node';
         // UA Format: Firebase/<wire_protocol>/<sdk_version>/<platform>/<device>
@@ -188,17 +195,8 @@ export class WebSocketConnection implements Transport {
         if (proxy) {
           options['proxy'] = { origin: proxy };
         }
-
-        this.mySock = new WebSocketImpl(this.connURL, [], options);
-      } else {
-        const options: { [k: string]: object } = {
-          headers: {
-            'X-Firebase-GMPID': this.applicationId || '',
-            'X-Firebase-AppCheck': this.appCheckToken || ''
-          }
-        };
-        this.mySock = new WebSocketImpl(this.connURL, [], options);
       }
+      this.mySock = new WebSocketImpl(this.connURL, [], options);
     } catch (e) {
       this.log_('Error instantiating WebSocket.');
       const error = e.message || e.data;
diff --git a/packages/database/test/websocketconnection.test.ts b/packages/database/test/websocketconnection.test.ts
new file mode 100644
index 00000000000..460ba6767a9
--- /dev/null
+++ b/packages/database/test/websocketconnection.test.ts
@@ -0,0 +1,48 @@
+/**
+ * @license
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { expect } from 'chai';
+
+import { APPLICATION_ID_PARAM } from '../src/realtime/Constants';
+import { WebSocketConnection } from '../src/realtime/WebSocketConnection';
+
+import { testRepoInfo } from './helpers/util';
+
+describe('WebSocketConnection', () => {
+  it('should add an applicationId to the query parameter', () => {
+    const repoInfo = testRepoInfo('https://test-ns.firebaseio.com');
+    const applicationId = 'myID';
+    const websocketConnection = new WebSocketConnection(
+      'connId',
+      repoInfo,
+      applicationId
+    );
+    const searchParams = new URL(websocketConnection.connURL).searchParams;
+    expect(searchParams.get(APPLICATION_ID_PARAM)).to.equal(applicationId);
+  });
+  it('should not add an applicationId to the query parameter if applicationId is empty', () => {
+    const repoInfo = testRepoInfo('https://test-ns.firebaseio.com');
+    const applicationId = '';
+    const websocketConnection = new WebSocketConnection(
+      'connId',
+      repoInfo,
+      applicationId
+    );
+    const searchParams = new URL(websocketConnection.connURL).searchParams;
+    expect(searchParams.get(APPLICATION_ID_PARAM)).to.be.null;
+  });
+});