From 776ff7171091af2f8421b362a1ef3150d72c245a Mon Sep 17 00:00:00 2001
From: "Dr. Brandon Wiley" <brandon@operatorfoundation.org>
Date: Thu, 11 Nov 2021 12:24:58 -0600
Subject: [PATCH] Now supports both socket and NWConnection backends

---
 Package.resolved                              |   8 +-
 .../TransmissionConnection.swift              | 509 ++++++++++--------
 2 files changed, 303 insertions(+), 214 deletions(-)

diff --git a/Package.resolved b/Package.resolved
index c90130a..50de3d4 100644
--- a/Package.resolved
+++ b/Package.resolved
@@ -15,8 +15,8 @@
         "repositoryURL": "https://github.com/OperatorFoundation/Chord",
         "state": {
           "branch": null,
-          "revision": "94ef383494fb4c1b8b0e0e368df6f35265ff62d1",
-          "version": "0.0.14"
+          "revision": "756eeb26b10016784e6dc7cfce33068b36e49f39",
+          "version": "0.0.15"
         }
       },
       {
@@ -60,8 +60,8 @@
         "repositoryURL": "https://github.com/OperatorFoundation/Transport",
         "state": {
           "branch": null,
-          "revision": "a95c7bd870c558a040f9a22a88badab5577bd4df",
-          "version": "2.3.8"
+          "revision": "514e391f7b927e4cc90acb0076a36700f584192f",
+          "version": "2.3.9"
         }
       }
     ]
diff --git a/Sources/TransmissionLinux/TransmissionConnection.swift b/Sources/TransmissionLinux/TransmissionConnection.swift
index 6377b2a..8a74628 100644
--- a/Sources/TransmissionLinux/TransmissionConnection.swift
+++ b/Sources/TransmissionLinux/TransmissionConnection.swift
@@ -1,37 +1,115 @@
 import Foundation
-import Socket
 import Datable
+import Transport
+import Logging
+import SwiftQueue
+
 import Chord
+import Socket
+import Net
 
 public class TransmissionConnection: Connection
 {
+    let id: Int
+    var connection: NetworkConnection
     var connectLock = DispatchGroup()
     var readLock = DispatchGroup()
     var writeLock = DispatchGroup()
+    let log: Logger?
+    let states: Queue<Bool> = Queue<Bool>()
 
-    let socket: Socket
     var buffer: Data = Data()
 
-    public init?(host: String, port: Int)
+    public init?(host: String, port: Int, type: ConnectionType = .tcp, logger: Logger? = nil)
     {
-        guard let socket = try? Socket.create() else {return nil}
-        self.socket = socket
-        
-        do
+        self.log = logger
+
+        switch type
         {
-            try self.socket.connect(to: host, port: Int32(port))
+            case .tcp:
+                guard let socket = try? Socket.create() else {return nil}
+                self.connection = .socket(socket)
+                self.id = Int(socket.socketfd)
+
+                do
+                {
+                    try socket.connect(to: host, port: Int32(port))
+                }
+                catch
+                {
+                    return nil
+                }
+            case .udp:
+                // FIXME
+                return nil
         }
-        catch
+    }
+
+    public required init?(transport: Transport.Connection, logger: Logger? = nil)
+    {
+        self.log = logger
+        maybeLog(message: "Initializing Transmission connection", logger: self.log)
+
+        self.id = Int.random(in: 0..<Int.max)
+
+        var mutableTransport = transport
+        self.connection = .transport(mutableTransport)
+        mutableTransport.stateUpdateHandler = self.handleState
+        mutableTransport.start(queue: .global())
+
+        guard let success = self.states.dequeue() else {return nil}
+        guard success else {return nil}
+    }
+
+    public init(socket: Socket)
+    {
+        self.connection = .socket(socket)
+        self.id = Int(socket.socketfd)
+        self.log = nil
+    }
+
+    func handleState(state: NWConnection.State)
+    {
+        connectLock.wait()
+
+        switch state
         {
-            return nil
+            case .ready:
+                self.states.enqueue(true)
+                return
+            case .cancelled:
+                self.states.enqueue(false)
+                self.failConnect()
+                return
+            case .failed(_):
+                self.states.enqueue(false)
+                self.failConnect()
+                return
+            case .waiting(_):
+                self.states.enqueue(false)
+                self.failConnect()
+                return
+            default:
+                return
         }
     }
-    
-    public init(socket: Socket)
+
+    func failConnect()
     {
-        self.socket = socket
+        maybeLog(message: "Failed to make a Transmission connection", logger: self.log)
+
+        switch self.connection
+        {
+            case .socket(let socket):
+                socket.close()
+                break
+            case .transport(var connection):
+                connection.stateUpdateHandler = nil
+                connection.cancel()
+                break
+        }
     }
-    
+
     public func read(size: Int) -> Data?
     {
         print("TransmissionLinux read called: \(#file), \(#line)")
@@ -52,27 +130,25 @@ public class TransmissionConnection: Connection
             return result
         }
 
-        do
+        guard let data = networkRead(size: size) else
         {
-            let _ = try self.socket.read(into: &buffer)
-
-            guard size <= buffer.count else
-            {
-                readLock.leave()
-                return nil
-            }
-
-            let result = Data(buffer[0..<size])
-            buffer = Data(buffer[size..<buffer.count])
-
             readLock.leave()
-            return result
+            return nil
         }
-        catch
+
+        buffer.append(data)
+
+        guard size <= buffer.count else
         {
             readLock.leave()
             return nil
         }
+
+        let result = Data(buffer[0..<size])
+        buffer = Data(buffer[size..<buffer.count])
+
+        readLock.leave()
+        return result
     }
 
     public func read(maxSize: Int) -> Data?
