Simple Swift TCP socket chat server

cnxshun.swift

import Foundation
import Network

internal struct FileHandleOutputStream: TextOutputStream {
    private let fileHandle: FileHandle
    let encoding: String.Encoding

    init(_ fileHandle: FileHandle, encoding: String.Encoding = .utf8) {
        self.fileHandle = fileHandle
        self.encoding = encoding
    }

    mutating func write(_ string: String) {
        if let data = string.data(using: encoding) {
            fileHandle.write(data)
        }
    }
}
internal var STDERR = FileHandleOutputStream(.standardError)
internal var STDOUT = FileHandleOutputStream(.standardOutput)

protocol ConnectionDelegate {
    func receive(data: Data)
    func messageSentSuccessfully()
    func messageFailedToSend(error: NWError)
    func connectionFailed(error: NWError)
    func connectionReady()
    func started()
    func connectionEnded()
    func stopped(error: NWError?)
}
class DebugDelegate: ConnectionDelegate {
    func receive(data: Data) {
        print(#function)
        print(data)
        print(String(bytes: data, encoding: .utf8))
    }

    func messageSentSuccessfully() {
        print(#function)
    }

    func messageFailedToSend(error: NWError) {
        print(#function)
        print(error)
    }

    func connectionFailed(error: NWError) {
        print(#function)
        print(error)
    }

    func connectionReady() {
        print(#function)
    }

    func start() {
        print(#function)
    }

    func connectionEnded() {
        print(#function)
    }

    func started() {
        print(#function)
    }

    func stopped(error: NWError?) {
        print(#function)
        print(error)
        onStop?(error)
    }

    var onStop: ((NWError?) -> Void)?
}
class ClosureDelegate: ConnectionDelegate {
    var receiveHandler: (Data) -> Void = { _ in }
    func receive(data: Data) {
        self.receiveHandler(data)
    }

    var messageSentSuccessfullyHandler: () -> Void = {}
    func messageSentSuccessfully() {
        self.messageSentSuccessfullyHandler()
    }

    var messageFailedToSendHandler: (NWError) -> Void = { _ in }
    func messageFailedToSend(error: NWError) {
        self.messageFailedToSendHandler(error)
    }

    var connectionFailedHandler: (NWError) -> Void = { _ in }
    func connectionFailed(error: NWError) {
        self.connectionFailedHandler(error)
    }

    var connectionReadyHandler: () -> Void = {}
    func connectionReady() {
        self.connectionReadyHandler()
    }

    var startedHandler: () -> Void = {}
    func started() {
        self.startedHandler()
    }

    var connectionEndedHandler: () -> Void = {}
    func connectionEnded() {
        self.connectionEndedHandler()
    }

    var stoppedHandler: (NWError?) -> Void = { _ in }
    func stopped(error: NWError?) {
        self.stoppedHandler(error)
    }
}
class Connection {
    private let nwConnection: NWConnection
    var delegate: ConnectionDelegate?

    init(connection: NWConnection) {
        self.nwConnection = connection
        self.id = Self.nextID
        Self.nextID += 1

        self.nwConnection.stateUpdateHandler = self.stateChanged(to:)
        self.setupReceive()
    }

    convenience init(host: NWEndpoint.Host, port: NWEndpoint.Port) {
        self.init(connection: NWConnection(host: host, port: port, using: .tcp))
    }


    private static var nextID: Int = 0
    let id: Int

    func start() {
        self.delegate?.started()
        self.nwConnection.start(queue: .main)
    }

    func send(data: Data) {
        self.nwConnection.send(content: data, contentContext: .defaultMessage, isComplete: true, completion: .contentProcessed { error in
            if let error = error {
                self.messageFailedToSend(error: error)
                return
            }
            self.messageSentSuccessfully()
        })
    }

    func stop() {
        self.nwConnection.send(content: nil, contentContext: .finalMessage, isComplete: true, completion: .idempotent)
        self.stop(error: nil)
    }

    private func messageFailedToSend(error: NWError) {
        self.delegate?.messageFailedToSend(error: error)
    }
    private func messageSentSuccessfully() {
        self.delegate?.messageSentSuccessfully()
    }

    private func stateChanged(to state: NWConnection.State) {
        switch state {
        case .setup:
            break
        case .waiting(let error):
            self.connectionFailed(error: error)
        case .preparing:
            break
        case .ready:
            self.connectionReady()
        case .failed(let error):
            self.connectionFailed(error: error)
        case .cancelled:
            break
        @unknown default:
            break
        }
    }

    private func connectionFailed(error: NWError) {
        self.delegate?.connectionFailed(error: error)
        self.stop(error: error)
    }
    private func connectionReady() {
        self.delegate?.connectionReady()
    }
    private func connectionEnded() {
        self.delegate?.connectionEnded()
        self.stop(error: nil)
    }

    private func setupReceive() {
        self.nwConnection.receive(minimumIncompleteLength: 1, maximumLength: 65536, completion: self.receiveMessage(data:context:isComplete:error:))
    }

    private func receiveMessage(data: Data?, context: NWConnection.ContentContext?, isComplete: Bool, error: NWError?) {
        if let data = data {
            self.delegate?.receive(data: data)
        }
        if isComplete {
            self.connectionEnded()
        } else if let error = error {
            self.connectionFailed(error: error)
        } else {
            self.setupReceive()
        }
    }

    private func stop(error: NWError?) {
        self.nwConnection.stateUpdateHandler = nil
        self.nwConnection.cancel()
        self.delegate?.stopped(error: error)
    }
}
class Client: ConnectionDelegate {
    func receive(data: Data) {
        if var str = String(bytes: data, encoding: .utf8) {
            if let last = str.last, last.isNewline { str = String(str.dropLast()) }
            let arr = str.split(separator: ":")
            switch arr.first {
            case "iam":
                print("[Server] \(arr.dropFirst().joined(separator: ":")) has joined the chat")
            case "iamnot":
                print("[Server] \(arr.dropFirst().joined(separator: ":")) has left the chat")
            case "server":
                print("[Server] \(arr.dropFirst().joined(separator: ":"))")
            case "msg":
                print("\(arr.dropFirst().first ?? "Unknown"): \(arr.dropFirst().dropFirst().joined(separator: ":"))")
            case "heartbeat":
                break
            default:
                print("[err] unknown message type")
            }
        } else {
            print("non-UTF8 data (\(data))")
        }
    }

    func messageSentSuccessfully() {

    }

    func messageFailedToSend(error: NWError) {
        print("Message failed to send", to: &STDERR)
    }

    func connectionFailed(error: NWError) {
        print("Connection failed", to: &STDERR)
    }

    func connectionReady() {

    }

    func connectionEnded() {

    }

    func started() {

    }

    func stopped(error: NWError?) {
        if error == nil {
            print("Exited")
            exit(EXIT_SUCCESS)
        } else {
            print("Exitied with error", to: &STDERR)
            exit(EXIT_FAILURE)
        }
    }

    init(host: String) {
        self.connection = Connection(host: NWEndpoint.Host.init(host), port: 12345)

        print("What's your name?")
        let input = readLine()
        if input == nil {
            print("Name must exist")
            exit(EXIT_FAILURE)
        }
        self.name = input!.replacingOccurrences(of: "[\\[\\]]", with: "", options: .regularExpression, range: nil)

        self.connection.delegate = self
    }

    let connection: Connection
    let name: String

    func start() {
        self.connection.start()
        self.connection.send(data: Data("iam:\(name)\n".utf8))
        DispatchQueue.init(label: "io", qos: .userInteractive).async {
            while true {
                let input = readLine()
                if let input = input {
                    self.connection.send(data: Data("msg:\(self.name):\(input)\n".utf8))
                } else {
                    self.connection.send(data: Data("iamnot:\(self.name)\n".utf8))
                    self.connection.stop()
                    break
                }
            }
        }
    }

    static func run(host: String) {
        let client = Client(host: host)
        client.start()
        dispatchMain()
    }
}
class Server {
    func receive(data: Data, from id: Int) {
        for (connectionID, connection) in self.connectionsByID {
            if connectionID == id { continue }
            connection.send(data: data)
        }
        if var str = String(bytes: data, encoding: .utf8) {
            if let last = str.last, last.isNewline { str = String(str.dropLast()) }
            let arr = str.split(separator: ":")
            switch arr.first {
            case "iam":
                print("[Server] \(arr.dropFirst().joined(separator: ":")) has joined the chat")
            case "iamnot":
                print("[Server] \(arr.dropFirst().joined(separator: ":")) has left the chat")
            case "server":
                print("[Server] \(arr.dropFirst().joined(separator: ":"))")
            case "msg":
                print("\(arr.dropFirst().first ?? "Unknown"): \(arr.dropFirst().dropFirst().joined(separator: ":"))")
            case "heartbeat":
                break
            default:
                print("[err] unknown message type")
            }
        } else {
            print("non-UTF8 data (\(data))")
        }
    }
    init() {
        self.listener = try! NWListener(using: .tcp, on: 12345)
        self.timer = DispatchSource.makeTimerSource(queue: .main)
    }

    let listener: NWListener
    let timer: DispatchSourceTimer

    func start() throws {
        print("server will start")
        self.listener.stateUpdateHandler = self.stateDidChange(to:)
        self.listener.newConnectionHandler = self.didAccept(nwConnection:)
        self.listener.start(queue: .main)

        self.timer.setEventHandler(handler: self.heartbeat)
        self.timer.schedule(deadline: .now() + 5.0, repeating: 5.0)
        self.timer.activate()

        DispatchQueue.init(label: "io", qos: .userInteractive).async {
            while true {
                let input = readLine()
                if let input = input {
                    let data = Data("server:\(input)".utf8)
                    for connection in self.connectionsByID.values {
                        connection.send(data: data)
                    }
                } else {
                    self.stop()
                    break
                }
            }
        }
    }

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

    private var connectionsByID: [Int: Connection] = [:]

    private func didAccept(nwConnection: NWConnection) {
        let connection = Connection(connection: nwConnection)
        self.connectionsByID[connection.id] = connection
        let d = ClosureDelegate()
        d.stoppedHandler = { _ in
            self.connectionDidStop(connection)
        }
        d.receiveHandler = { data in
            self.receive(data: data, from: connection.id)
        }
        connection.delegate = d
        connection.start()
    }

    private func connectionDidStop(_ connection: Connection) {
        self.connectionsByID.removeValue(forKey: connection.id)
        print("server did close connection \(connection.id)")
    }

    private func stop() {
        self.listener.stateUpdateHandler = nil
        self.listener.newConnectionHandler = nil
        self.listener.cancel()
        for connection in self.connectionsByID.values {
            connection.stop()
        }
        self.connectionsByID.removeAll()
        self.timer.cancel()
    }

    private func heartbeat() {
        let timestamp = Date()
        for connection in self.connectionsByID.values {
            let data = "heartbeat:connection: \(connection.id), timestamp: \(timestamp)\n"
            connection.send(data: Data(data.utf8))
        }
    }

    static func run() {
        let listener = Server()
        try! listener.start()
        dispatchMain()
    }
}

func main() {
    switch CommandLine.arguments.dropFirst() {
    case ["client"]: Client.run(host: CommandLine.arguments.dropFirst().dropFirst().first ?? "127.0.0.1")
    case ["server"]: Server.run()
    default:
        print("usage: NWTest server | client")
        exit(EXIT_FAILURE)
    }
}

main()