Malicious site protection navigation detection (#3707)

Task/Issue URL: https://app.asana.com/0/1206329551987282/1207151848931030
Tech Design URL: https://app.asana.com/0/1206329551987282/1207273224076495/f

**Description**:
This PR adds the navigation logic for detecting a malicious site and navigating to a special error page if the site is malicious.
The original idea in the tech design was to intercept the Request in `decidePolicyForNavigationAction` and check whether a site is malicious cancelling the request accordingly.
We noticed that the above approach increases the page load time of websites due to the logic check.
I opted for an approach where in `decidePolicyForNavigationAction` we start the detection task in parallel without waiting and in `decidePolicyForNavigationResponse` we evaluate the task’s result.
Another approach I thought of was to perform the logic in the background in `didStartProvisionalNavigation`. The problem with this approach is that is called only for navigation that starts from the main frame so it
would not be possible to intercept malicious iFrame URLs.

DuckDuckGo/AppLifecycle/AppStateTransitions.swift

@@ -92,7 +92,7 @@ extension Resuming {
 
 }
 
-extension Testing {
+extension AppTesting {
 
     func apply(event: AppEvent) -> any AppState { self }
 

DuckDuckGo/AppLifecycle/AppStates/Testing.swift → DuckDuckGo/AppLifecycle/AppStates/AppTesting.swift

@@ -1,5 +1,5 @@
 //
-//  Testing.swift
+//  AppTesting.swift
 //  DuckDuckGo
 //
 //  Copyright © 2024 DuckDuckGo. All rights reserved.
@@ -21,7 +21,7 @@ import Core
 import UIKit
 
 @MainActor
-struct Testing: AppState {
+struct AppTesting: AppState {
 
     init(application: UIApplication) {
         Pixel.isDryRun = true

DuckDuckGo/MaliciousSiteProtection/MaliciousSiteProtectionManager.swift

@@ -18,38 +18,20 @@
 //
 
 import Foundation
+import MaliciousSiteProtection
 
 final class MaliciousSiteProtectionManager: MaliciousSiteDetecting {
 
     func evaluate(_ url: URL) async -> ThreatKind? {
         try? await Task.sleep(interval: 0.3)
-        return .none
-    }
-
-}
-
-// MARK: - To Remove
-
-// These entities are copied from BSK and they will be used to mock the library
-import SpecialErrorPages
-
-protocol MaliciousSiteDetecting {
-    func evaluate(_ url: URL) async -> ThreatKind?
-}
-
-public enum ThreatKind: String, CaseIterable, CustomStringConvertible {
-    public var description: String { rawValue }
-
-    case phishing
-    case malware
-}
-
-public extension ThreatKind {
 
-    var errorPageType: SpecialErrorKind {
-        switch self {
-        case .malware: .phishing // WIP in BSK
-        case .phishing: .phishing
+        switch url.absoluteString {
+        case "http://privacy-test-pages.site/security/badware/phishing.html":
+            return .phishing
+        case "http://privacy-test-pages.site/security/badware/malware.html":
+            return .malware
+        default:
+            return .none
         }
     }
 

DuckDuckGo/SpecialErrorPage/Model/SpecialErrorModel.swift

@@ -19,6 +19,7 @@
 
 import Foundation
 import SpecialErrorPages
+import WebKit
 
 struct SpecialErrorModel: Equatable {
     let url: URL
@@ -29,3 +30,8 @@ struct SSLSpecialError {
     let type: SSLErrorType
     let error: SpecialErrorModel
 }
+
+struct MaliciousSiteDetectionNavigationResponse: Equatable {
+    let navigationAction: WKNavigationAction
+    let errorData: SpecialErrorData
+}

DuckDuckGo/SpecialErrorPage/SpecialErrorPageInterfaces/SpecialErrorPageActionHandler.swift

@@ -18,6 +18,7 @@
 //
 
 import Foundation
+import SpecialErrorPages
 
 /// A type that defines actions for handling special error pages.
 ///
@@ -26,11 +27,29 @@ import Foundation
 /// advanced information related to the error.
 protocol SpecialErrorPageActionHandler {
     /// Handles the action of navigating to the site associated with the error page
+    /// - Parameters:
+    ///  - url: The URL that the user wants to visit.
+    ///  - errorData: The special error information.
+    @MainActor
+    func visitSite(url: URL, errorData: SpecialErrorData)
+
+    /// Handles the action of navigating to the site associated with the error page
+    @MainActor
     func visitSite()
 
     /// Handles the action of leaving the site associated with the error page
+    @MainActor
     func leaveSite()
 
     /// Handles the action of requesting more detailed information about the error
+    @MainActor
     func advancedInfoPresented()
 }
+
+extension SpecialErrorPageActionHandler {
+
+    func visitSite(url: URL, errorData: SpecialErrorData) { }
+
+    func visitSite() { }
+
+}

DuckDuckGo/SpecialErrorPage/SpecialErrorPageInterfaces/SpecialErrorPageContextHandling.swift

@@ -32,6 +32,9 @@ protocol SpecialErrorPageContextHandling: AnyObject {
     /// The URL that failed to load, if any.
     var failedURL: URL? { get }
 
+    /// A boolean value indicating whether the WebView request requires showing a special error page.
+    var isSpecialErrorPageRequest: Bool { get }
+
     /// Attaches a web view to the special error page handling.
     func attachWebView(_ webView: WKWebView)
 

DuckDuckGo/SpecialErrorPage/SpecialErrorPageInterfaces/WebViewNavigationHandling.swift

@@ -19,6 +19,7 @@
 
 import Foundation
 import WebKit
+import MaliciousSiteProtection
 
 // MARK: - WebViewNavigation
 
@@ -41,15 +42,24 @@ protocol WebViewNavigationHandling: AnyObject {
     /// - Parameters:
     ///   - navigationAction: Details about the action that triggered the navigation request.
     ///   - webView: The web view from which the navigation request began.
-    /// - Returns: A Boolean value that indicates whether the navigation action was handled.
-    func handleSpecialErrorNavigation(navigationAction: WKNavigationAction, webView: WKWebView) async -> Bool
+    @MainActor
+    func handleDecidePolicy(for navigationAction: WKNavigationAction, webView: WKWebView)
+
+    /// Decides whether to to navigate to new content after the response to the navigation request is known or cancel the navigation and show a special error page based on the specified action information.
+    /// - Parameters:
+    ///   - navigationResponse: Descriptive information about the navigation response.
+    ///   - webView: The web view from which the navigation request began.
+    /// - Returns: A Boolean value that indicates whether to cancel or allow the navigation.
+    @MainActor
+    func handleDecidePolicy(for navigationResponse: WKNavigationResponse, webView: WKWebView) async -> Bool
 
     /// Handles authentication challenges received by the web view.
     ///
     /// - Parameters:
     ///   - webView: The web view that receives the authentication challenge.
     ///   - challenge: The authentication challenge.
     ///   - completionHandler: A completion handler block to execute with the response.
+    @MainActor
     func handleWebView(_ webView: WKWebView, didReceive challenge: URLAuthenticationChallenge, completionHandler: @escaping (URLSession.AuthChallengeDisposition, URLCredential?) -> Void)
 
     /// Handles failures during provisional navigation.
@@ -58,12 +68,14 @@ protocol WebViewNavigationHandling: AnyObject {
     ///   - webView: The `WKWebView` instance that failed the navigation.
     ///   - navigation: The navigation object for the operation.
     ///   - error: The error that occurred.
+    @MainActor
     func handleWebView(_ webView: WKWebView, didFailProvisionalNavigation navigation: WebViewNavigation, withError error: NSError)
 
     /// Handles the successful completion of a navigation in the web view.
     ///
     /// - Parameters:
     ///   - webView: The web view that loaded the content.
     ///   - navigation: The navigation object that finished.
+    @MainActor
     func handleWebView(_ webView: WKWebView, didFinish navigation: WebViewNavigation)
 }

DuckDuckGo/SpecialErrorPage/SpecialErrorPageNavigationHandler+MaliciousSite.swift

@@ -22,26 +22,49 @@ import BrowserServicesKit
 import Core
 import SpecialErrorPages
 import WebKit
+import MaliciousSiteProtection
 
 enum MaliciousSiteProtectionNavigationResult: Equatable {
-    case navigationHandled(SpecialErrorModel)
+    case navigationHandled(NavigationType)
     case navigationNotHandled
+
+    enum NavigationType: Equatable {
+        case mainFrame(MaliciousSiteDetectionNavigationResponse)
+        case iFrame(maliciousURL: URL, error: SpecialErrorData)
+    }
 }
 
 protocol MaliciousSiteProtectionNavigationHandling: AnyObject {
-    /// Decides whether to cancel navigation to prevent opening the YouTube app from the web view.
+    /// Creates a task for detecting malicious sites based on the provided navigation action.
     ///
     /// - Parameters:
-    ///   - navigationAction: The navigation action to evaluate.
-    ///   - webView: The web view where navigation is occurring.
-    /// - Returns: `true` if the navigation should be canceled, `false` otherwise.
-    func handleMaliciousSiteProtectionNavigation(for navigationAction: WKNavigationAction, webView: WKWebView) async -> MaliciousSiteProtectionNavigationResult
+    ///   - navigationAction: The `WKNavigationAction` object that contains information about
+    ///     the navigation event.
+    ///   - webView: The web view from which the navigation request began.
+    @MainActor
+    func makeMaliciousSiteDetectionTask(for navigationAction: WKNavigationAction, webView: WKWebView)
+
+    /// Retrieves a task for detecting malicious sites based on the provided navigation response.
+    ///
+    /// - Parameters:
+    ///   - navigationResponse: The `WKNavigationResponse` object that contains information about
+    ///     the navigation event.
+    ///   - webView: The web view from which the navigation request began.
+    /// - Returns: A `Task<MaliciousSiteProtectionNavigationResult, Never>?` that represents the
+    ///   asynchronous operation for detecting malicious sites. If the task cannot be found,
+    ///   the function returns `nil`.
+    @MainActor
+    func getMaliciousSiteDectionTask(for navigationResponse: WKNavigationResponse, webView: WKWebView) -> Task<MaliciousSiteProtectionNavigationResult, Never>?
 }
 
 final class MaliciousSiteProtectionNavigationHandler {
     private let maliciousSiteProtectionManager: MaliciousSiteDetecting
     private let storageCache: StorageCache
 
+    @MainActor private(set) var maliciousURLExemptions: [URL: ThreatKind] = [:]
+    @MainActor private(set) var bypassedMaliciousSiteThreatKind: ThreatKind?
+    @MainActor private(set) var maliciousSiteDetectionTasks: [URL: Task<MaliciousSiteProtectionNavigationResult, Never>] = [:]
+
     init(
         maliciousSiteProtectionManager: MaliciousSiteDetecting = MaliciousSiteProtectionManager(),
         storageCache: StorageCache = AppDependencyProvider.shared.storageCache
@@ -56,10 +79,50 @@ final class MaliciousSiteProtectionNavigationHandler {
 extension MaliciousSiteProtectionNavigationHandler: MaliciousSiteProtectionNavigationHandling {
 
     @MainActor
-    func handleMaliciousSiteProtectionNavigation(for navigationAction: WKNavigationAction, webView: WKWebView) async -> MaliciousSiteProtectionNavigationResult {
-        // Implement logic to use `maliciousSiteProtectionManager.evaluate(url)`
-        // Return navigationNotHandled for the time being
-        return .navigationNotHandled
+    func makeMaliciousSiteDetectionTask(for navigationAction: WKNavigationAction, webView: WKWebView) {
+
+        guard
+            let url = navigationAction.request.url,
+            url != URL(string: "about:blank")
+        else {
+            return
+        }
+
+        handleMaliciousExemptions(for: navigationAction.navigationType, url: url)
+
+        guard !shouldBypassMaliciousSiteProtection(for: url) else {
+            return
+        }
+
+        let threatDetectionTask: Task<MaliciousSiteProtectionNavigationResult, Never> = Task {
+            guard let threatKind = await maliciousSiteProtectionManager.evaluate(url) else {
+                return .navigationNotHandled
+            }
+
+            if navigationAction.isTargetingMainFrame {
+                let errorData = SpecialErrorData.maliciousSite(kind: threatKind, url: url)
+                let response = MaliciousSiteDetectionNavigationResponse(navigationAction: navigationAction, errorData: errorData)
+                return .navigationHandled(.mainFrame(response))
+            } else {
+                // Extract the URL of the source frame (the iframe) that initiated the navigation action
+                let iFrameTopURL = navigationAction.sourceFrame.safeRequest?.url ?? url
+                let errorData = SpecialErrorData.maliciousSite(kind: threatKind, url: iFrameTopURL)
+                return .navigationHandled(.iFrame(maliciousURL: url, error: errorData))
+            }
+        }
+
+        maliciousSiteDetectionTasks[url] = threatDetectionTask
+    }
+
+    @MainActor
+    func getMaliciousSiteDectionTask(for navigationResponse: WKNavigationResponse, webView: WKWebView) -> Task<MaliciousSiteProtectionNavigationResult, Never>? {
+
+        guard let url = navigationResponse.response.url else {
+            assertionFailure("Could not find Malicious Site Detection Task for URL")
+            return nil
+        }
+
+        return maliciousSiteDetectionTasks.removeValue(forKey: url)
     }
 
 }
@@ -68,7 +131,10 @@ extension MaliciousSiteProtectionNavigationHandler: MaliciousSiteProtectionNavig
 
 extension MaliciousSiteProtectionNavigationHandler: SpecialErrorPageActionHandler {
 
-    func visitSite() {
+    func visitSite(url: URL, errorData: SpecialErrorData) {
+        maliciousURLExemptions[url] = errorData.threatKind
+        bypassedMaliciousSiteThreatKind = errorData.threatKind
+
         // Fire Pixel
     }
 
@@ -81,3 +147,36 @@ extension MaliciousSiteProtectionNavigationHandler: SpecialErrorPageActionHandle
     }
 
 }
+
+// MARK: - Private
+
+private extension MaliciousSiteProtectionNavigationHandler {
+
+    @MainActor
+    func handleMaliciousExemptions(for navigationType: WKNavigationType, url: URL) {
+        // TODO: check storing redirects
+        // Re-set the flag every time we load a web page
+        bypassedMaliciousSiteThreatKind = maliciousURLExemptions[url]
+    }
+
+    @MainActor
+    func shouldBypassMaliciousSiteProtection(for url: URL) -> Bool {
+        bypassedMaliciousSiteThreatKind != .none || url.isDuckDuckGo || url.isDuckURLScheme
+    }
+
+}
+
+// MARK: - Helpers
+
+private extension SpecialErrorData {
+
+    var threatKind: ThreatKind? {
+        switch self {
+        case .ssl:
+            return nil
+        case let .maliciousSite(threatKind, _):
+            return threatKind
+        }
+    }
+
+}

DuckDuckGo/SpecialErrorPage/SpecialErrorPageNavigationHandler+SSL.swift

@@ -71,7 +71,11 @@ extension SSLErrorPageNavigationHandler: SSLSpecialErrorPageNavigationHandling {
             return nil
         }
 
-        let errorData = SpecialErrorData.ssl(type: errorType, domain: host, eTldPlus1: storageCache.tld.eTLDplus1(host))
+        let errorData = SpecialErrorData.ssl(
+            type: errorType,
+            domain: host,
+            eTldPlus1: storageCache.tld.eTLDplus1(host)
+        )
 
         return SSLSpecialError(type: errorType, error: SpecialErrorModel(url: failedURL, errorData: errorData))
     }