From 91349fd780de9d77f2d62765c02ae01d92a5e1ab Mon Sep 17 00:00:00 2001
From: Davide De Rosa <keeshux@gmail.com>
Date: Wed, 24 Oct 2018 09:29:39 +0200
Subject: [PATCH] Take shouldChangeProtocol out of GenericSocket

Behavior is not exactly similar in UDP and TCP.
---
 .../Sources/AppExtension/GenericSocket.swift  |  2 --
 .../Transport/NETCPInterface.swift            |  1 -
 .../AppExtension/TunnelKitProvider.swift      | 25 ++++++++++++-------
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/TunnelKit/Sources/AppExtension/GenericSocket.swift b/TunnelKit/Sources/AppExtension/GenericSocket.swift
index 87aa9e2..4a59c02 100644
--- a/TunnelKit/Sources/AppExtension/GenericSocket.swift
+++ b/TunnelKit/Sources/AppExtension/GenericSocket.swift
@@ -44,8 +44,6 @@ protocol LinkProducer {
 protocol GenericSocketDelegate: class {
     func socketDidTimeout(_ socket: GenericSocket)
 
-    func socketShouldChangeProtocol(_ socket: GenericSocket) -> Bool
-    
     func socketDidBecomeActive(_ socket: GenericSocket)
 
     func socket(_ socket: GenericSocket, didShutdownWithFailure failure: Bool)
diff --git a/TunnelKit/Sources/AppExtension/Transport/NETCPInterface.swift b/TunnelKit/Sources/AppExtension/Transport/NETCPInterface.swift
index e74f4ae..3fd949c 100644
--- a/TunnelKit/Sources/AppExtension/Transport/NETCPInterface.swift
+++ b/TunnelKit/Sources/AppExtension/Transport/NETCPInterface.swift
@@ -79,7 +79,6 @@ class NETCPSocket: NSObject, GenericSocket {
                 return
             }
             guard _self.isActive else {
-                _ = _self.delegate?.socketShouldChangeProtocol(_self)
                 _self.delegate?.socketDidTimeout(_self)
                 return
             }
diff --git a/TunnelKit/Sources/AppExtension/TunnelKitProvider.swift b/TunnelKit/Sources/AppExtension/TunnelKitProvider.swift
index f8f25a1..dfd3005 100644
--- a/TunnelKit/Sources/AppExtension/TunnelKitProvider.swift
+++ b/TunnelKit/Sources/AppExtension/TunnelKitProvider.swift
@@ -400,14 +400,14 @@ extension TunnelKitProvider: GenericSocketDelegate {
         log.debug("Socket timed out waiting for activity, cancelling...")
         reasserting = true
         socket.shutdown()
-    }
-    
-    func socketShouldChangeProtocol(_ socket: GenericSocket) -> Bool {
-        guard strategy.tryNextProtocol() else {
-            disposeTunnel(error: ProviderError.exhaustedProtocols)
-            return false
+
+        // fallback: TCP connection timeout suggests falling back
+        if let _ = socket as? NETCPSocket {
+            guard tryNextProtocol() else {
+                // disposeTunnel
+                return
+            }
         }
-        return true
     }
     
     func socketDidBecomeActive(_ socket: GenericSocket) {
@@ -448,9 +448,9 @@ extension TunnelKitProvider: GenericSocketDelegate {
         // clean up
         finishTunnelDisconnection(error: shutdownError)
 
-        // treat negotiation timeout as socket timeout, UDP is connection-less
+        // fallback: UDP is connection-less, treat negotiation timeout as socket timeout
         if didTimeoutNegotiation {
-            guard socketShouldChangeProtocol(socket) else {
+            guard tryNextProtocol() else {
                 // disposeTunnel
                 return
             }
@@ -574,6 +574,13 @@ extension TunnelKitProvider: SessionProxyDelegate {
 }
 
 extension TunnelKitProvider {
+    private func tryNextProtocol() -> Bool {
+        guard strategy.tryNextProtocol() else {
+            disposeTunnel(error: ProviderError.exhaustedProtocols)
+            return false
+        }
+        return true
+    }
     
     // MARK: Logging