Bug Fix: Phishing Detection Dataset Discrepancies

Sources/PhishingDetection/PhishingDetectionDataStore.swift

@@ -136,29 +136,43 @@ public class PhishingDetectionDataStore: PhishingDetectionDataSaving {
     }
 
     private func loadHashPrefix() -> Set<String> {
-        guard let data = fileStorageManager.read(from: hashPrefixFilename) else { return dataProvider.loadEmbeddedHashPrefixes() }
+        guard let data = fileStorageManager.read(from: hashPrefixFilename) else {
+            return dataProvider.loadEmbeddedHashPrefixes()
+        }
         let decoder = JSONDecoder()
         do {
-            return Set(try decoder.decode(Set<String>.self, from: data))
+            if loadRevisionFromDisk() < dataProvider.embeddedRevision {
+                return dataProvider.loadEmbeddedHashPrefixes()
+            }
+            let onDiskHashPrefixes = Set(try decoder.decode(Set<String>.self, from: data))
+            return onDiskHashPrefixes
         } catch {
             Logger.phishingDetectionDataStore.error("Error decoding \(self.hashPrefixFilename): \(error.localizedDescription)")
             return dataProvider.loadEmbeddedHashPrefixes()
         }
     }
 
     private func loadFilterSet() -> Set<Filter> {
-        guard let data = fileStorageManager.read(from: filterSetFilename) else { return dataProvider.loadEmbeddedFilterSet() }
+        guard let data = fileStorageManager.read(from: filterSetFilename) else {
+            return dataProvider.loadEmbeddedFilterSet()
+        }
         let decoder = JSONDecoder()
         do {
-            return Set(try decoder.decode(Set<Filter>.self, from: data))
+            if loadRevisionFromDisk() < dataProvider.embeddedRevision {
+                return dataProvider.loadEmbeddedFilterSet()
+            }
+            let onDiskFilterSet = Set(try decoder.decode(Set<Filter>.self, from: data))
+            return onDiskFilterSet
         } catch {
             Logger.phishingDetectionDataStore.error("Error decoding \(self.filterSetFilename): \(error.localizedDescription)")
             return dataProvider.loadEmbeddedFilterSet()
         }
     }
 
-    private func loadRevision() -> Int {
-        guard let data = fileStorageManager.read(from: revisionFilename) else { return dataProvider.embeddedRevision }
+    private func loadRevisionFromDisk() -> Int {
+        guard let data = fileStorageManager.read(from: revisionFilename) else {
+            return dataProvider.embeddedRevision
+        }
         let decoder = JSONDecoder()
         do {
             return try decoder.decode(Int.self, from: data)
@@ -167,22 +181,39 @@ public class PhishingDetectionDataStore: PhishingDetectionDataSaving {
             return dataProvider.embeddedRevision
         }
     }
+
+    private func loadRevision() -> Int {
+        guard let data = fileStorageManager.read(from: revisionFilename) else {
+            return dataProvider.embeddedRevision
+        }
+        let decoder = JSONDecoder()
+        do {
+            let loadedRevision = try decoder.decode(Int.self, from: data)
+            if loadedRevision < dataProvider.embeddedRevision {
+                return dataProvider.embeddedRevision
+            }
+            return loadedRevision
+        } catch {
+            Logger.phishingDetectionDataStore.error("Error decoding \(self.revisionFilename): \(error.localizedDescription)")
+            return dataProvider.embeddedRevision
+        }
+    }
 }
 
 extension PhishingDetectionDataStore {
     public func saveFilterSet(set: Set<Filter>) {
-        writeFilterSet()
         self.filterSet = set
+        writeFilterSet()
     }
 
     public func saveHashPrefixes(set: Set<String>) {
-        writeHashPrefixes()
         self.hashPrefixes = set
+        writeHashPrefixes()
     }
 
     public func saveRevision(_ revision: Int) {
-        writeRevision()
         self.currentRevision = revision
+        writeRevision()
     }
 }
 

Tests/PhishingDetectionTests/PhishingDetectionDataStoreTests.swift

@@ -61,6 +61,60 @@
         XCTAssertEqual(actualHashPrefix, expectedHashPrefix)
     }
 
