chore: improve rewriting and checking

Author: Justineo

rewrite/figma.ts

@@ -1,9 +1,12 @@
-import type { Group } from './config'
 import { GROUPS } from './config'
+import { applyGroups } from './shared'
 
 async function rewriteScript() {
   const current = document.currentScript as HTMLScriptElement
-  const src = current.src
+  const src = current?.src
+  if (!current || !src) {
+    return
+  }
 
   const desc = Object.getOwnPropertyDescriptor(Document.prototype, 'currentScript')
 
@@ -14,39 +17,17 @@ async function rewriteScript() {
     current.replaceWith(script)
   }
 
-  function applyGroup(content: string, group: Group) {
-    const markers = group.markers || []
-
-    if (!markers.every((marker) => content.includes(marker))) {
-      return content
-    }
-
-    let out = content
-    for (const { pattern, replacer } of group.replacements) {
-      if (typeof pattern === 'string') {
-        // @ts-ignore
-        out = out.replaceAll(pattern, replacer)
-      } else {
-        // @ts-ignore
-        out = out.replace(pattern, replacer)
-      }
-    }
-    return out
-  }
-
   try {
-    const original = await (await fetch(src)).text()
-    let content = original
+    const response = await fetch(src)
+    const original = await response.text()
 
-    for (const group of GROUPS) {
-      content = applyGroup(content, group)
-    }
+    const { content: afterRules, changed } = applyGroups(original, GROUPS)
 
-    if (content !== original) {
+    if (changed) {
       console.log(`Rewrote script: ${src}`)
     }
 
-    content = content.replaceAll('delete window.figma', 'window.figma = undefined')
+    const content = afterRules.replaceAll('delete window.figma', 'window.figma = undefined')
 
     Object.defineProperty(document, 'currentScript', {
       configurable: true,

rewrite/shared.ts

@@ -0,0 +1,46 @@
+import type { Group, Replacement } from './config'
+
+export function groupMatches(content: string, group: Group) {
+  const markers = group.markers || []
+  return markers.every((marker) => content.includes(marker))
+}
+
+export function replacementWouldChange(content: string, replacement: Replacement) {
+  if (typeof replacement.pattern === 'string') {
+    return content.includes(replacement.pattern)
+  }
+  const re = new RegExp(replacement.pattern.source, replacement.pattern.flags)
+  return re.test(content)
+}
+
+export function applyGroups(content: string, groups: Group[]) {
+  let out = content
+  for (const group of groups) {
+    if (!groupMatches(out, group)) {
+      continue
+    }
+    for (const { pattern, replacer } of group.replacements) {
+      if (typeof pattern === 'string') {
+        // @ts-ignore
+        out = out.replaceAll(pattern, replacer)
+      } else {
+        // @ts-ignore
+        out = out.replace(pattern, replacer)
+      }
+    }
+  }
+  return { content: out, changed: out !== content }
+}
+
+export function analyze(content: string, groups: Group[]) {
+  for (const group of groups) {
+    if (!groupMatches(content, group)) {
+      continue
+    }
+    if (group.replacements.some((r) => replacementWouldChange(content, r))) {
+      return { matched: true, rewritable: true }
+    }
+    return { matched: true, rewritable: false }
+  }
+  return { matched: false, rewritable: false }
+}

scripts/check-rewrite.ts

@@ -1,10 +1,9 @@
-import { matchFile, REWRITE_PATTERN } from '@/rewrite/config'
+import { GROUPS } from '@/rewrite/config'
+import { analyze } from '@/rewrite/shared'
 import { chromium } from 'playwright-chromium'
-
 import rules from '../public/rules/figma.json'
 
 const redirectRule = rules.find((rule) => rule.action.type === 'redirect')
-
 const ASSETS_PATTERN = new RegExp(redirectRule?.condition?.regexFilter || /a^/)
 const MAX_RETRIES = 3
 
@@ -16,16 +15,16 @@ async function runCheck() {
 
   try {
     page.on('response', async (response) => {
-      if (response.request().resourceType() === 'script') {
-        const url = response.url()
-        if (!ASSETS_PATTERN.test(url)) {
-          return
-        }
-
-        const content = await response.text()
-        scripts.push({ url, content })
-        console.log(`Captured script: ${url}`)
+      if (response.request().resourceType() !== 'script') {
+        return
+      }
+      const url = response.url()
+      if (!ASSETS_PATTERN.test(url)) {
+        return
       }
+      const content = await response.text()
+      scripts.push({ url, content })
+      console.log(`Captured script: ${url}`)
     })
 
     try {
@@ -43,7 +42,7 @@ async function runCheck() {
     try {
       await page.waitForURL(/^(?!.*login).*$/, { timeout: 10000 })
       console.log('Logged in successfully.')
-    } catch (error) {
+    } catch {
       console.error('Login failed. Please check your credentials.')
       return false
     }
@@ -53,16 +52,20 @@ async function runCheck() {
 
     let matched: string | null = null
     let rewritable = false
-    scripts.forEach(({ url, content }) => {
-      if (matchFile(url, content)) {
-        matched = url
-        console.log(`Matched script: ${url}`)
-        if (REWRITE_PATTERN.test(content)) {
-          rewritable = true
-          console.log(`Rewritable script: ${url}`)
-        }
+
+    for (const { url, content } of scripts) {
+      const { matched: m, rewritable: r } = analyze(content, GROUPS)
+      if (!m) {
+        continue
       }
-    })
+      matched = url
+      console.log(`Matched script: ${url}`)
+      if (r) {
+        rewritable = true
+        console.log(`Rewritable script: ${url}`)
+        break
+      }
+    }
 
     if (!matched) {
       console.log('❌ No matched script found.')
@@ -72,11 +75,11 @@ async function runCheck() {
     console.log(`✅ Matched script: ${matched}`)
 
     if (!rewritable) {
-      console.log(`❌ Rewrite pattern not found.`)
+      console.log('❌ Rewrite pattern not found.')
       return false
     }
 
-    console.log(`✅ Rewrite pattern found.`)
+    console.log('✅ Rewrite pattern found.')
     return true
   } finally {
     await browser.close()
@@ -90,14 +93,12 @@ async function main() {
     }
 
     const success = await runCheck()
-
     if (success) {
       process.exit(0)
     }
 
     if (attempt < MAX_RETRIES) {
       console.log(`Attempt ${attempt} failed. Waiting before retry...`)
-      // Wait a bit before retrying
       await new Promise((resolve) => setTimeout(resolve, 3000))
     }
   }