[Both] Adds support for installing apps over USB

AltServer/Connections/ConnectionManager.swift

@@ -54,10 +54,13 @@ class ConnectionManager
     private lazy var listener = self.makeListener()
     private let dispatchQueue = DispatchQueue(label: "com.rileytestut.AltServer.connections", qos: .utility)
     
-    private var connections = [NWConnection]()
+    private var connections = [ClientConnection]()
+    private var notificationConnections = [ALTDevice: NotificationConnection]()
     
     private init()
     {
+        NotificationCenter.default.addObserver(self, selector: #selector(ConnectionManager.deviceDidConnect(_:)), name: .deviceManagerDeviceDidConnect, object: nil)
+        NotificationCenter.default.addObserver(self, selector: #selector(ConnectionManager.deviceDidDisconnect(_:)), name: .deviceManagerDeviceDidDisconnect, object: nil)
     }
     
     func start()
@@ -77,6 +80,16 @@ class ConnectionManager
         default: break
         }
     }
+    
+    func disconnect(_ connection: ClientConnection)
+    {
+        connection.disconnect()
+        
+        if let index = self.connections.firstIndex(where: { $0 === connection })
+        {
+            self.connections.remove(at: index)
+        }
+    }
 }
 
 private extension ConnectionManager
@@ -127,67 +140,18 @@ private extension ConnectionManager
         }
         
         listener.newConnectionHandler = { [weak self] (connection) in
-            self?.awaitRequest(from: connection)
+            self?.prepare(connection)
         }
         
         return listener
     }
     