@@ -100,21 +176,64 @@ public class TransmissionConnection: Connection
         {
             // Buffer is empty, so we need to do a network read
 
-            do
+            var data: Data?
+
+            switch self.connection
             {
-                let readSize = try self.socket.read(into: &buffer)
+                case .socket(let socket):
+                    do
+                    {
+                        data = Data(repeating: 0, count: maxSize)
+                        let bytesRead = try socket.read(into: &data!)
+                        if (bytesRead < size)
+                        {
+                            data = Data(data![..<bytesRead])
+                        }
+                    }
+                    catch
+                    {
+                        readLock.leave()
+                        return nil
+                    }
 
-                let result = Data(buffer[0..<readSize])
-                buffer = Data(buffer[readSize..<buffer.count])
+                case .transport(let transport):
+                    let transportLock = DispatchGroup()
+                    transportLock.enter()
+                    transport.receive(minimumIncompleteLength: 1, maximumLength: maxSize)
+                    {
+                        maybeData, maybeContext, isComplete, maybeError in
 
-                readLock.leave()
-                return result
+                        guard let transportData = maybeData else
+                        {
+                            data = nil
+                            return
+                        }
+
+                        guard maybeError == nil else
+                        {
+                            data = nil
+                            return
+                        }
+
+                        data = transportData
+                    }
             }
-            catch
+
+            guard let bytes = data else
             {
                 readLock.leave()
                 return nil
             }
+
+            buffer.append(bytes)
+
+            let targetSize = min(maxSize, buffer.count)
+
+            let result = Data(buffer[0..<targetSize])
+            buffer = Data(buffer[targetSize..<buffer.count])
+
+            readLock.leave()
+            return result
         }
     }
 
@@ -134,20 +253,14 @@ public class TransmissionConnection: Connection
     public func write(data: Data) -> Bool
     {
         print("TransmissionLinux write called: \(#file), \(#line)")
+
         writeLock.enter()
 
-        do
-        {
-            try self.socket.write(from: data)
-        }
-        catch
-        {
-            writeLock.leave()
-            return false
-        }
+        let success = networkWrite(data: data)
 
         writeLock.leave()
-        return true
+
+        return success
     }
 
     public func readWithLengthPrefix(prefixSizeInBits: Int) -> Data?
