Swift networking with Apple's NWListener and NWConnection

NetworkConnection.swift

//
//  NetworkConnection.swift
//  Listen
//
//  Created by Ruurd Adema on 13/12/2019.
//  Copyright © 2019 Ruurd Adema. All rights reserved.
//

import Foundation
import Network

protocol NetworkConnectionDelegate: AnyObject
{
    func connectionOpened(connection: NetworkConnection)
    func connectionClosed(connection: NetworkConnection)
    func connectionError(connection: NetworkConnection, error: Error)
    func connectionReceivedData(connection: NetworkConnection, data: Data)
}

class NetworkConnection
{
    private static var nextID: Int = 0
    
    weak var delegate: NetworkConnectionDelegate?
    private let nwConnection: NWConnection
    let id: Int
    private var queue: DispatchQueue?
    
    init(nwConnection: NWConnection)
    {
        self.nwConnection = nwConnection
        self.id = NetworkConnection.nextID
        NetworkConnection.nextID += 1
    }
    
    func start(queue: DispatchQueue)
    {
        self.queue = queue
        self.nwConnection.stateUpdateHandler = self.onStateDidChange(to:)
        self.doReceive()
        self.nwConnection.start(queue: queue)
    }

    func send(data: Data)
    {
        let sizePrefix = withUnsafeBytes(of: UInt16(data.count).bigEndian) { Data($0) }
                
        log.info("Send \(data.count) bytes")
        
        self.nwConnection.batch {
            self.nwConnection.send(content: sizePrefix, completion: .contentProcessed( { error in
                if let error = error {
                    self.delegate?.connectionError(connection: self, error: error)
                    return
                }
            }))
            
            self.nwConnection.send(content: data, completion: .contentProcessed( { error in
                if let error = error {
                    self.delegate?.connectionError(connection: self, error: error)
                    return
                }
            }))
        }
    }

    private func onStateDidChange(to state: NWConnection.State)
    {
        switch state {
            case .setup:
                break
            case .waiting(let error):
                self.delegate?.connectionError(connection: self, error: error)
            case .preparing:
                break
            case .ready:
                self.delegate?.connectionOpened(connection: self)
            case .failed(let error):
                self.delegate?.connectionError(connection: self, error: error)
            case .cancelled:
                break
            default:
                break
        }
    }

    func close()
    {
        self.nwConnection.stateUpdateHandler = nil
        self.nwConnection.cancel()
        delegate?.connectionClosed(connection: self)
    }

    private func doReceive()
    {
        self.nwConnection.receive(minimumIncompleteLength: MemoryLayout<UInt16>.size, maximumLength: MemoryLayout<UInt16>.size) { (sizePrefixData, _, isComplete, error) in
            var sizePrefix: UInt16 = 0
            
            // Decode the size prefix
            if let data = sizePrefixData, !data.isEmpty
            {
                sizePrefix = data.withUnsafeBytes
                {
                    $0.bindMemory(to: UInt16.self)[0].bigEndian
                }
            }
            
            if isComplete
            {
                self.close()
            }
            else if let error = error
            {
                self.delegate?.connectionError(connection: self, error: error)
            }
            else
            {
                // If there is nothing to read
                if sizePrefix == 0
                {
                    log.error("Received size prefix of 0")
                    self.doReceive()
                    return
                }
                
                log.info("Read \(sizePrefix) bytes")
                
                // At this point we received a valid message and a valid size prefix 
                self.nwConnection.receive(minimumIncompleteLength: Int(sizePrefix), maximumLength: Int(sizePrefix)) { (data, _, isComplete, error) in
                    if let data = data, !data.isEmpty
                    {
                        self.delegate?.connectionReceivedData(connection: self, data: data)
                    }
                    if isComplete
                    {
                        self.close()
                    }
                    else if let error = error
                    {
                        self.delegate?.connectionError(connection: self, error: error)
                    }
                    else
                    {
                        self.doReceive()
                    }
                }
            }
        }
    }
}

NetworkServer.swift

