From daa31da8af00b03f1d1f3a0f6f10a31a6de4880e Mon Sep 17 00:00:00 2001
From: Niklas Berglund <niklas.berglund@gmail.com>
Date: Thu, 28 Mar 2024 15:02:10 +0100
Subject: [PATCH] Add custom DNS test

---
 .../Networking/Networking.swift               | 61 +++++++++++++++++++
 ios/MullvadVPNUITests/RelayTests.swift        | 30 +++++++++
 .../Test base classes/BaseUITestCase.swift    |  1 -
 3 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/ios/MullvadVPNUITests/Networking/Networking.swift b/ios/MullvadVPNUITests/Networking/Networking.swift
index 120110969b48..4af87b5523a9 100644
--- a/ios/MullvadVPNUITests/Networking/Networking.swift
+++ b/ios/MullvadVPNUITests/Networking/Networking.swift
@@ -15,6 +15,11 @@ enum NetworkingError: Error {
     case internalError(reason: String)
 }
 
+struct DNSServerEntry: Decodable {
+    let organization: String
+    let mullvad_dns: Bool
+}
+
 /// Class with methods for verifying network connectivity
 class Networking {
     /// Get IP address of the iOS device under test
@@ -165,4 +170,60 @@ class Networking {
     public static func verifyCannotReachAdServingDomain() throws {
         XCTAssertFalse(try Self.canConnectSocket(host: try Self.getAdServingDomain(), port: "80"))
     }
+
+    /// Verify that the expected DNS server is used by verifying provider name and whether it is a Mullvad DNS server or not
+    public static func verifyDNSServerProvider(_ providerName: String, isMullvad: Bool) throws {
+        guard let mullvadDNSLeakURL = URL(string: "https://am.i.mullvad.net/dnsleak") else {
+            throw NetworkingError.internalError(reason: "Failed to create URL object")
+        }
+
+        var request = URLRequest(url: mullvadDNSLeakURL)
+        request.setValue("application/json", forHTTPHeaderField: "accept")
+
+        var requestData: Data?
+        var requestResponse: URLResponse?
+        var requestError: Error?
+        let completionHandlerInvokedExpectation = XCTestExpectation(
+            description: "Completion handler for the request is invoked"
+        )
+
+        do {
+            let dataTask = URLSession.shared.dataTask(with: request) { data, response, error in
+                requestData = data
+                requestResponse = response
+                requestError = error
+                completionHandlerInvokedExpectation.fulfill()
+            }
+
+            dataTask.resume()
+
+            let waitResult = XCTWaiter.wait(for: [completionHandlerInvokedExpectation], timeout: 30)
+
+            if waitResult != .completed {
+                XCTFail("Failed to verify DNS server provider - timeout")
+            } else {
+                if let response = requestResponse as? HTTPURLResponse {
+                    if response.statusCode != 200 {
+                        XCTFail("Failed to verify DNS server provider - unexpected server response")
+                    }
+                }
+
+                if let error = requestError {
+                    XCTFail("Failed to verify DNS server provider - encountered error \(error.localizedDescription)")
+                }
+
+                if let requestData = requestData {
+                    let dnsServerEntries = try JSONDecoder().decode([DNSServerEntry].self, from: requestData)
+                    XCTAssertGreaterThanOrEqual(dnsServerEntries.count, 1)
+
+                    for dnsServerEntry in dnsServerEntries {
+                        XCTAssertEqual(dnsServerEntry.organization, providerName)
+                        XCTAssertEqual(dnsServerEntry.mullvad_dns, isMullvad)
+                    }
+                }
+            }
+        } catch {
+            XCTFail("Failed to verify DNS server provider - couldn't serialize JSON")
+        }
+    }
 }
diff --git a/ios/MullvadVPNUITests/RelayTests.swift b/ios/MullvadVPNUITests/RelayTests.swift
index 5df5b05a068f..d0f8ea20375a 100644
--- a/ios/MullvadVPNUITests/RelayTests.swift
+++ b/ios/MullvadVPNUITests/RelayTests.swift
@@ -197,4 +197,34 @@ class RelayTests: LoggedInWithTimeUITestCase {
 
         return relayIPAddress
     }
+
+    func testCustomDNS() throws {
+        let dnsServerIPAddress = "8.8.8.8"
+        let dnsServerProviderName = "GOOGLE"
+
+        TunnelControlPage(app)
+            .tapSecureConnectionButton()
+
+        allowAddVPNConfigurations()
+
+        HeaderBar(app)
+            .tapSettingsButton()
+
+        SettingsPage(app)
+            .tapVPNSettingsCell()
+
+        VPNSettingsPage(app)
+            .tapDNSSettingsCell()
+
+        DNSSettingsPage(app)
+            .tapEditButton()
+            .tapAddAServer()
+            .tapEnterIPAddressTextField()
+            .enterText(dnsServerIPAddress)
+            .dismissKeyboard()
+            .tapUseCustomDNSSwitch()
+            .tapDoneButton()
+
+        try Networking.verifyDNSServerProvider(dnsServerProviderName, isMullvad: false)
+    }
 }
diff --git a/ios/MullvadVPNUITests/Test base classes/BaseUITestCase.swift b/ios/MullvadVPNUITests/Test base classes/BaseUITestCase.swift
index ae4011cfc7ef..11b0cf6ee5f0 100644
--- a/ios/MullvadVPNUITests/Test base classes/BaseUITestCase.swift	
+++ b/ios/MullvadVPNUITests/Test base classes/BaseUITestCase.swift	
@@ -129,7 +129,6 @@ class BaseUITestCase: XCTestCase {
 
     func logoutIfLoggedIn() {
         if isLoggedIn() {
-            // First dismiss settings modal if presented
             if isPresentingSettings() {
                 SettingsPage(app)
                     .swipeDownToDismissModal()