From 1bfcf07f796c7a22b5862a0953182bee0aaa2634 Mon Sep 17 00:00:00 2001
From: "Dr. Brandon Wiley" <brandon@operatorfoundation.org>
Date: Thu, 11 Jan 2018 13:03:09 -0600
Subject: [PATCH] Adjusted callback API

---
 Sources/Transport/PacketTunnelProvider.swift | 13 ++------
 Sources/Transport/UDPSession.swift           | 32 ++++++++++++++++++++
 2 files changed, 35 insertions(+), 10 deletions(-)
 create mode 100644 Sources/Transport/UDPSession.swift

diff --git a/Sources/Transport/PacketTunnelProvider.swift b/Sources/Transport/PacketTunnelProvider.swift
index fe5f71c..1a29cbe 100644
--- a/Sources/Transport/PacketTunnelProvider.swift
+++ b/Sources/Transport/PacketTunnelProvider.swift
@@ -10,14 +10,8 @@ import Foundation
 import NetworkExtension
 
 extension NEPacketTunnelProvider: PacketTunnelProvider {
-    public func createTCPConnectionThroughTunnel(to remoteEndpoint: NWEndpoint, enableTLS: Bool, tlsParameters TLSParameters: NWTLSParameters?, delegate: Any?, stateCallback: @escaping (NWTCPConnectionState, Error?) -> Void) -> TCPConnection? {
-        let tunnel = createTCPConnectionThroughTunnel(to: remoteEndpoint, enableTLS: enableTLS, tlsParameters: TLSParameters, delegate: delegate)
-        let _ = tunnel.observe(\NWTCPConnection.state) {
-            (conn, change) in
-            
-            stateCallback(change.newValue!, nil)
-        }
-        return tunnel
+    public func createTCPConnectionThroughTunnel(to remoteEndpoint: NWEndpoint, enableTLS: Bool, tlsParameters TLSParameters: NWTLSParameters?, delegate: Any?) -> TCPConnection? {
+        return createTCPConnectionThroughTunnel(to: remoteEndpoint, enableTLS: enableTLS, tlsParameters: TLSParameters, delegate: delegate)
     }
 }
 
@@ -25,6 +19,5 @@ public protocol PacketTunnelProvider {
     func createTCPConnectionThroughTunnel(to remoteEndpoint: NWEndpoint,
                                           enableTLS: Bool,
                                           tlsParameters TLSParameters: NWTLSParameters?,
-                                          delegate: Any?,
-                                          stateCallback: @escaping (NWTCPConnectionState, Error?) -> Void) -> TCPConnection?
+                                          delegate: Any?) -> TCPConnection?
 }
diff --git a/Sources/Transport/UDPSession.swift b/Sources/Transport/UDPSession.swift
new file mode 100644
index 0000000..5e50eb7
--- /dev/null
+++ b/Sources/Transport/UDPSession.swift
@@ -0,0 +1,32 @@
+//
+//  UDPSession.swift
+//  Transport
+//
+//  Created by Brandon Wiley on 1/11/18.
+//
+
+import Foundation
+import NetworkExtension
+
+extension NWUDPSession: UDPSession {}
+
+public protocol UDPSession {
+    var state: NWUDPSessionState { get }
+    var isViable: Bool { get }
+    
+    var resolvedEndpoint: NWEndpoint? { get }
+    func tryNextResolvedEndpoint()
+    
+    func setReadHandler(_ handler: @escaping ([Data]?, Error?) -> Void,
+                        maxDatagrams: Int)
+    func writeDatagram(_ datagram: Data,
+                       completionHandler: @escaping (Error?) -> Void)
+    func writeMultipleDatagrams(_ datagramArray: [Data],
+                                completionHandler: @escaping (Error?) -> Void)
+    var maximumDatagramLength: Int { get }
+    
+    func cancel()
+    
+    var hasBetterPath: Bool { get }
+    init(upgradeFor session: NWUDPSession)
+}