//
//  NetworkServer.swift
//  Listen
//
//  Created by Ruurd Adema on 13/12/2019.
//  Copyright © 2019 Ruurd Adema. All rights reserved.
//

import Foundation
import Network

protocol NetworkServerDelegate: AnyObject
{
    func serverBecameReady()
    func connectionOpened(id: Int)
    func connectionReceivedData(id: Int, data: Data)
}

class NetworkServer: NetworkConnectionDelegate
{
    weak var delegate: NetworkServerDelegate?
    let listener: NWListener
    
    private var serverQueue: DispatchQueue?
    
    private var connectionsByID: [Int: NetworkConnection] = [:]
    
    init()
    {
        self.listener = try! NWListener(using: .tcp, on: 12345)
    }
    
    init(name: String?, type: String, domain: String?, txtRecord: NWTXTRecord)
    {
        self.listener = try! NWListener(using: .tcp, on: 0)
        self.listener.service = NWListener.Service(name: name, type: type, domain: domain, txtRecord: txtRecord)
    }
    
    func start(queue: DispatchQueue) throws
    {
        self.serverQueue = queue
        
        self.listener.stateUpdateHandler = self.onStateDidChange(to:)
        self.listener.newConnectionHandler = self.onNewConnectionAccepted(nwConnection:)
        self.listener.start(queue: queue)
    }

    func onStateDidChange(to newState: NWListener.State)
    {
        switch newState {
        case .setup:
            break
        case .waiting:
            break
        case .ready:
            self.delegate?.serverBecameReady()
            break
        case .failed(let error):
            print("server did fail, error: \(error)")
            self.stop()
        case .cancelled:
            break
        default:
            break
        }
    }

    private func onNewConnectionAccepted(nwConnection: NWConnection)
    {
        let connection = NetworkConnection(nwConnection: nwConnection)
        self.connectionsByID[connection.id] = connection
        connection.delegate = self
        connection.start(queue: self.serverQueue!)
        
        log.info("Server accepted connection \(connection)")
    }

    private func stop()
    {
        self.listener.stateUpdateHandler = nil
        self.listener.newConnectionHandler = nil
        self.listener.cancel()
        closeAllConnections()
    }

    private func heartbeat() {
        let timestamp = Date()
        print("server heartbeat, timestamp: \(timestamp)")
        for connection in self.connectionsByID.values {
            let data = "heartbeat, connection: \(connection.id), timestamp: \(timestamp)\r\n"
            connection.send(data: Data(data.utf8))
        }
    }
    
    func closeAllConnections()
    {
        for connection in self.connectionsByID.values{
            connection.close()
        }
        // Theoretically next call shouldn't have to remove anything as the connections
        // removed them selves by calling our connectionClosed function
        self.connectionsByID.removeAll()
    }
    
    func sendToAll(data: Data)
    {
        for conn in self.connectionsByID {
            conn.value.send(data: data)
        }
    }
    
    func sendTo(id: Int, data: Data)
    {
        self.connectionsByID[id]?.send(data: data)
    }
    
    func sendToAllExcept(id: Int, data: Data)
    {
        for conn in self.connectionsByID
        {
            if (conn.value.id != id)
            {
                conn.value.send(data: data)
            }
        }
    }
    
    var port: NWEndpoint.Port?
    {
        get {
            return self.listener.port
        }
    }
    
    // MARK: NetworkConnectionDelegate
    func connectionOpened(connection: NetworkConnection)
    {
        log.info("Server connection opened")
        self.delegate?.connectionOpened(id: connection.id)
    }
    
    func connectionClosed(connection: NetworkConnection)
    {
        self.connectionsByID.removeValue(forKey: connection.id)
        log.info("Server connection closed (\(connection))")
    }
    
    func connectionError(connection: NetworkConnection, error: Error)
    {
        log.error("Server connection error: \(error) (\(connection))")
    }
    
    func connectionReceivedData(connection: NetworkConnection, data: Data)
    {
        self.delegate?.connectionReceivedData(id: connection.id, data: data)
    }
}