+    func testWhenEmbeddedRevisionNewerThanOnDisk_ThenLoadEmbedded() async {
+        let encoder = JSONEncoder()
+        // On Disk Data Setup
+        fileStorageManager.write(data: "1".utf8data, to: "revision.txt")
+        let onDiskFilterSet = Set([Filter(hashValue: "other", regex: "other")])
+        let filterSetData = try! encoder.encode(Array(onDiskFilterSet))
+        let onDiskHashPrefix = Set(["faffa"])
+        let hashPrefixData = try! encoder.encode(Array(onDiskHashPrefix))
+        fileStorageManager.write(data: filterSetData, to: "filterSet.json")
+        fileStorageManager.write(data: hashPrefixData, to: "hashPrefixes.json")
+
+        // Embedded Data Setup
+        mockDataProvider.embeddedRevision = 5
+        let embeddedFilterSet = Set([Filter(hashValue: "some", regex: "some")])
+        let embeddedHashPrefix = Set(["sassa"])
+        mockDataProvider.shouldReturnFilterSet(set: embeddedFilterSet)
+        mockDataProvider.shouldReturnHashPrefixes(set: embeddedHashPrefix)
+
+        let actualRevision = dataStore.currentRevision
+        let actualFilterSet = dataStore.filterSet
+        let actualHashPrefix = dataStore.hashPrefixes
+
+        XCTAssertEqual(actualFilterSet, embeddedFilterSet)
+        XCTAssertEqual(actualHashPrefix, embeddedHashPrefix)
+        XCTAssertEqual(actualRevision, 5)
+    }
+
+    func testWhenEmbeddedRevisionOlderThanOnDisk_ThenDontLoadEmbedded() async {
+        let encoder = JSONEncoder()
+        // On Disk Data Setup
+        fileStorageManager.write(data: "6".utf8data, to: "revision.txt")
+        let onDiskFilterSet = Set([Filter(hashValue: "other", regex: "other")])
+        let filterSetData = try! encoder.encode(Array(onDiskFilterSet))
+        let onDiskHashPrefix = Set(["faffa"])
+        let hashPrefixData = try! encoder.encode(Array(onDiskHashPrefix))
+        fileStorageManager.write(data: filterSetData, to: "filterSet.json")
+        fileStorageManager.write(data: hashPrefixData, to: "hashPrefixes.json")
+
+        // Embedded Data Setup
+        mockDataProvider.embeddedRevision = 1
+        let embeddedFilterSet = Set([Filter(hashValue: "some", regex: "some")])
+        let embeddedHashPrefix = Set(["sassa"])
+        mockDataProvider.shouldReturnFilterSet(set: embeddedFilterSet)
+        mockDataProvider.shouldReturnHashPrefixes(set: embeddedHashPrefix)
+
+        let actualRevision = dataStore.currentRevision
+        let actualFilterSet = dataStore.filterSet
+        let actualHashPrefix = dataStore.hashPrefixes
+
+        XCTAssertEqual(actualFilterSet, onDiskFilterSet)
+        XCTAssertEqual(actualHashPrefix, onDiskHashPrefix)
+        XCTAssertEqual(actualRevision, 6)
+    }
+
     func testWriteAndLoadData() async {
         // Get and write data
         let expectedHashPrefixes = Set(["aabb"])
@@ -93,4 +147,51 @@
             XCTFail("Failed to decode stored PhishingDetection data")
         }
     }
+
+    func testLazyLoadingDoesNotReturnStaleData() async {
+        clearDatasets()
+
+        // Set up initial data
+        let initialFilterSet = Set([Filter(hashValue: "initial", regex: "initial")])
+        let initialHashPrefixes = Set(["initialPrefix"])
+        mockDataProvider.shouldReturnFilterSet(set: initialFilterSet)
+        mockDataProvider.shouldReturnHashPrefixes(set: initialHashPrefixes)
+
+        // Access the lazy-loaded properties to trigger loading
+        let loadedFilterSet = dataStore.filterSet
+        let loadedHashPrefixes = dataStore.hashPrefixes
+
+        // Validate loaded data matches initial data
+        XCTAssertEqual(loadedFilterSet, initialFilterSet)
+        XCTAssertEqual(loadedHashPrefixes, initialHashPrefixes)
+
+        // Update in-memory data
+        let updatedFilterSet = Set([Filter(hashValue: "updated", regex: "updated")])
+        let updatedHashPrefixes = Set(["updatedPrefix"])
+        dataStore.saveFilterSet(set: updatedFilterSet)
+        dataStore.saveHashPrefixes(set: updatedHashPrefixes)
+
+        // Access lazy-loaded properties again
+        let reloadedFilterSet = dataStore.filterSet
+        let reloadedHashPrefixes = dataStore.hashPrefixes
+
+        // Validate reloaded data matches updated data
+        XCTAssertEqual(reloadedFilterSet, updatedFilterSet)
+        XCTAssertEqual(reloadedHashPrefixes, updatedHashPrefixes)
+
+        // Validate on-disk data is also updated
+        let storedFilterSetData = fileStorageManager.read(from: "filterSet.json")
+        let storedHashPrefixesData = fileStorageManager.read(from: "hashPrefixes.json")
+
+        let decoder = JSONDecoder()
+        if let storedFilterSet = try? decoder.decode(Set<Filter>.self, from: storedFilterSetData!),
+           let storedHashPrefixes = try? decoder.decode(Set<String>.self, from: storedHashPrefixesData!) {
+
+            XCTAssertEqual(storedFilterSet, updatedFilterSet)
+            XCTAssertEqual(storedHashPrefixes, updatedHashPrefixes)
+        } else {
+            XCTFail("Failed to decode stored PhishingDetection data after update")
+        }
+    }
+
 }