Merge pull request #41 from browserstack/PPLT-2916

Executing customRequests in Playwright session

Author: pranavj1001

lib/config/constants.js

@@ -1,6 +1,7 @@
 'use strict';
 
 const env = process.env.NODE_ENV || 'dev';
+const customRequestEnabled = process.env.CUSTOM_REQUEST_ENABLED === 'true';
 const config = require('./config.json')[env];
 const logger = require('../util/loggerFactory');
 const {
@@ -252,7 +253,8 @@ module.exports = {
   kEnableIncomingQueue,
   kEnableOutgoingQueue,
   kUpstreamRestart,
+  customRequestEnabled,
   PROXY_LOCKED,
   PROXY_RESTART,
-  HTTPLOG
+  HTTPLOG,
 };

lib/core/CustomRequestHandler.js

@@ -0,0 +1,91 @@
+'use strict';
+
+const logger = require('../util/loggerFactory');
+
+/**
+ * Handles the custom requests made to existing playwright connection.
+ * This class is implemented as a Singleton to maintain a map of commands for which
+ * responses can be resolved once received from the Playwright server.
+ * @class
+ */
+class CustomRequestHandler {
+  /**
+   * Creates an instance of CustomRequestHandler.
+   * @constructor
+   */
+  constructor() {
+    if (!CustomRequestHandler.instance) {
+      // Initialize the instance if it doesn't exist
+      CustomRequestHandler.instance = this;
+      // Initialize the map {} as part of the instance
+      this.customRequestList = {};
+    }
+    // Return the existing instance if it already exists
+    return CustomRequestHandler.instance;
+  }
+
+  /**
+   * Static method to get the single instance of the class.
+   * @returns {CustomRequestHandler} The single instance of the CustomRequestHandler class.
+   */
+  static getInstance() {
+    if (!CustomRequestHandler.instance) {
+      // Create a new instance if it doesn't exist
+      CustomRequestHandler.instance = new CustomRequestHandler();
+    }
+    // Return the existing instance
+    return CustomRequestHandler.instance;
+  }
+
+  /**
+   * Checks if the custom request list is empty.
+   * @returns {boolean} Returns true if the custom request list is empty, otherwise false.
+   */
+  isCustomRequestListEmpty() {
+    for (const prop in this.customRequestList) {
+      if (this.customRequestList.hasOwnProperty(prop)) {
+        return false;
+      }
+    }
+
+    return true;
+  }
+
+  /**
+   * Adds an item to the custom request list.
+   * @param {string} request_id - The ID of the request to be added.
+   */
+  addCustomRequest(request_id) {
+    let resolveFunc;
+    let rejectFunc;
+    let promise = new Promise((resolve, reject) => {
+      resolveFunc = resolve;
+      rejectFunc = reject;
+    });
+    this.customRequestList[request_id] = {
+      resolve: resolveFunc,
+      reject: rejectFunc,
+      promise: promise,
+    };
+    logger.info(`Added request '${request_id}' to the customRequestList.`);
+  }
+
+  /**
+   * Gets the items in the custom request list.
+   * @returns {Object} The custom request list.
+   */
+  getList() {
+    return this.customRequestList;
+  }
+
+  /**
+   * Resets the instance of the CustomRequestHandler class.
+   * Only for testing purposes. Do not use it in production code.
+   * @static
+   */
+  static resetInstance() {
+    CustomRequestHandler.instance = null;
+  }
+}
+
+module.exports = CustomRequestHandler;

lib/core/IncomingWebSocket.js

@@ -117,8 +117,7 @@ class IncomingWebSocket extends EventEmitter {
    */
   close(code, msg) {
     this.teardown = true;
-    if(code >= 1000 && code < 1004)
-      this.socket.close(code, msg);
+    if (code >= 1000 && code < 1004) this.socket.close(code, msg);
     this.socket.terminate();
   }
 

lib/core/OutgoingWebSocket.js

@@ -24,10 +24,12 @@ const {
   PROXY_RESTART,
   DISALLOWED_HEADERS,
   OUTGOING,
+  customRequestEnabled,
 } = require('../config/constants');
 const { extractConnectionId } = require('../util/util');
 const { incrReconnectionCount } = require('../util/metrics');
 const { isNotUndefined } = require('../util/typeSanity');
+const CustomRequestHandler = require('./CustomRequestHandler');
 
 /**
  * Outgoing WebSocket connection is the connection object
@@ -107,6 +109,20 @@ class OutgoingWebSocket extends EventEmitter {
       this.emit(kEnableIncomingQueue);
       return;
     }
+
+    const customReqInstance = CustomRequestHandler.getInstance();
+    if (customRequestEnabled && !customReqInstance.isCustomRequestListEmpty()) {
+      let resp;
+      try {
+        resp = JSON.parse(msg);
+        if (resp && customReqInstance.getList().hasOwnProperty(resp.id)) {
+          customReqInstance.customRequestList[resp.id].resolve(msg);
+          return;
+        }
+      } catch (error) {
+        logger.error(`Error parsing JSON: ${error}`);
+      }
+    }
     this.emit(kMessageReceived, msg);
   }
 
@@ -136,6 +152,19 @@ class OutgoingWebSocket extends EventEmitter {
    * Triggers when error occured on socket.
    */
   errorHandler(error) {
+    const customReqInstance = CustomRequestHandler.getInstance();
+    if (customRequestEnabled && !customReqInstance.isCustomRequestListEmpty()) {
+      let resp;
+      try {
+        resp = JSON.parse(error);
+        if (resp && customReqInstance.getList().hasOwnProperty(resp.id)) {
+          customReqInstance.customRequestList[resp.id].reject(error);
+          return;
+        }
+      } catch (error) {
+        logger.error(`Error parsing JSON: ${error}`);
+      }
+    }
     this.emit(kError, error);
   }
 
@@ -164,12 +193,9 @@ class OutgoingWebSocket extends EventEmitter {
    *  Closes the socket connection.
    */
   close(code, msg) {
-    if(code == 1006)
-      this.socket.terminate();
-    else if(code == 1005)
-      this.socket.close();
-    else
-      this.socket.close(code, msg);
+    if (code == 1006) this.socket.terminate();
+    else if (code == 1005) this.socket.close();
+    else this.socket.close(code, msg);
   }
 
   /**

lib/core/Proxy.js

@@ -1,7 +1,13 @@
 'use strict';
 
 const WebSocket = require('ws');
-const { config, kCleanup, kAddNewContext, HTTPLOG } = require('../config/constants');
+const {
+  config,
+  kCleanup,
+  kAddNewContext,
+  HTTPLOG,
+  customRequestEnabled,
+} = require('../config/constants');
 const logger = require('../util/loggerFactory');
 const Context = require('./Context');
 const {
@@ -16,6 +22,12 @@ const { setMetrics } = require('../util/metrics');
 const AlertManager = require('../util/AlertManager');
 const Instrumentation = require('../util/Instrumentation');
 const ErrorHandler = require('../util/ErrorHandler');
+const CustomRequestHandler = require('./CustomRequestHandler');
+const responseHeaders = {
+  'content-type': 'application/json; charset=utf-8',
+  accept: 'application/json',
+  'WWW-Authenticate': 'Basic realm="WS Reconnect Proxy"',
+}
 
 /**
  * Proxy is the entrypoint and instantiates the context among the socket connection.
@@ -64,9 +76,58 @@ class Proxy {
       headers: request.headers,
     };
     logger.info(`${HTTPLOG} Received http request ${options}`);
-    if(request.url.indexOf('/status') > -1){
-      response.writeHead(200, {'content-type': 'application/json; charset=utf-8', 'accept': 'application/json', 'WWW-Authenticate': 'Basic realm="WS Reconnect Proxy"'});
-      response.end(JSON.stringify({"status" : "Running"}));
+    if (request.url.indexOf('/status') > -1) {
+      response.writeHead(200, responseHeaders);
+      response.end(JSON.stringify({ status: 'Running' }));
+      return;
+    } else if (
+      customRequestEnabled &&
+      request.url.indexOf('/customRequest') > -1 &&
+      request.method == 'POST'
+    ) {
+      try {
+        logger.info(`Handling request to execute custom command in server`);
+        let body = '';
+
+        // Read data from the request
+        request.on('data', (chunk) => {
+          body += chunk.toString(); // Convert Buffer to string
+        });
+
+        // When the request ends, process the received data
+        request.on('end', () => {
+          body = JSON.parse(body);
+          const command = body.command;
+          const commandId = body.command.id;
+          //Create singleton object and map the command id with pending promise
+          const customReqInstance = CustomRequestHandler.getInstance();
+          customReqInstance.addCustomRequest(commandId);
+
+          //Send to playwright server
+          const sessionId = [...this.contexts.keys()][0];
+          const sessionContext = this.contexts.get(sessionId);
+          sessionContext.outgoingSocket.send(JSON.stringify(command));
+
+          //Get the resolved promise and returning it to end user
+          customReqInstance.customRequestList[commandId].promise
+            .then((result) => {
+              delete customReqInstance.customRequestList[commandId];
+              response.writeHead(200, responseHeaders);
+              response.end(
+                JSON.stringify({ status: 'success', value: result })
+              );
+            })
+            .catch((err) => {
+              delete customReqInstance.customRequestList[commandId];
+              response.writeHead(500, responseHeaders);
+              response.end(JSON.stringify({ status: 'failure', value: err }));
+            });
+        });
+      } catch (err) {
+        logger.error(`Error while handling custom request ${err}`);
+        response.writeHead(500, responseHeaders);
+        response.end(JSON.stringify({ status: 'failure', value: err.message }));
+      }
       return;
     }
     const proxyReq = http.request(options, (proxyResponse) => {
@@ -75,13 +136,14 @@ class Proxy {
         end: true,
       });
     });
-    proxyReq.on('error', (error)=>{
+
+    proxyReq.on('error', (error) => {
       logger.error(`${request.url} received error ${error}`);
     });
-    proxyReq.on('timeout', ()=>{
+    proxyReq.on('timeout', () => {
       logger.info(`${request.url} timed out`);
     });
-    proxyReq.on('drain', ()=>{
+    proxyReq.on('drain', () => {
       logger.info(`${request.url} drained out`);
     });
     request.pipe(proxyReq, {