VPN: Fix DNS resolution API

Signed-off-by: Roopesh Chander <roop@roopc.net>

2 files changed, 25 insertions, 9 deletions

diff --git a/WireGuard/WireGuard/VPN/DNSResolver.swift b/WireGuard/WireGuard/VPN/DNSResolver.swift
index 89b1a83..e027852 100644
--- a/WireGuard/WireGuard/VPN/DNSResolver.swift
+++ b/WireGuard/WireGuard/VPN/DNSResolver.swift
@@ -5,19 +5,34 @@ import Network
 import Foundation
 
 class DNSResolver {
-    let endpoints: [Endpoint]
+    let endpoints: [Endpoint?]
 
-    init(endpoints: [Endpoint]) {
+    init(endpoints: [Endpoint?]) {
         self.endpoints = endpoints
     }
 
-    func resolve(completionHandler: @escaping ([Endpoint?]) -> Void) {
+    func resolve(completionHandler: @escaping ([Endpoint?]?) -> Void) {
         let endpoints = self.endpoints
         DispatchQueue.global(qos: .userInitiated).async {
-            var resolvedEndpoints = Array<Endpoint?>(repeating: nil, count: endpoints.count)
-            for (i, endpoint) in endpoints.enumerated() {
-                let resolvedEndpoint = DNSResolver.resolveSync(endpoint: endpoint)
-                resolvedEndpoints[i] = resolvedEndpoint
+            var resolvedEndpoints: [Endpoint?] = []
+            var isError = false
+            for endpoint in endpoints {
+                if let endpoint = endpoint {
+                    if let resolvedEndpoint = DNSResolver.resolveSync(endpoint: endpoint) {
+                        resolvedEndpoints.append(resolvedEndpoint)
+                    } else {
+                        isError = true
+                        break
+                    }
+                } else {
+                    resolvedEndpoints.append(nil)
+                }
+            }
+            if (isError) {
+                DispatchQueue.main.async {
+                    completionHandler(nil)
+                }
+                return
             }
             DispatchQueue.main.async {
                 completionHandler(resolvedEndpoints)
diff --git a/WireGuard/WireGuard/VPN/TunnelsManager.swift b/WireGuard/WireGuard/VPN/TunnelsManager.swift
index 65e8373..039862d 100644
--- a/WireGuard/WireGuard/VPN/TunnelsManager.swift
+++ b/WireGuard/WireGuard/VPN/TunnelsManager.swift
@@ -256,13 +256,14 @@ class TunnelContainer: NSObject {
     fileprivate func activate(completionHandler: @escaping (Bool) -> Void) {
         assert(status == .inactive)
         guard let tunnelConfiguration = tunnelConfiguration() else { fatalError() }
-        let endpoints = tunnelConfiguration.peers.compactMap { $0.endpoint }
+        let endpoints = tunnelConfiguration.peers.map { $0.endpoint }
         let dnsResolver = DNSResolver(endpoints: endpoints)
         assert(self.dnsResolver == nil)
         self.dnsResolver = dnsResolver
         status = .resolvingEndpointDomains
         dnsResolver.resolve { [weak self] endpoints in
-            guard (!endpoints.contains { $0 == nil }) else {
+            guard let endpoints = endpoints else {
+                // TODO: Show error message
                 completionHandler(false)
                 return
             }