From 79aabd549e274515e306aa49dccc175c69ca6308 Mon Sep 17 00:00:00 2001
From: Sam Symons <sam@samsymons.com>
Date: Mon, 11 Dec 2023 11:18:52 -0800
Subject: [PATCH] Fix an IPv6 regression.

---
 .../PacketTunnelProvider.swift                |  2 +-
 .../WireGuardKit/Endpoint.swift               | 15 +++-
 .../EndpointTests.swift                       | 73 +++++++++++++++++++
 3 files changed, 86 insertions(+), 4 deletions(-)
 create mode 100644 Tests/NetworkProtectionTests/EndpointTests.swift

diff --git a/Sources/NetworkProtection/PacketTunnelProvider.swift b/Sources/NetworkProtection/PacketTunnelProvider.swift
index 74a0f9c62..5748e1587 100644
--- a/Sources/NetworkProtection/PacketTunnelProvider.swift
+++ b/Sources/NetworkProtection/PacketTunnelProvider.swift
@@ -989,7 +989,7 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
     }
 
     private func handleGetServerAddress(completionHandler: ((Data?) -> Void)? = nil) {
-        let response = lastSelectedServerInfo?.endpoint.map { ExtensionMessageString($0.host.description) }
+        let response = lastSelectedServerInfo?.endpoint.map { ExtensionMessageString($0.host.hostWithoutPort) }
         completionHandler?(response?.rawValue)
     }
 
diff --git a/Sources/NetworkProtection/WireGuardKit/Endpoint.swift b/Sources/NetworkProtection/WireGuardKit/Endpoint.swift
index 9d9502dd3..f4259e8f8 100644
--- a/Sources/NetworkProtection/WireGuardKit/Endpoint.swift
+++ b/Sources/NetworkProtection/WireGuardKit/Endpoint.swift
@@ -30,7 +30,16 @@ extension Endpoint: Hashable {
 extension Endpoint: CustomStringConvertible {
 
     public var description: String {
-        "\(host):\(port)"
+        switch host {
+        case .name(let hostname, _):
+            return "\(hostname):\(port)"
+        case .ipv4(let address):
+            return "\(address):\(port)"
+        case .ipv6(let address):
+            return "[\(address)]:\(port)"
+        @unknown default:
+            fatalError()
+        }
     }
 
     public init?(from string: String) {
@@ -80,8 +89,8 @@ extension Endpoint {
     }
 }
 
-extension NWEndpoint.Host: CustomStringConvertible {
-    public var description: String {
+extension NWEndpoint.Host {
+    public var hostWithoutPort: String {
         switch self {
         case .name(let hostname, _):
             return hostname
diff --git a/Tests/NetworkProtectionTests/EndpointTests.swift b/Tests/NetworkProtectionTests/EndpointTests.swift
new file mode 100644
index 000000000..b732948c8
--- /dev/null
+++ b/Tests/NetworkProtectionTests/EndpointTests.swift
@@ -0,0 +1,73 @@
+//
+//  File.swift
+//  DuckDuckGo
+//
+//  Copyright © 2023 DuckDuckGo. All rights reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+import XCTest
+import Network
+@testable import NetworkProtection
+
+final class EndpointTests: XCTestCase {
+
+    func testEndpointWithHostName_ShouldIncludePort() {
+        let endpoint = self.endpointWithHostName()
+        XCTAssertEqual(endpoint.description, "https://duckduckgo.com:443")
+    }
+
+    func testEndpointWithIPv4Address_ShouldIncludePort() {
+        let endpoint = self.endpointWithIPv4()
+        XCTAssertEqual(endpoint.description, "52.250.42.157:443")
+    }
+
+    func testEndpointWithIPv6Address_ShouldIncludePort() {
+        let endpoint = self.endpointWithIPv6()
+        XCTAssertEqual(endpoint.description, "[2001:db8:85a3::8a2e:370:7334]:443")
+    }
+
+    func testParsingEndpointFromIPv4Address() {
+        let address = "52.250.42.157:443"
+        let endpoint = Endpoint(from: address)!
+        XCTAssertEqual(endpoint.description, address)
+    }
+
+    func testParsingEndpointFromIPv6Address() {
+        let address = "[2001:0db8:85a3:0000:0000:8a2e:0370]:443"
+        let endpoint = Endpoint(from: address)!
+        XCTAssertEqual(endpoint.description, "2001:0db8:85a3:0000:0000:8a2e:0370:443")
+    }
+
+    private func endpointWithHostName() -> Endpoint {
+        let host = NWEndpoint.Host.name("https://duckduckgo.com", nil)
+        let port = NWEndpoint.Port.https
+        return Endpoint(host: host, port: port)
+    }
+
+    private func endpointWithIPv4() -> Endpoint {
+        let address = IPv4Address("52.250.42.157")!
+        let host = NWEndpoint.Host.ipv4(address)
+        let port = NWEndpoint.Port.https
+        return Endpoint(host: host, port: port)
+    }
+
+    private func endpointWithIPv6() -> Endpoint {
+        let address = IPv6Address("2001:0db8:85a3:0000:0000:8a2e:0370:7334")!
+        let host = NWEndpoint.Host.ipv6(address)
+        let port = NWEndpoint.Port.https
+        return Endpoint(host: host, port: port)
+    }
+
+}