From 8f608076b30939b050e97d33bd3b6d7d4c036e3f Mon Sep 17 00:00:00 2001
From: Anna Henningsen <anna@addaleax.net>
Date: Mon, 9 Oct 2023 18:03:58 +0200
Subject: [PATCH] chore(export-to-language): reokace usage of vm module
 COMPASS-7312

---
 package-lock.json                             |  6 ++-
 .../compass-export-to-language/package.json   |  4 +-
 .../count-aggregation-stages-in-string.js     | 45 -------------------
 .../count-aggregation-stages-in-string.ts     |  9 ++++
 4 files changed, 16 insertions(+), 48 deletions(-)
 delete mode 100644 packages/compass-export-to-language/src/modules/count-aggregation-stages-in-string.js
 create mode 100644 packages/compass-export-to-language/src/modules/count-aggregation-stages-in-string.ts

diff --git a/package-lock.json b/package-lock.json
index 75e064e9619..3b10e631d59 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -44163,7 +44163,6 @@
       "dependencies": {
         "@mongosh/node-runtime-worker-thread": "^2.0.0",
         "clipboard": "^2.0.6",
-        "electron": "^25.8.4",
         "kerberos": "^2.0.1",
         "keytar": "^7.9.0",
         "mongodb-client-encryption": "^6.0.0",
@@ -46016,7 +46015,8 @@
         "@mongodb-js/compass-logging": "^1.2.1",
         "@mongodb-js/compass-maybe-protect-connection-string": "^0.13.0",
         "@mongodb-js/mongodb-redux-common": "^2.0.12",
-        "bson-transpilers": "^2.0.4"
+        "bson-transpilers": "^2.0.4",
+        "ejson-shell-parser": "^1.2.4"
       },
       "devDependencies": {
         "@mongodb-js/eslint-config-compass": "^1.0.9",
@@ -46048,6 +46048,7 @@
         "@mongodb-js/compass-maybe-protect-connection-string": "^0.13.0",
         "@mongodb-js/mongodb-redux-common": "*",
         "bson-transpilers": "*",
+        "ejson-shell-parser": "^1.2.4",
         "react": "^17.0.2"
       }
     },
@@ -58959,6 +58960,7 @@
         "chai": "^4.3.6",
         "compass-preferences-model": "^2.15.0",
         "depcheck": "^1.4.1",
+        "ejson-shell-parser": "^1.2.4",
         "enzyme": "^3.11.0",
         "eslint": "^7.25.0",
         "hadron-app-registry": "^9.0.11",
diff --git a/packages/compass-export-to-language/package.json b/packages/compass-export-to-language/package.json
index b0ca4760b24..8c21b08efa0 100644
--- a/packages/compass-export-to-language/package.json
+++ b/packages/compass-export-to-language/package.json
@@ -61,6 +61,7 @@
     "@mongodb-js/compass-maybe-protect-connection-string": "^0.13.0",
     "@mongodb-js/mongodb-redux-common": "*",
     "bson-transpilers": "*",
+    "ejson-shell-parser": "^1.2.4",
     "react": "^17.0.2"
   },
   "dependencies": {
@@ -68,7 +69,8 @@
     "@mongodb-js/compass-logging": "^1.2.1",
     "@mongodb-js/compass-maybe-protect-connection-string": "^0.13.0",
     "@mongodb-js/mongodb-redux-common": "^2.0.12",
-    "bson-transpilers": "^2.0.4"
+    "bson-transpilers": "^2.0.4",
+    "ejson-shell-parser": "^1.2.4"
   },
   "devDependencies": {
     "@mongodb-js/eslint-config-compass": "^1.0.9",
diff --git a/packages/compass-export-to-language/src/modules/count-aggregation-stages-in-string.js b/packages/compass-export-to-language/src/modules/count-aggregation-stages-in-string.js
deleted file mode 100644
index 035ad3e7d72..00000000000
--- a/packages/compass-export-to-language/src/modules/count-aggregation-stages-in-string.js
+++ /dev/null
@@ -1,45 +0,0 @@
-import vm from 'vm';
-
-let dummySandbox;
-
-// We do want to report the number of stages in an aggregation
-// for telemetry, but we only receive them as a string from
-// the consumers of this package. Instead, we evaluate
-// the code in a dummy sandbox environment and just count
-// the number of stages in the result. This is a little inefficient,
-// but ultimately a simpler solution than pulling in a query
-// parser here.
-export function countAggregationStagesInString(str) {
-  if (!dummySandbox) {
-    dummySandbox = vm.createContext(Object.create(null), {
-      codeGeneration: { strings: false, wasm: false },
-      microtaskMode: 'afterEvaluate',
-    });
-    vm.runInContext(
-      [
-        'BSONRegExp',
-        'DBRef',
-        'Decimal128',
-        'Double',
-        'Int32',
-        'Long',
-        'Int64',
-        'MaxKey',
-        'MinKey',
-        'ObjectID',
-        'ObjectId',
-        'BSONSymbol',
-        'Timestamp',
-        'Code',
-        'Buffer',
-        'Binary',
-      ]
-        .map((name) => `function ${name}() {}`)
-        .join('\n'),
-      dummySandbox
-    );
-  }
-
-  return vm.runInContext('(' + str + ')', dummySandbox, { timeout: 100 })
-    .length;
-}
diff --git a/packages/compass-export-to-language/src/modules/count-aggregation-stages-in-string.ts b/packages/compass-export-to-language/src/modules/count-aggregation-stages-in-string.ts
new file mode 100644
index 00000000000..19608536ad3
--- /dev/null
+++ b/packages/compass-export-to-language/src/modules/count-aggregation-stages-in-string.ts
@@ -0,0 +1,9 @@
+import parseShellBSON, { ParseMode } from 'ejson-shell-parser';
+
+export function countAggregationStagesInString(source: string): number {
+  const parsed = parseShellBSON(source, { mode: ParseMode.Loose });
+  if (!Array.isArray(parsed)) {
+    throw new Error('Source expression is not an aggregation stage array');
+  }
+  return parsed.length;
+}