From 1472616bb528d707f427b70f64f055ab4f10d69f Mon Sep 17 00:00:00 2001
From: Fernando Bunn <bunn@idevzilla.com>
Date: Wed, 18 Dec 2024 18:45:28 -0300
Subject: [PATCH] Validate duck.ai URL

---
 .../AIChat/AIChatWebViewController.swift      |  2 +-
 .../AIChat/Sources/AIChat/URL+Extension.swift | 13 ++++++
 ...temTests.swift => URLExtensionTests.swift} | 44 ++++++++++++++++++-
 3 files changed, 56 insertions(+), 3 deletions(-)
 rename LocalPackages/AIChat/Tests/{URLQueryItemTests.swift => URLExtensionTests.swift} (67%)

diff --git a/LocalPackages/AIChat/Sources/AIChat/AIChatWebViewController.swift b/LocalPackages/AIChat/Sources/AIChat/AIChatWebViewController.swift
index bb04b03f84..5a15b6546d 100644
--- 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
index 681695a62b..1e0978f346 100644
--- 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/LocalPackages/AIChat/Tests/URLQueryItemTests.swift b/LocalPackages/AIChat/Tests/URLExtensionTests.swift
similarity index 67%
rename from LocalPackages/AIChat/Tests/URLQueryItemTests.swift
rename to LocalPackages/AIChat/Tests/URLExtensionTests.swift
index f6a9a548f9..fc72d2a759 100644
--- a/LocalPackages/AIChat/Tests/URLQueryItemTests.swift
+++ b/LocalPackages/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 {