From 6fd6495d4fa889882b45af844e38cba488c6bc31 Mon Sep 17 00:00:00 2001
From: Alvin Wang <Dogfalo@users.noreply.github.com>
Date: Thu, 5 Sep 2024 08:46:16 -0700
Subject: [PATCH] fix(autocapture): prevent potential memory exhaustion when
 generating combinations (#863)

* fix(autocapture): prevent potential memory exhaustion when generating combinations for CSS Selector
---
 packages/plugin-autocapture-browser/src/helpers.ts     | 1 +
 packages/plugin-autocapture-browser/src/libs/finder.ts | 7 +++++--
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/packages/plugin-autocapture-browser/src/helpers.ts b/packages/plugin-autocapture-browser/src/helpers.ts
index 764964d23..4883f067d 100644
--- a/packages/plugin-autocapture-browser/src/helpers.ts
+++ b/packages/plugin-autocapture-browser/src/helpers.ts
@@ -131,6 +131,7 @@ export const getSelector = (element: Element, logger?: Logger): string => {
   try {
     selector = finder(element, {
       className: (name: string) => name !== constants.AMPLITUDE_VISUAL_TAGGING_HIGHLIGHT_CLASS,
+      maxNumberOfTries: 1000,
     });
     return selector;
   } catch (error) {
diff --git a/packages/plugin-autocapture-browser/src/libs/finder.ts b/packages/plugin-autocapture-browser/src/libs/finder.ts
index cbaa89bc6..3fec33ded 100644
--- a/packages/plugin-autocapture-browser/src/libs/finder.ts
+++ b/packages/plugin-autocapture-browser/src/libs/finder.ts
@@ -132,10 +132,13 @@ function bottomUpSearch(
 }
 
 function findUniquePath(stack: Knot[][], fallback?: () => Path | null): Path | null {
-  const paths = sort(combinations(stack));
-  if (paths.length > config.threshold) {
+  // Check first the total number of combinations first since generating the combinations can cause memory exhaustion
+  const numCombinations = stack.reduce((acc, i) => acc * i.length, 1);
+  if (numCombinations > config.threshold) {
     return fallback ? fallback() : null;
   }
+
+  const paths = sort(combinations(stack));
   for (const candidate of paths) {
     if (unique(candidate)) {
       return candidate;