From 753927f36b354d986859d9a75bca1944a35e8ea0 Mon Sep 17 00:00:00 2001
From: Roopesh Chander <roop@roopc.net>
Date: Tue, 5 May 2020 17:31:27 +0530
Subject: [PATCH] Fix how NETunnelInterface handles IP protocol number

The IP protocol number passed to NEPacketTunnelFlow is determined per
packet based on the IP header, instead of determining it based on
whether IPv6 settings are available or not.
---
 .../Transport/NETunnelInterface.swift         | 34 ++++++++++++-------
 .../AppExtension/OpenVPNTunnelProvider.swift  |  2 +-
 2 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/TunnelKit/Sources/AppExtension/Transport/NETunnelInterface.swift b/TunnelKit/Sources/AppExtension/Transport/NETunnelInterface.swift
index 429881b..f9d03dd 100644
--- a/TunnelKit/Sources/AppExtension/Transport/NETunnelInterface.swift
+++ b/TunnelKit/Sources/AppExtension/Transport/NETunnelInterface.swift
@@ -41,21 +41,12 @@ import NetworkExtension
 public class NETunnelInterface: TunnelInterface {
     private weak var impl: NEPacketTunnelFlow?
     
-    private let protocolNumber: NSNumber
+    private static let protocolNumberForIPv4 = NSNumber(value: AF_INET)
+    private static let protocolNumberForIPv6 = NSNumber(value: AF_INET6)
 
     /// :nodoc:
-    public init(impl: NEPacketTunnelFlow, isIPv6: Bool) {
+    public init(impl: NEPacketTunnelFlow) {
         self.impl = impl
-        #if os(macOS)
-        if #available(OSX 10.15, *) {
-            protocolNumber = (isIPv6 ? AF_INET6 : AF_INET) as NSNumber
-        } else {
-            // Force IPv4 on Mojave otherwise it breaks
-            protocolNumber = AF_INET as NSNumber
-        }
-        #else
-        protocolNumber = (isIPv6 ? AF_INET6 : AF_INET) as NSNumber
-        #endif
     }
     
     // MARK: TunnelInterface
@@ -85,14 +76,31 @@ public class NETunnelInterface: TunnelInterface {
     
     /// :nodoc:
     public func writePacket(_ packet: Data, completionHandler: ((Error?) -> Void)?) {
+        let protocolNumber = NETunnelInterface.ipProtocolNumber(inPacket: packet)
         impl?.writePackets([packet], withProtocols: [protocolNumber])
         completionHandler?(nil)
     }
     
     /// :nodoc:
     public func writePackets(_ packets: [Data], completionHandler: ((Error?) -> Void)?) {
-        let protocols = [NSNumber](repeating: protocolNumber, count: packets.count)
+        let protocols = packets.map {
+            NETunnelInterface.ipProtocolNumber(inPacket: $0)
+        }
         impl?.writePackets(packets, withProtocols: protocols)
         completionHandler?(nil)
     }
+
+    private static func ipProtocolNumber(inPacket packet: Data) -> NSNumber {
+        // 'packet' contains the decrypted incoming IP packet data
+
+        // The first 4 bits identify the IP version
+        let ipVersion = ((packet[0] & 0xf0) >> 4)
+        assert(ipVersion == 4 || ipVersion == 6)
+
+        if ipVersion == 6 {
+            return NETunnelInterface.protocolNumberForIPv6
+        } else {
+            return NETunnelInterface.protocolNumberForIPv4
+        }
+    }
 }
diff --git a/TunnelKit/Sources/Protocols/OpenVPN/AppExtension/OpenVPNTunnelProvider.swift b/TunnelKit/Sources/Protocols/OpenVPN/AppExtension/OpenVPNTunnelProvider.swift
index 7b80a13..979a79b 100644
--- a/TunnelKit/Sources/Protocols/OpenVPN/AppExtension/OpenVPNTunnelProvider.swift
+++ b/TunnelKit/Sources/Protocols/OpenVPN/AppExtension/OpenVPNTunnelProvider.swift
@@ -565,7 +565,7 @@ extension OpenVPNTunnelProvider: OpenVPNSessionDelegate {
             
             log.info("Tunnel interface is now UP")
             
-            session.setTunnel(tunnel: NETunnelInterface(impl: self.packetFlow, isIPv6: options.ipv6 != nil))
+            session.setTunnel(tunnel: NETunnelInterface(impl: self.packetFlow))
 
             self.pendingStartHandler?(nil)
             self.pendingStartHandler = nil