ATB parameter in MV3

shared/js/background/atb.es6.js

@@ -8,6 +8,7 @@ const settings = require('./settings.es6')
 const parseUserAgentString = require('../shared-utils/parse-user-agent-string.es6')
 const load = require('./load.es6')
 const browserWrapper = require('./wrapper.es6')
+const { setOrUpdateATBdnrRule } = require('./declarative-net-request')
 
 const ATB_ERROR_COHORT = 'v1-1'
 const ATB_FORMAT_RE = /(v\d+-\d(?:[a-z_]{2})?)$/
@@ -40,6 +41,7 @@ const ATB = (() => {
             if (!atbSetting) {
                 atbSetting = ATB_ERROR_COHORT
                 settings.updateSetting('atb', ATB_ERROR_COHORT)
+                setOrUpdateATBdnrRule(ATB_ERROR_COHORT, regExpAboutPage)
                 errorParam = '&e=1'
             }
 
@@ -52,8 +54,10 @@ const ATB = (() => {
 
                 if (res.data.updateVersion) {
                     settings.updateSetting('atb', res.data.updateVersion)
+                    setOrUpdateATBdnrRule(res.data.updateVersion, regExpAboutPage)
                 } else if (atbSetting === ATB_ERROR_COHORT) {
                     settings.updateSetting('atb', res.data.version)
+                    setOrUpdateATBdnrRule(res.data.version, regExpAboutPage)
                 }
             })
         },
@@ -92,6 +96,7 @@ const ATB = (() => {
             const url = ddgAtbURL + randomValue + '&browser=' + parseUserAgentString().browser
             return load.JSONfromExternalFile(url).then((res) => {
                 settings.updateSetting('atb', res.data.version)
+                setOrUpdateATBdnrRule(res.data.version, regExpAboutPage)
             }, () => {
                 console.log('couldn\'t reach atb.js for initial server call, trying again')
                 numTries += 1
@@ -164,6 +169,7 @@ const ATB = (() => {
 
                     if (atb) {
                         settings.updateSetting('atb', atb)
+                        setOrUpdateATBdnrRule(atb, regExpAboutPage)
                     }
 
                     ATB.finalizeATB(params)

shared/js/background/declarative-net-request.js

@@ -15,7 +15,8 @@ import {
 } from '@duckduckgo/ddg2dnr/lib/utils'
 import {
     SERVICE_WORKER_INITIATED_ALLOWING_PRIORITY,
-    USER_ALLOWLISTED_PRIORITY
+    USER_ALLOWLISTED_PRIORITY,
+    ATB_PARAM_PRIORITY
 } from '@duckduckgo/ddg2dnr/lib/rulePriorities'
 
 export const SETTING_PREFIX = 'declarative_net_request-'
@@ -32,6 +33,7 @@ const ruleIdRangeByConfigName = {
 // only require one declarativeNetRequest rule, so hardcode the rule IDs here.
 export const USER_ALLOWLIST_RULE_ID = 20001
 export const SERVICE_WORKER_INITIATED_ALLOWING_RULE_ID = 20002
+export const ATB_PARAM_RULE_ID = 20003
 
 /**
  * A dummy etag rule is saved with the declarativeNetRequest rules generated for
@@ -341,6 +343,12 @@ export async function getMatchDetails (ruleId) {
         }
     }
 
+    if (ruleId === ATB_PARAM_RULE_ID) {
+        return {
+            type: 'atbParam'
+        }
+    }
+
     for (const [configName, [ruleIdStart, ruleIdEnd]]
         of Object.entries(ruleIdRangeByConfigName)) {
         if (ruleId >= ruleIdStart && ruleId <= ruleIdEnd) {
@@ -465,6 +473,43 @@ export async function ensureServiceWorkerInitiatedRequestException () {
     })
 }
 
+/**
+ * Creates a DNR rule for ATB parameters
+ * @param {string} atb
+ * @param {object} regexFilter
+ */
+export async function setOrUpdateATBdnrRule (atb, regexFilter) {
+    if (!(atb && regexFilter)) {
+        return
+    }
+
+    // format regex for json and dnr
+    regexFilter = regexFilter.toString().replace(/\\/g, '\\').replace(/^\//, '').replace(/\/$/, '')
+
+    const atbRule = generateDNRRule({
+        id: ATB_PARAM_RULE_ID,
+        priority: ATB_PARAM_PRIORITY,
+        actionType: 'redirect',
+        redirect: {
+            transform: {
+                queryTransform: {
+                    addOrReplaceParams: [{ key: 'atb', value: atb }]
+                }
+            }
+        },
+        resourceTypes: ['main_frame'],
+        requestDomains: ['duckduckgo.com'],
+        regexFilter
+    })
+
+    if (atbRule?.id === ATB_PARAM_RULE_ID) {
+        await chrome.declarativeNetRequest.updateDynamicRules({
+            removeRuleIds: [atbRule.id],
+            addRules: [atbRule]
+        })
+    }
+}
+
 if (browserWrapper.getManifestVersion() === 3) {
     tdsStorage.onUpdate('config', onConfigUpdate)
     tdsStorage.onUpdate('tds', onConfigUpdate)