-    func disconnect(_ connection: NWConnection)
+    func prepare(_ connection: NWConnection)
     {
-        switch connection.state
-        {
-        case .cancelled, .failed:
-            print("Disconnecting from \(connection.endpoint)...")
-            
-            if let index = self.connections.firstIndex(where: { $0 === connection })
-            {
-                self.connections.remove(at: index)
-            }
-            
-        default:
-            // State update handler will call this method again.
-            connection.cancel()
-        }
-    }
-    
-    func process(data: Data?, error: NWError?, from connection: NWConnection) throws -> Data
-    {
-        do
-        {
-            do
-            {
-                guard let data = data else { throw error ?? ALTServerError(.unknown) }
-                return data
-            }
-            catch let error as NWError
-            {
-                print("Error receiving data from connection \(connection)", error)
-                
-                throw ALTServerError(.lostConnection)
-            }
-            catch
-            {
-                throw error
-            }
-        }
-        catch let error as ALTServerError
-        {
-            throw error
-        }
-        catch
-        {
-            preconditionFailure("A non-ALTServerError should never be thrown from this method.")
-        }
-    }
-}
-
-private extension ConnectionManager
-{
-    func awaitRequest(from connection: NWConnection)
-    {
-        guard !self.connections.contains(where: { $0 === connection }) else { return }
-        self.connections.append(connection)
+        let clientConnection = ClientConnection(connection: .wireless(connection))
+        
+        guard !self.connections.contains(where: { $0 === clientConnection }) else { return }
+        self.connections.append(clientConnection)
         
         connection.stateUpdateHandler = { [weak self] (state) in
             switch state
@@ -196,17 +160,17 @@ private extension ConnectionManager
                 
             case .ready:
                 print("Connected to client:", connection.endpoint)
-                self?.handleRequest(for: connection)
+                self?.handleRequest(for: clientConnection)
                 
             case .waiting:
                 print("Waiting for connection...")
                 
             case .failed(let error):
                 print("Failed to connect to service \(connection.endpoint).", error)
-                self?.disconnect(connection)
+                self?.disconnect(clientConnection)
                 
             case .cancelled:
-                self?.disconnect(connection)
+                self?.disconnect(clientConnection)
                 
             @unknown default: break
             }
@@ -214,17 +178,85 @@ private extension ConnectionManager
         
         connection.start(queue: self.dispatchQueue)
     }
-    
-    func handleRequest(for connection: NWConnection)
+}
+
+private extension ConnectionManager
+{
+    func startNotificationConnection(to device: ALTDevice)
     {
-        self.receiveRequest(from: connection) { (result) in
+        ALTDeviceManager.shared.startNotificationConnection(to: device) { (connection, error) in
+            guard let connection = connection else { return }
+            
+            let notifications: [CFNotificationName] = [.wiredServerConnectionAvailableRequest, .wiredServerConnectionStartRequest]
+            connection.startListening(forNotifications: notifications.map { String($0.rawValue) }) { (success, error) in
+                guard success else { return }
+                
+                connection.receivedNotificationHandler = { [weak self, weak connection] (notification) in
+                    guard let self = self, let connection = connection else { return }
+                    self.handle(notification, for: connection)
+                }
+                
+                self.notificationConnections[device] = connection
+            }
+        }
+    }
+    
+    func stopNotificationConnection(to device: ALTDevice)
+    {
+        guard let connection = self.notificationConnections[device] else { return }
+        connection.disconnect()
+        
+        self.notificationConnections[device] = nil
+    }
+    
+    func handle(_ notification: CFNotificationName, for connection: NotificationConnection)
+    {
+        switch notification
+        {
+        case .wiredServerConnectionAvailableRequest:
+            connection.sendNotification(.wiredServerConnectionAvailableResponse) { (success, error) in
+                if let error = error, !success
+                {
+                    print("Error sending wired server connection response.", error)
+                }
+                else
+                {
+                    print("Sent wired server connection available response!")
+                }
+            }
+            
+        case .wiredServerConnectionStartRequest:
+            ALTDeviceManager.shared.startWiredConnection(to: connection.device) { (wiredConnection, error) in
+                if let wiredConnection = wiredConnection
+                {
+                    print("Started wired server connection!")
+                    
+                    let clientConnection = ClientConnection(connection: .wired(wiredConnection))
+                    self.handleRequest(for: clientConnection)
+                }
+                else if let error = error
+                {
+                    print("Error starting wired server connection.", error)
+                }
+            }
+            
+        default: break
+        }
+    }
+}
+
+private extension ConnectionManager
+{
+    func handleRequest(for connection: ClientConnection)
+    {
+        connection.receiveRequest() { (result) in
             print("Received initial request with result:", result)
             
             switch result
             {
             case .failure(let error):
                 let response = ErrorResponse(error: ALTServerError(error))
-                self.send(response, to: connection, shouldDisconnect: true) { (result) in
+                connection.send(response, shouldDisconnect: true) { (result) in
                     print("Sent error response with result:", result)
                 }
                 
@@ -236,34 +268,34 @@ private extension ConnectionManager
                 
             case .success:
                 let response = ErrorResponse(error: ALTServerError(.unknownRequest))
-                self.send(response, to: connection, shouldDisconnect: true) { (result) in
+                connection.send(response, shouldDisconnect: true) { (result) in
                     print("Sent unknown request response with result:", result)
                 }
             }
         }
     }
     
-    func handleAnisetteDataRequest(_ request: AnisetteDataRequest, for connection: NWConnection)
+    func handleAnisetteDataRequest(_ request: AnisetteDataRequest, for connection: ClientConnection)
     {
         AnisetteDataManager.shared.requestAnisetteData { (result) in
             switch result
             {
             case .failure(let error):
                 let errorResponse = ErrorResponse(error: ALTServerError(error))
-                self.send(errorResponse, to: connection, shouldDisconnect: true) { (result) in
+                connection.send(errorResponse, shouldDisconnect: true) { (result) in
                     print("Sent anisette data error response with result:", result)
                 }
                 
             case .success(let anisetteData):
                 let response = AnisetteDataResponse(anisetteData: anisetteData)
-                self.send(response, to: connection, shouldDisconnect: true) { (result) in
+                connection.send(response, shouldDisconnect: true) { (result) in
                     print("Sent anisette data response with result:", result)
                 }
             }
         }
     }
     
-    func handlePrepareAppRequest(_ request: PrepareAppRequest, for connection: NWConnection)
+    func handlePrepareAppRequest(_ request: PrepareAppRequest, for connection: ClientConnection)
     {
         var temporaryURL: URL?
         
@@ -278,19 +310,19 @@ private extension ConnectionManager
             switch result
             {
             case .failure(let error):
-                print("Failed to process request from \(connection.endpoint).", error)
+                print("Failed to process request from \(connection).", error)
                 
                 let response = ErrorResponse(error: ALTServerError(error))
-                self.send(response, to: connection, shouldDisconnect: true) { (result) in
-                    print("Sent install app error response to \(connection.endpoint) with result:", result)
+                connection.send(response, shouldDisconnect: true) { (result) in
+                    print("Sent install app error response to \(connection) with result:", result)
                 }
                 
             case .success:
-                print("Processed request from \(connection.endpoint).")
+                print("Processed request from \(connection).")
                 
                 let response = InstallationProgressResponse(progress: 1.0)
-                self.send(response, to: connection, shouldDisconnect: true) { (result) in
-                    print("Sent install app response to \(connection.endpoint) with result:", result)
+                connection.send(response, shouldDisconnect: true) { (result) in
+                    print("Sent install app response to \(connection) with result:", result)
                 }
             }
         }
@@ -306,7 +338,7 @@ private extension ConnectionManager
                 
                 print("Awaiting begin installation request...")
                 
-                self.receiveRequest(from: connection) { (result) in
+                connection.receiveRequest() { (result) in
                     print("Received begin installation request with result:", result)
                     
                     switch result
@@ -326,8 +358,8 @@ private extension ConnectionManager
                         
                     case .success:
                         let response = ErrorResponse(error: ALTServerError(.unknownRequest))
-                        self.send(response, to: connection, shouldDisconnect: true) { (result) in
-                            print("Sent unknown request error response to \(connection.endpoint) with result:", result)
+                        connection.send(response, shouldDisconnect: true) { (result) in
+                            print("Sent unknown request error response to \(connection) with result:", result)
                         }
                     }
                 }
@@ -335,17 +367,15 @@ private extension ConnectionManager
         }
     }
     
-    func receiveApp(for request: PrepareAppRequest, from connection: NWConnection, completionHandler: @escaping (Result<URL, ALTServerError>) -> Void)
+    func receiveApp(for request: PrepareAppRequest, from connection: ClientConnection, completionHandler: @escaping (Result<URL, ALTServerError>) -> Void)
     {
-        connection.receive(minimumIncompleteLength: request.contentSize, maximumLength: request.contentSize) { (data, _, _, error) in
+        connection.receiveData(expectedBytes: request.contentSize) { (result) in
             do
             {
                 print("Received app data!")
                 
-                let data = try self.process(data: data, error: error, from: connection)
-                
-                print("Processed app data!")
-                
+                let data = try result.get()
+                                
                 guard ALTDeviceManager.shared.availableDevices.contains(where: { $0.identifier == request.udid }) else { throw ALTServerError(.deviceNotFound) }
                 
                 print("Writing app data...")
@@ -366,7 +396,7 @@ private extension ConnectionManager
         }
     }
     
-    func installApp(at fileURL: URL, toDeviceWithUDID udid: String, connection: NWConnection, completionHandler: @escaping (Result<Void, ALTServerError>) -> Void)
+    func installApp(at fileURL: URL, toDeviceWithUDID udid: String, connection: ClientConnection, completionHandler: @escaping (Result<Void, ALTServerError>) -> Void)
     {
         let serialQueue = DispatchQueue(label: "com.altstore.ConnectionManager.installQueue", qos: .default)
         var isSending = false
@@ -397,7 +427,7 @@ private extension ConnectionManager
                 print("Progress:", progress.fractionCompleted)
                 let response = InstallationProgressResponse(progress: progress.fractionCompleted)
                 
-                self.send(response, to: connection) { (result) in                    
+                connection.send(response) { (result) in                    
                     serialQueue.async {
                         isSending = false
                     }
@@ -405,92 +435,19 @@ private extension ConnectionManager
             }
         })
     }
+}
 
-    func send<T: Encodable>(_ response: T, to connection: NWConnection, shouldDisconnect: Bool = false, completionHandler: @escaping (Result<Void, ALTServerError>) -> Void)
+private extension ConnectionManager
+{
+    @objc func deviceDidConnect(_ notification: Notification)
     {
-        func finish(_ result: Result<Void, ALTServerError>)
-        {
-            completionHandler(result)
-            
-            if shouldDisconnect
-            {
-                // Add short delay to prevent us from dropping connection too quickly.
-                DispatchQueue.global().asyncAfter(deadline: .now() + 1.0) {
-                    self.disconnect(connection)
-                }
-            }
-        }
-        
-        do
-        {            
-            let data = try JSONEncoder().encode(response)
-            let responseSize = withUnsafeBytes(of: Int32(data.count)) { Data($0) }
-
-            connection.send(content: responseSize, completion: .contentProcessed { (error) in
-                do
-                {
-                    if let error = error
-                    {
-                        throw error
-                    }
-
-                    connection.send(content: data, completion: .contentProcessed { (error) in
-                        if error != nil
-                        {
-                            finish(.failure(.init(.lostConnection)))
-                        }
-                        else
-                        {
-                            finish(.success(()))
-                        }
-                    })
-                }
-                catch
-                {
-                    finish(.failure(.init(.lostConnection)))
-                }
-            })
-        }
-        catch
-        {
-            finish(.failure(.init(.invalidResponse)))
-        }
+        guard let device = notification.object as? ALTDevice else { return }
+        self.startNotificationConnection(to: device)
     }
     
-    func receiveRequest(from connection: NWConnection, completionHandler: @escaping (Result<ServerRequest, ALTServerError>) -> Void)
+    @objc func deviceDidDisconnect(_ notification: Notification)
     {
-        let size = MemoryLayout<Int32>.size
-        
-        print("Receiving request size")
-        connection.receive(minimumIncompleteLength: size, maximumLength: size) { (data, _, _, error) in
-            do
-            {
-                let data = try self.process(data: data, error: error, from: connection)
-                
-                print("Receiving request...")
-                
-                let expectedBytes = Int(data.withUnsafeBytes { $0.load(as: Int32.self) })
-                connection.receive(minimumIncompleteLength: expectedBytes, maximumLength: expectedBytes) { (data, _, _, error) in
-                    do
-                    {
-                        let data = try self.process(data: data, error: error, from: connection)
-                        
-                        let request = try JSONDecoder().decode(ServerRequest.self, from: data)
-                        
-                        print("Received installation request:", request)
-                        
-                        completionHandler(.success(request))
-                    }
-                    catch
-                    {
-                        completionHandler(.failure(ALTServerError(error)))
-                    }
-                }
-            }
-            catch
-            {
-                completionHandler(.failure(ALTServerError(error)))
-            }
-        }
+        guard let device = notification.object as? ALTDevice else { return }
+        self.stopNotificationConnection(to: device)
     }
 }