@@ -159,109 +272,61 @@ public class TransmissionConnection: Connection
         switch prefixSizeInBits
         {
             case 8:
-                var lengthData = Data(repeating: 0, count: 1)
-
-                do
+                guard let lengthData = networkRead(size: prefixSizeInBits/8) else
                 {
-                    let readSize = try self.socket.read(into: &lengthData)
-
-                    guard readSize == 1 else
-                    {
-                        readLock.leave()
-                        return nil
-                    }
-
-                    guard let length8 = UInt8(maybeNetworkData: lengthData) else
-                    {
-                        readLock.leave()
-                        return nil
-                    }
-
-                    maybeLength = Int(length8)
+                    readLock.leave()
+                    return nil
                 }
-                catch
+
+                guard let boundedLength = UInt8(maybeNetworkData: lengthData) else
                 {
                     readLock.leave()
                     return nil
                 }
-            case 16:
-                var lengthData = Data(repeating: 0, count: 2)
 
-                do
+                maybeLength = Int(boundedLength)
+            case 16:
+                guard let lengthData = networkRead(size: prefixSizeInBits/8) else
                 {
-                    let readSize = try self.socket.read(into: &lengthData)
-
-                    guard readSize == 2 else
-                    {
-                        readLock.leave()
-                        return nil
-                    }
-
-                    guard let length16 = UInt16(maybeNetworkData: lengthData) else
-                    {
-                        readLock.leave()
-                        return nil
-                    }
-
-                    maybeLength = Int(length16)
+                    readLock.leave()
+                    return nil
                 }
-                catch
+
+                guard let boundedLength = UInt16(maybeNetworkData: lengthData) else
                 {
                     readLock.leave()
                     return nil
                 }
-            case 32:
-                var lengthData = Data(repeating: 0, count: 4)
 
-                do
+                maybeLength = Int(boundedLength)
+            case 32:
+                guard let lengthData = networkRead(size: prefixSizeInBits/8) else
                 {
-                    let readSize = try self.socket.read(into: &lengthData)
-
-                    guard readSize == 4 else
-                    {
-                        readLock.leave()
-                        return nil
-                    }
-
-                    guard let length32 = UInt32(maybeNetworkData: lengthData) else
-                    {
-                        readLock.leave()
-                        return nil
-                    }
-
-                    maybeLength = Int(length32)
+                    readLock.leave()
+                    return nil
                 }
-                catch
+
+                guard let boundedLength = UInt32(maybeNetworkData: lengthData) else
                 {
                     readLock.leave()
                     return nil
                 }
-            case 64:
-                var lengthData = Data(repeating: 0, count: 8)
 
-                do
+                maybeLength = Int(boundedLength)
+            case 64:
+                guard let lengthData = networkRead(size: prefixSizeInBits/8) else
                 {
-                    let readSize = try self.socket.read(into: &lengthData)
-
-                    guard readSize == 8 else
-                    {
-                        readLock.leave()
-                        return nil
-                    }
-
-                    guard let length64 = UInt64(maybeNetworkData: lengthData) else
-                    {
-                        readLock.leave()
-                        return nil
-                    }
-
-                    maybeLength = Int(length64)
+                    readLock.leave()
+                    return nil
                 }
-                catch
+
+                guard let boundedLength = UInt64(maybeNetworkData: lengthData) else
                 {
                     readLock.leave()
                     return nil
                 }
+
+                maybeLength = Int(boundedLength)
             default:
                 readLock.leave()
                 return nil
@@ -273,25 +338,13 @@ public class TransmissionConnection: Connection
             return nil
         }
 
-        var data = Data(repeating: 0, count: length)
-
-        do
-        {
-            let readSize = try self.socket.read(into: &data)
-
-            guard readSize == length else
-            {
-                readLock.leave()
-                return nil
-            }
-
-            return data
-        }
-        catch
+        guard let data = networkRead(size: length) else
         {
             readLock.leave()
             return nil
         }
+
+        return data
     }
 
     public func writeWithLengthPrefix(data: Data, prefixSizeInBits: Int) -> Bool
@@ -300,107 +353,143 @@ public class TransmissionConnection: Connection
 
         let length = data.count
 
