From 6fb5df97cc41d6162143ebcdc06e86e7002fb170 Mon Sep 17 00:00:00 2001
From: "yanlong.wang" <yanlong.wang@naiver.org>
Date: Thu, 1 Aug 2024 19:34:39 +0800
Subject: [PATCH] fix: abuse of flooding elements

---
 backend/functions/src/services/puppeteer.ts | 26 +++++++++++++++++----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/backend/functions/src/services/puppeteer.ts b/backend/functions/src/services/puppeteer.ts
index 1286a62..59e4468 100644
--- a/backend/functions/src/services/puppeteer.ts
+++ b/backend/functions/src/services/puppeteer.ts
@@ -50,6 +50,7 @@ export interface PageSnapshot {
     imgs?: ImgBrief[];
     pdfs?: string[];
     maxElemDepth?: number;
+    elemCount?: number;
     childFrames?: PageSnapshot[];
 }
 
@@ -117,14 +118,21 @@ function briefPDFs() {
         return x.src === 'about:blank' ? document.location.href : x.src;
     });
 }
-function getMaxDepthUsingTreeWalker(root) {
+function getMaxDepthAndCountUsingTreeWalker(root) {
   let maxDepth = 0;
   let currentDepth = 0;
+  let elementCount = 0;
 
-  const treeWalker = document.createTreeWalker(root, NodeFilter.SHOW_ELEMENT, null, false);
+  const treeWalker = document.createTreeWalker(
+    root,
+    NodeFilter.SHOW_ELEMENT,
+    (node) => (node.nodeName.toLowerCase() === 'svg') ? NodeFilter.FILTER_REJECT : NodeFilter.FILTER_ACCEPT,
+    false
+  );
 
   while (true) {
     maxDepth = Math.max(maxDepth, currentDepth);
+    elementCount++; // Increment the count for the current node
 
     if (treeWalker.firstChild()) {
       currentDepth++;
@@ -140,7 +148,10 @@ function getMaxDepthUsingTreeWalker(root) {
     }
   }
 
-  return maxDepth + 1;
+  return {
+    maxDepth: maxDepth + 1,
+    elementCount: elementCount
+  };
 }
 
 function giveSnapshot(stopActiveSnapshot) {
@@ -153,7 +164,7 @@ function giveSnapshot(stopActiveSnapshot) {
     } catch (err) {
         void 0;
     }
-
+    const domAnalysis = getMaxDepthAndCountUsingTreeWalker(document.documentElement);
     const r = {
         title: document.title,
         href: document.location.href,
@@ -162,7 +173,8 @@ function giveSnapshot(stopActiveSnapshot) {
         parsed: parsed,
         imgs: [],
         pdfs: briefPDFs(),
-        maxElemDepth: getMaxDepthUsingTreeWalker(document.documentElement)
+        maxElemDepth: domAnalysis.maxDepth,
+        elemCount: domAnalysis.elementCount,
     };
     if (parsed && parsed.content) {
         const elem = document.createElement('div');
@@ -478,6 +490,10 @@ document.addEventListener('load', handlePageLoad);
                 page.emit('abuse', { url, page, sn, reason: `DoS attack suspected: DOM tree too deep` });
                 return;
             }
+            if (s?.elemCount && s.elemCount > 15_000) {
+                page.emit('abuse', { url, page, sn, reason: `DoS attack suspected: too many DOM elements` });
+                return;
+            }
             snapshot = s;
             nextSnapshotDeferred.resolve(s);
             nextSnapshotDeferred = Defer();