From fbe73b238737713f547c14c289a81a3aa3cc3fd1 Mon Sep 17 00:00:00 2001
From: Ben White <ben@posthog.com>
Date: Tue, 26 Sep 2023 09:25:48 +0200
Subject: [PATCH] Extended masking function to allow passing custom fn

---
 packages/rrweb-snapshot/src/snapshot.ts |  2 +-
 packages/rrweb-snapshot/src/types.ts    |  2 +-
 packages/rrweb/src/record/mutation.ts   |  5 ++++-
 packages/rrweb/src/utils.ts             | 27 ++++++++++++++++++++-----
 packages/rrweb/test/html/mask-text.html |  4 ++++
 packages/rrweb/test/integration.test.ts | 25 +++++++++++++++++++++++
 6 files changed, 57 insertions(+), 8 deletions(-)

diff --git a/packages/rrweb-snapshot/src/snapshot.ts b/packages/rrweb-snapshot/src/snapshot.ts
index 02619296c8..0963e6cfef 100644
--- a/packages/rrweb-snapshot/src/snapshot.ts
+++ b/packages/rrweb-snapshot/src/snapshot.ts
@@ -575,7 +575,7 @@ function serializeTextNode(
     needMaskingText(n, maskTextClass, maskTextSelector)
   ) {
     textContent = maskTextFn
-      ? maskTextFn(textContent)
+      ? maskTextFn(textContent, n.parentElement)
       : textContent.replace(/[\S]/g, '*');
   }
 
diff --git a/packages/rrweb-snapshot/src/types.ts b/packages/rrweb-snapshot/src/types.ts
index 9edb4dd6d4..e573dfc1e0 100644
--- a/packages/rrweb-snapshot/src/types.ts
+++ b/packages/rrweb-snapshot/src/types.ts
@@ -153,7 +153,7 @@ export type DataURLOptions = Partial<{
   quality: number;
 }>;
 
-export type MaskTextFn = (text: string) => string;
+export type MaskTextFn = (text: string, element: HTMLElement | null) => string;
 export type MaskInputFn = (text: string, element: HTMLElement) => string;
 
 export type KeepIframeSrcFn = (src: string) => boolean;
diff --git a/packages/rrweb/src/record/mutation.ts b/packages/rrweb/src/record/mutation.ts
index 097d1a8fd5..d3c33e70c0 100644
--- a/packages/rrweb/src/record/mutation.ts
+++ b/packages/rrweb/src/record/mutation.ts
@@ -30,6 +30,7 @@ import {
   isSerializedStylesheet,
   inDom,
   getShadowHost,
+  getElementFromNode,
 } from '../utils';
 
 type DoubleLinkedListNode = {
@@ -508,6 +509,8 @@ export default class MutationBuffer {
     switch (m.type) {
       case 'characterData': {
         const value = m.target.textContent;
+        const el = getElementFromNode(m.target)
+
         if (
           !isBlocked(m.target, this.blockClass, this.blockSelector, false) &&
           value !== m.oldValue
@@ -520,7 +523,7 @@ export default class MutationBuffer {
                 this.maskTextSelector,
               ) && value
                 ? this.maskTextFn
-                  ? this.maskTextFn(value)
+                  ? this.maskTextFn(value, el)
                   : value.replace(/[\S]/g, '*')
                 : value,
             node: m.target,
diff --git a/packages/rrweb/src/utils.ts b/packages/rrweb/src/utils.ts
index 604c8810e2..f476562037 100644
--- a/packages/rrweb/src/utils.ts
+++ b/packages/rrweb/src/utils.ts
@@ -215,6 +215,23 @@ export function getWindowWidth(): number {
   );
 }
 
+/**
+ * Returns the given node as an HTMLElement if it is one, otherwise the parent node as an HTMLElement
+ * @param node - node to check
+ * @returns HTMLElement or null
+ */
+
+export function getElementFromNode(node: Node | null,): HTMLElement | null {
+  if (!node) {
+    return null;
+  }
+  const el: HTMLElement | null =
+    node.nodeType === node.ELEMENT_NODE
+      ? (node as HTMLElement)
+      : node.parentElement;
+  return el
+}
+
 /**
  * Checks if the given element set to be blocked by rrweb
  * @param node - node to check
@@ -232,11 +249,11 @@ export function isBlocked(
   if (!node) {
     return false;
   }
-  const el: HTMLElement | null =
-    node.nodeType === node.ELEMENT_NODE
-      ? (node as HTMLElement)
-      : node.parentElement;
-  if (!el) return false;
+  const el = getElementFromNode(node)
+
+  if (!el) {
+    return false
+  }
 
   try {
     if (typeof blockClass === 'string') {
diff --git a/packages/rrweb/test/html/mask-text.html b/packages/rrweb/test/html/mask-text.html
index 2abaaaa511..135034b6af 100644
--- a/packages/rrweb/test/html/mask-text.html
+++ b/packages/rrweb/test/html/mask-text.html
@@ -16,5 +16,9 @@
         <div>mask3</div>
       </div>
     </div>
+
+    <p data-unmask-example="true">
+      unmask1
+    </p>
   </body>
 </html>
diff --git a/packages/rrweb/test/integration.test.ts b/packages/rrweb/test/integration.test.ts
index edfc8a97af..0fccd169f0 100644
--- a/packages/rrweb/test/integration.test.ts
+++ b/packages/rrweb/test/integration.test.ts
@@ -1170,6 +1170,31 @@ describe('record integration tests', function (this: ISuite) {
     assertSnapshot(snapshots);
   });
 
+  it('should unmask texts using maskTextFn', async () => {
+    const page: puppeteer.Page = await browser.newPage();
+    await page.goto('about:blank');
+    await page.setContent(
+      getHtml.call(this, 'mask-text.html', {
+        maskTextSelector: '*',
+        maskTextFn: (t: string, el: HTMLElement) => {
+          return el.matches('[data-unmask-example="true"]') ? t : t.replace(/[a-z]/g, '*');
+        },
+      }),
+    );
+
+    await page.type('#password', 'secr3t');
+
+    // Change type to text (simulate "show password")
+    await page.click('#show-password');
+    await page.type('#password', 'XY');
+    await page.click('#show-password');
+
+    const snapshots = (await page.evaluate(
+      'window.snapshots',
+    )) as eventWithTime[];
+    assertSnapshot(snapshots);
+  });
+
   it('can mask character data mutations', async () => {
     const page: puppeteer.Page = await browser.newPage();
     await page.goto('about:blank');