+        var maybeLengthData: Data? = nil
+
         switch prefixSizeInBits
         {
             case 8:
-                let length8 = UInt8(length)
-
-                guard let lengthData = length8.maybeNetworkData else
-                {
-                    writeLock.leave()
-                    return false
-                }
-
-                do
-                {
-                    try self.socket.write(from: lengthData)
-                    try self.socket.write(from: data)
+                let boundedLength = UInt8(length)
+                maybeLengthData = boundedLength.maybeNetworkData
+            case 16:
+                let boundedLength = UInt16(length)
+                maybeLengthData = boundedLength.maybeNetworkData
+            case 32:
+                let boundedLength = UInt32(length)
+                maybeLengthData = boundedLength.maybeNetworkData
+            case 64:
+                let boundedLength = UInt64(length)
+                maybeLengthData = boundedLength.maybeNetworkData
+            default:
+                maybeLengthData = nil
+        }
 
-                    writeLock.leave()
-                    return true
-                }
-                catch
-                {
-                    writeLock.leave()
-                    return false
-                }
+        guard let lengthData = maybeLengthData else
+        {
+            writeLock.leave()
+            return false
+        }
 
-            case 16:
-                let length16 = UInt16(length)
+        let atomicData = lengthData + data
+        let success = networkWrite(data: atomicData)
+        writeLock.leave()
+        return success
+    }
+    
+    public func identifier() -> Int
+    {
+        return self.id
+    }
 
-                guard let lengthData = length16.maybeNetworkData else
-                {
-                    writeLock.leave()
-                    return false
-                }
+    func networkRead(size: Int) -> Data?
+    {
+        var data: Data?
 
+        switch self.connection
+        {
+            case .socket(let socket):
                 do
                 {
-                    try self.socket.write(from: lengthData)
-                    try self.socket.write(from: data)
-
-                    writeLock.leave()
-                    return true
+                    data = Data(repeating: 0, count: size)
+                    let bytesRead = try socket.read(into: &data!)
+                    if (bytesRead < size)
+                    {
+                        data = Data(data![..<bytesRead])
+                    }
                 }
                 catch
                 {
-                    writeLock.leave()
-                    return false
-                }
-
-            case 32:
-                let length32 = UInt32(length)
-
-                guard let lengthData = length32.maybeNetworkData else
-                {
-                    writeLock.leave()
-                    return false
+                    data = nil
+                    break
                 }
 
-                do
+            case .transport(let transport):
+                let transportLock = DispatchGroup()
+                transportLock.enter()
+                transport.receive(minimumIncompleteLength: size, maximumLength: size)
                 {
-                    try self.socket.write(from: lengthData)
-                    try self.socket.write(from: data)
+                    maybeData, maybeContext, isComplete, maybeError in
 
-                    writeLock.leave()
-                    return true
-                }
-                catch
-                {
-                    writeLock.leave()
-                    return false
-                }
+                    guard let transportData = maybeData else
+                    {
+                        data = nil
+                        return
+                    }
 
-            case 64:
-                let length64 = UInt8(length)
+                    guard maybeError == nil else
+                    {
+                        data = nil
+                        return
+                    }
 
-                guard let lengthData = length64.maybeNetworkData else
-                {
-                    writeLock.leave()
-                    return false
+                    data = transportData
                 }
+        }
+
+        return data
+    }
 
+    func networkWrite(data: Data) -> Bool
+    {
+        var success = false
+        switch self.connection
+        {
+            case .socket(let socket):
                 do
                 {
-                    try self.socket.write(from: lengthData)
-                    try self.socket.write(from: data)
-
-                    writeLock.leave()
-                    return true
+                    try socket.write(from: data)
+                    success = true
+                    break
                 }
                 catch
                 {
-                    writeLock.leave()
-                    return false
+                    success = false
+                    break
                 }
-
-            default:
-                writeLock.leave()
-                return false
+            case .transport(let transport):
+                let lock = DispatchGroup()
+                lock.enter()
+                transport.send(content: data, contentContext: .defaultMessage, isComplete: false, completion: NWConnection.SendCompletion.contentProcessed(
+                    {
+                        error in
+
+                        guard error == nil else
+                        {
+                            success = false
+                            lock.leave()
+                            return
+                        }
+
+                        success = true
+                        lock.leave()
+                        return
+                    }))
+                lock.wait()
         }
+
+        return success
     }
-    
-    public func identifier() -> Int {
-        return Int(self.socket.socketfd)
+}
+
+enum NetworkConnection
+{
+    case socket(Socket)
+    case transport(Transport.Connection)
+}
+
+public func maybeLog(message: String, logger: Logger? = nil) {
+    if logger != nil {
+        logger!.debug("\(message)")
+    } else {
+        print(message)
     }
 }