Validate duck.ai URL

duckduckgo · Dec 18, 2024 · 1472616 · 1472616
1 parent cf2d339
commit 1472616
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 3 deletions.
diff --git a/LocalPackages/AIChat/Sources/AIChat/AIChatWebViewController.swift b/LocalPackages/AIChat/Sources/AIChat/AIChatWebViewController.swift
@@ -109,7 +109,7 @@ extension AIChatWebViewController: WKNavigationDelegate {
 
     func webView(_ webView: WKWebView, decidePolicyFor navigationAction: WKNavigationAction) async -> WKNavigationActionPolicy {
         if let url = navigationAction.request.url {
-            if url == chatModel.aiChatURL || navigationAction.targetFrame?.isMainFrame == false {
+            if url.isDuckAIURL || navigationAction.targetFrame?.isMainFrame == false {
                 return .allow
             } else {
                 delegate?.aiChatWebViewController(self, didRequestToLoad: url)

diff --git a/LocalPackages/AIChat/Sources/AIChat/URL+Extension.swift b/LocalPackages/AIChat/Sources/AIChat/URL+Extension.swift
@@ -32,4 +32,17 @@ extension URL {
         urlComponents.queryItems = queryItems
         return urlComponents.url ?? self
     }
+
+    var isDuckAIURL: Bool {
+        guard let host = self.host, host == "duckduckgo.com" else {
+            return false
+        }
+
+        guard let urlComponents = URLComponents(url: self, resolvingAgainstBaseURL: false),
+              let queryItems = urlComponents.queryItems else {
+            return false
+        }
+
+        return queryItems.contains { $0.name == "ia" && $0.value == "chat" }
+    }
 }
diff --git a/...ages/AIChat/Tests/URLQueryItemTests.swift → ...ages/AIChat/Tests/URLExtensionTests.swift b/...ages/AIChat/Tests/URLQueryItemTests.swift → ...ages/AIChat/Tests/URLExtensionTests.swift
@@ -1,5 +1,5 @@
 //
-//  URLQueryItemTests.swift
+//  URLExtensionTests.swift
 //  DuckDuckGo
 //
 //  Copyright © 2022 DuckDuckGo. All rights reserved.
@@ -20,7 +20,7 @@
 import XCTest
 @testable import AIChat
 
-class URLQueryItemTests: XCTestCase {
+class URLExtensionTests: XCTestCase {
 
     func testAddingQueryItemToEmptyURL() {
         let url = URL(string: "https://example.com")!
@@ -81,6 +81,46 @@ class URLQueryItemTests: XCTestCase {
         XCTAssertEqual(result.host, "example.com")
         XCTAssertEqual(result.queryItemsDictionary, ["key": ""])
     }
+
+    func testIsDuckAIURLWithValidURL() {
+        if let url = URL(string: "https://duckduckgo.com/?ia=chat") {
+            XCTAssertTrue(url.isDuckAIURL, "The URL should be identified as a DuckDuckGo AI URL.")
+        } else {
+            XCTFail("Failed to create URL from string.")
+        }
+    }
+
+    func testIsDuckAIURLWithInvalidDomain() {
+        if let url = URL(string: "https://example.com/?ia=chat") {
+            XCTAssertFalse(url.isDuckAIURL, "The URL should not be identified as a DuckDuckGo AI URL due to the domain.")
+        } else {
+            XCTFail("Failed to create URL from string.")
+        }
+    }
+
+    func testIsDuckAIURLWithMissingQueryItem() {
+        if let url = URL(string: "https://duckduckgo.com/") {
+            XCTAssertFalse(url.isDuckAIURL, "The URL should not be identified as a DuckDuckGo AI URL due to missing query item.")
+        } else {
+            XCTFail("Failed to create URL from string.")
+        }
+    }
+
+    func testIsDuckAIURLWithDifferentQueryItem() {
+        if let url = URL(string: "https://duckduckgo.com/?ia=search") {
+            XCTAssertFalse(url.isDuckAIURL, "The URL should not be identified as a DuckDuckGo AI URL due to different query item value.")
+        } else {
+            XCTFail("Failed to create URL from string.")
+        }
+    }
+
+    func testIsDuckAIURLWithAdditionalQueryItems() {
+        if let url = URL(string: "https://duckduckgo.com/?ia=chat&other=param") {
+            XCTAssertTrue(url.isDuckAIURL, "The URL should be identified as a DuckDuckGo AI URL even with additional query items.")
+        } else {
+            XCTFail("Failed to create URL from string.")
+        }
+    }
 }
 
 extension URL {