From b6e6691eaf04dd6eebfb636a1f0dff11eb1436a7 Mon Sep 17 00:00:00 2001
From: Fernando Bunn <bunn@idevzilla.com>
Date: Sun, 17 Nov 2024 23:40:02 +0000
Subject: [PATCH] Add DDG domain name check for fireproofing storage (#3549)

Task/Issue URL: https://app.asana.com/0/1204167627774280/1204888209490841/f

Description:
Add DDG domain name check for fireproofing storage
---
 .../FireproofingURLExtensions.swift           |  4 +-
 .../Tab/Services/WebsiteDataStore.swift       |  4 +-
 .../Tab/Services/WebsiteDataStoreTests.swift  | 37 ++++++++++++-------
 3 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/DuckDuckGo/Fireproofing/Extensions/FireproofingURLExtensions.swift b/DuckDuckGo/Fireproofing/Extensions/FireproofingURLExtensions.swift
index da4b9df331..7cc8ccab39 100644
--- a/DuckDuckGo/Fireproofing/Extensions/FireproofingURLExtensions.swift
+++ b/DuckDuckGo/Fireproofing/Extensions/FireproofingURLExtensions.swift
@@ -24,7 +24,7 @@ private typealias URLPatterns = [String: [NSRegularExpression]]
 
 extension URL {
 
-    static let cookieDomain = "duckduckgo.com"
+    static let duckduckgoDomain = "duckduckgo.com"
 
     private static let loginPattern = regex("login|sign-in|signin|session")
 
@@ -53,7 +53,7 @@ extension URL {
 
     var canFireproof: Bool {
         guard let host = self.host, self.navigationalScheme?.isHypertextScheme == true else { return false }
-        return (host != Self.cookieDomain)
+        return (host != Self.duckduckgoDomain)
     }
 
     var showFireproofStatus: Bool {
diff --git a/DuckDuckGo/Tab/Services/WebsiteDataStore.swift b/DuckDuckGo/Tab/Services/WebsiteDataStore.swift
index e4c167a284..3da5e9bc43 100644
--- a/DuckDuckGo/Tab/Services/WebsiteDataStore.swift
+++ b/DuckDuckGo/Tab/Services/WebsiteDataStore.swift
@@ -133,7 +133,7 @@ internal class WebCacheManager {
         let removableRecords = allRecords.filter { record in
             // For Local Storage, only remove records that *exactly match* the display name.
             // Subdomains or root domains should be excluded.
-            !fireproofDomains.fireproofDomains.contains(record.displayName)
+            !URL.duckduckgoDomain.contains(record.displayName) && !fireproofDomains.fireproofDomains.contains(record.displayName)
         }
 
         await websiteDataStore.removeData(ofTypes: WKWebsiteDataStore.allWebsiteDataTypesExceptCookies, for: removableRecords)
@@ -155,7 +155,7 @@ internal class WebCacheManager {
 
         // Don't clear fireproof domains
         let cookiesToRemove = cookies.filter { cookie in
-            !self.fireproofDomains.isFireproof(cookieDomain: cookie.domain) && ![URL.cookieDomain, SubscriptionCookieManager.cookieDomain].contains(cookie.domain)
+            !self.fireproofDomains.isFireproof(cookieDomain: cookie.domain) && ![URL.duckduckgoDomain, SubscriptionCookieManager.cookieDomain].contains(cookie.domain)
         }
 
         for cookie in cookiesToRemove {
diff --git a/UnitTests/Tab/Services/WebsiteDataStoreTests.swift b/UnitTests/Tab/Services/WebsiteDataStoreTests.swift
index 2b145f870b..45652fe8d0 100644
--- a/UnitTests/Tab/Services/WebsiteDataStoreTests.swift
+++ b/UnitTests/Tab/Services/WebsiteDataStoreTests.swift
@@ -20,6 +20,18 @@ import XCTest
 @testable import DuckDuckGo_Privacy_Browser
 
 final class WebCacheManagerTests: XCTestCase {
+    /// The webView is necessary to manage the shared state of WKWebsiteDataRecord
+    var webView: WKWebView?
+
+    override func setUp() {
+        super.setUp()
+        webView = WKWebView()
+    }
+
+    override func tearDown() {
+        webView = nil
+        super.tearDown()
+    }
 
     func testWhenCookiesHaveSubDomainsOnSubDomainsAndWildcardsThenAllCookiesRetained() {
         let logins = MockPreservedLogins(domains: [
@@ -89,7 +101,7 @@ final class WebCacheManagerTests: XCTestCase {
 
     }
 
-    func testWhenClearedThenDDGCookiesAreRetained() {
+    @MainActor func testWhenClearedThenDDGCookiesAndStorageAreRetained() async {
         let logins = MockPreservedLogins(domains: [
             "example.com"
         ])
@@ -104,16 +116,17 @@ final class WebCacheManagerTests: XCTestCase {
             MockDataRecord(recordName: "duckduckgo.com")
         ]
 
-        let expect = expectation(description: #function)
         let webCacheManager = WebCacheManager(fireproofDomains: logins, websiteDataStore: dataStore)
-        Task {
-            await webCacheManager.clear()
-            expect.fulfill()
-        }
-        wait(for: [expect], timeout: 30.0)
 
+        // Await the clear function directly
+        await webCacheManager.clear()
+
+        // Assertions after the async operation
         XCTAssertEqual(cookieStore.cookies.count, 1)
         XCTAssertEqual(cookieStore.cookies[0].domain, "duckduckgo.com")
+
+        XCTAssertEqual(dataStore.records.count, 1)
+        XCTAssertEqual(dataStore.records.first?.displayName, "duckduckgo.com")
     }
 
     func testWhenClearedThenCookiesForLoginsAreRetained() {
@@ -184,15 +197,11 @@ final class WebCacheManagerTests: XCTestCase {
             }
         }
 
-        func removeData(ofTypes dataTypes: Set<String>, for records: [WKWebsiteDataRecord]) async {
+        func removeData(ofTypes dataTypes: Set<String>, for recordsToRemove: [WKWebsiteDataRecord]) async {
             removeDataCalledCount += 1
 
-            // In the real implementation, records will be selectively removed or edited based on their Fireproof status. For simplicity in this test,
-            // only remove records if all data types are removed, so that we can tell whether records for given domains still exist in some form.
-            if dataTypes == WKWebsiteDataStore.allWebsiteDataTypes() {
-                self.records = records.filter {
-                    dataTypes == $0.dataTypes
-                }
+            self.records = self.records.filter { record in
+                !recordsToRemove.contains(where: { $0 == record && $0.dataTypes.isSubset(of: dataTypes)})
             }
         }