Keep DuckDuckGo Settings after Fire Button (#2444)

duckduckgo · Feb 6, 2024 · cef5f51 · cef5f51
1 parent 8a39dc9
commit cef5f51
Show file tree

Hide file tree

Showing 5 changed files with 61 additions and 15 deletions.
diff --git a/Core/CookieStorage.swift b/Core/CookieStorage.swift
@@ -114,7 +114,7 @@ public class CookieStorage {
         }
 
         persistedCookiesByDomain.keys.forEach {
-            guard $0 != "duckduckgo.com" else { return } // DDG cookies are for SERP settings only
+            guard !URL.isDuckDuckGo(domain: $0) else { return } // DDG cookies are for SERP settings only
 
             if !preservedLogins.isAllowed(cookieDomain: $0) {
                 persistedCookiesByDomain.removeValue(forKey: $0)

diff --git a/Core/WebCacheManager.swift b/Core/WebCacheManager.swift
@@ -56,10 +56,6 @@ extension WebCacheManagerDataStore {
 }
 
 public class WebCacheManager {
-
-    private struct Constants {
-        static let cookieDomainsToPreserve = ["duckduckgo.com", "surveys.duckduckgo.com"]
-    }
 
     public static var shared = WebCacheManager()
 
@@ -223,7 +219,7 @@ public class WebCacheManager {
 
         func keep(_ cookie: HTTPCookie) -> Bool {
             return logins.isAllowed(cookieDomain: cookie.domain) ||
-                Constants.cookieDomainsToPreserve.contains(cookie.domain)
+                URL.isDuckDuckGo(domain: cookie.domain)
         }
 
         let dataStore = WKWebsiteDataStore.default()
@@ -270,7 +266,7 @@ public class WebCacheManager {
                     // Remove legacy HTTPCookieStorage cookies
                     let storageCookies = HTTPCookieStorage.shared.cookies ?? []
                     let storageCookiesToRemove = storageCookies.filter {
-                        !logins.isAllowed(cookieDomain: $0.domain) && !Constants.cookieDomainsToPreserve.contains($0.domain)
+                        !logins.isAllowed(cookieDomain: $0.domain) && !URL.isDuckDuckGo(domain: $0.domain)
                     }
 
                     let protectedStorageCookiesCount = storageCookies.count - storageCookiesToRemove.count
@@ -389,7 +385,7 @@ extension WKWebsiteDataStore: WebCacheManagerDataStore {
     public func preservedCookies(_ preservedLogins: PreserveLogins) async -> [HTTPCookie] {
         let allCookies = await self.httpCookieStore.allCookies()
         return allCookies.filter {
-            preservedLogins.isAllowed(cookieDomain: $0.domain)
+            URL.isDuckDuckGo(domain: $0.domain) || preservedLogins.isAllowed(cookieDomain: $0.domain)
         }
     }
 

diff --git a/DuckDuckGoTests/CookieStorageTests.swift b/DuckDuckGoTests/CookieStorageTests.swift
@@ -51,6 +51,12 @@ public class CookieStorageTests: XCTestCase {
 
         XCTAssertEqual(2, storage.cookies.count)
 
+        storage.updateCookies([
+            make("usedev1.duckduckgo.com", name: "x", value: "1"),
+        ], keepingPreservedLogins: logins)
+
+        XCTAssertEqual(3, storage.cookies.count)
+
     }
 
     func testWhenUpdatedThenCookiesWithFutureExpirationAreNotRemoved() {

diff --git a/DuckDuckGoTests/FireButtonReferenceTests.swift b/DuckDuckGoTests/FireButtonReferenceTests.swift
@@ -44,7 +44,7 @@ final class FireButtonReferenceTests: XCTestCase {
         return url.host!
     }
 
-    func testCookieStorage() {
+    func testClearData() async throws {
         let preservedLogins = PreserveLogins.shared
         preservedLogins.clearAll()
 
@@ -59,12 +59,54 @@ final class FireButtonReferenceTests: XCTestCase {
         }
 
         let cookieStorage = CookieStorage()
+        let idManager = DataStoreIdManager()
         for test in referenceTests {
-            guard let cookie = cookie(for: test) else {
-                XCTFail("Cookie should exist for test \(test.name)")
-                return
+            let cookie = try XCTUnwrap(cookie(for: test))
+
+            // Set directly to avoid logic to remove non-preserved cookies
+            cookieStorage.cookies = [
+                cookie
+            ]
+
+            idManager.allocateNewContainerId()
+            await withCheckedContinuation { continuation in
+                WebCacheManager.shared.clear(cookieStorage: cookieStorage, logins: preservedLogins, dataStoreIdManager: idManager) {
+                    continuation.resume()
+                }
             }
 
+            let testCookie = cookieStorage.cookies.filter { $0.name == test.cookieName }.first
+
+            if test.expectCookieRemoved {
+                XCTAssertNil(testCookie, "Cookie should not exist for test: \(test.name)")
+            } else {
+                XCTAssertNotNil(testCookie, "Cookie should exist for test: \(test.name)")
+            }
+
+            // Reset cache
+            cookieStorage.cookies = []
+        }
+
+    }
+
+    func testCookieStorage() throws {
+        let preservedLogins = PreserveLogins.shared
+        preservedLogins.clearAll()
+
+        for site in testData.fireButtonFireproofing.fireproofedSites {
+            let sanitizedSite = sanitizedSite(site)
+            os_log("Adding %s to fireproofed sites", sanitizedSite)
+            preservedLogins.addToAllowed(domain: sanitizedSite)
+        }
+
+        let referenceTests = testData.fireButtonFireproofing.tests.filter {
+            $0.exceptPlatforms.contains("ios-browser") == false
+        }
+
+        let cookieStorage = CookieStorage()
+        for test in referenceTests {
+            let cookie = try XCTUnwrap(cookie(for: test))
+
             cookieStorage.updateCookies([
                 cookie
             ], keepingPreservedLogins: preservedLogins)

diff --git a/DuckDuckGoTests/WebCacheManagerTests.swift b/DuckDuckGoTests/WebCacheManagerTests.swift
@@ -128,7 +128,8 @@ class WebCacheManagerTests: XCTestCase {
 
         let dataStore = MockDataStore()
         let cookieStore = MockHTTPCookieStore(cookies: [
-            .make(domain: "duckduckgo.com")
+            .make(domain: "duckduckgo.com"),
+            .make(domain: "subdomain.duckduckgo.com")
         ])
 
         dataStore.cookieStore = cookieStore
@@ -139,8 +140,9 @@ class WebCacheManagerTests: XCTestCase {
         }
         wait(for: [expect], timeout: 5.0)
 
-        XCTAssertEqual(cookieStore.cookies.count, 1)
-        XCTAssertEqual(cookieStore.cookies[0].domain, "duckduckgo.com")
+        XCTAssertEqual(cookieStore.cookies.count, 2)
+        XCTAssertTrue(cookieStore.cookies.contains(where: { $0.domain == "duckduckgo.com" }))
+        XCTAssertTrue(cookieStore.cookies.contains(where: { $0.domain == "subdomain.duckduckgo.com" }))
     }
 
     @MainActor