From e6537d5c8f79d83d07e2682c25eb39dfeac5319b Mon Sep 17 00:00:00 2001
From: Bug Magnet <marco.nikic@mullvad.net>
Date: Fri, 1 Dec 2023 11:38:44 +0100
Subject: [PATCH] Block all traffic when device is connecting or reconnecting

---
 .../Actor/PacketTunnelActor+ErrorState.swift        | 13 +++++++++----
 .../Actor/PacketTunnelActor+KeyPolicy.swift         | 12 ++++++++++++
 .../PacketTunnelActor+NetworkReachability.swift     |  3 +++
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/ios/PacketTunnelCore/Actor/PacketTunnelActor+ErrorState.swift b/ios/PacketTunnelCore/Actor/PacketTunnelActor+ErrorState.swift
index 077bfdbb25f4..c9f144a0f2cc 100644
--- a/ios/PacketTunnelCore/Actor/PacketTunnelActor+ErrorState.swift
+++ b/ios/PacketTunnelCore/Actor/PacketTunnelActor+ErrorState.swift
@@ -43,7 +43,7 @@ extension PacketTunnelActor {
 
         if let blockedState = makeBlockedState(reason: reason) {
             state = .error(blockedState)
-            await configureAdapterForErrorState()
+            await blockAllTrafficUntilDeviceIsConnected(shouldStopTunnel: true)
         }
     }
 
@@ -110,7 +110,7 @@ extension PacketTunnelActor {
     /**
      Configure tunnel with empty WireGuard configuration that consumes all traffic on device emulating a firewall blocking all traffic.
      */
-    private func configureAdapterForErrorState() async {
+    func blockAllTrafficUntilDeviceIsConnected(shouldStopTunnel: Bool = false) async {
         do {
             let configurationBuilder = ConfigurationBuilder(
                 privateKey: PrivateKey(),
@@ -119,11 +119,16 @@ extension PacketTunnelActor {
             var config = try configurationBuilder.makeConfiguration()
             config.dns = [IPv4Address.loopback]
             config.interfaceAddresses = [IPAddressRange(from: "10.64.0.1/8")!]
+            // Randomize the port and the last octet of the local address to reduce risks of accidentally addressing a service running on localhost.
+            let randomPort = String(describing: (1 ... (UInt16.max - 1)).randomElement()!)
+            let randomOctet = String(describing: (1 ... 254).randomElement()!)
             config.peer = TunnelPeer(
-                endpoint: .ipv4(IPv4Endpoint(string: "127.0.0.1:9090")!),
+                endpoint: .ipv4(IPv4Endpoint(string: "127.0.0.\(randomOctet):\(randomPort)")!),
                 publicKey: PrivateKey().publicKey
             )
-            try? await tunnelAdapter.stop()
+            if shouldStopTunnel {
+                try? await tunnelAdapter.stop()
+            }
             try await tunnelAdapter.start(configuration: config)
         } catch {
             logger.error(error: error, message: "Unable to configure the tunnel for error state.")
diff --git a/ios/PacketTunnelCore/Actor/PacketTunnelActor+KeyPolicy.swift b/ios/PacketTunnelCore/Actor/PacketTunnelActor+KeyPolicy.swift
index 5ab13353e915..6239ce182d44 100644
--- a/ios/PacketTunnelCore/Actor/PacketTunnelActor+KeyPolicy.swift
+++ b/ios/PacketTunnelCore/Actor/PacketTunnelActor+KeyPolicy.swift
@@ -48,6 +48,9 @@ extension PacketTunnelActor {
         case var .connecting(connState):
             if mutateConnectionState(&connState) {
                 state = .connecting(connState)
+                Task {
+                    await blockAllTrafficUntilDeviceIsConnected()
+                }
             }
 
         case var .connected(connState):
@@ -58,6 +61,9 @@ extension PacketTunnelActor {
         case var .reconnecting(connState):
             if mutateConnectionState(&connState) {
                 state = .reconnecting(connState)
+                Task {
+                    await blockAllTrafficUntilDeviceIsConnected()
+                }
             }
 
         case var .error(blockedState):
@@ -127,6 +133,9 @@ extension PacketTunnelActor {
         case var .connecting(connState):
             if setCurrentKeyPolicy(&connState.keyPolicy) {
                 state = .connecting(connState)
+                Task {
+                    await blockAllTrafficUntilDeviceIsConnected()
+                }
                 return true
             }
 
@@ -139,6 +148,9 @@ extension PacketTunnelActor {
         case var .reconnecting(connState):
             if setCurrentKeyPolicy(&connState.keyPolicy) {
                 state = .reconnecting(connState)
+                Task {
+                    await blockAllTrafficUntilDeviceIsConnected()
+                }
                 return true
             }
 
diff --git a/ios/PacketTunnelCore/Actor/PacketTunnelActor+NetworkReachability.swift b/ios/PacketTunnelCore/Actor/PacketTunnelActor+NetworkReachability.swift
index 2d3b600239ae..916b91be5549 100644
--- a/ios/PacketTunnelCore/Actor/PacketTunnelActor+NetworkReachability.swift
+++ b/ios/PacketTunnelCore/Actor/PacketTunnelActor+NetworkReachability.swift
@@ -65,6 +65,9 @@ extension PacketTunnelActor {
         case var .reconnecting(connState):
             if mutateConnectionState(&connState) {
                 state = .reconnecting(connState)
+                Task {
+                    await blockAllTrafficUntilDeviceIsConnected()
+                }
             }
 
         case var .disconnecting(connState):