echod echo tcp/udp server with TLS-PSK support, using apple Network.framework

main.swift

//
//  main.swift
//  echod
//
//  Created by Ivo Vacek on 27/12/2018.
//  Copyright © 2018 Ivo Vacek. All rights reserved.
//

import Network
import Foundation
import Security

// see https://viewsourcecode.org/snaptoken/kilo/02.enteringRawMode.html
func enableRawMode(fileHandle: FileHandle) -> termios {
    var raw = termios()
    tcgetattr(fileHandle.fileDescriptor, &raw)
    let original = raw
    raw.c_lflag &= ~(UInt(ECHO | ICANON))
    tcsetattr(fileHandle.fileDescriptor, TCSAFLUSH, &raw);
    return original
}

func restoreRawMode(fileHandle: FileHandle, originalTerm: termios) {
    var term = originalTerm
    tcsetattr(fileHandle.fileDescriptor, TCSAFLUSH, &term);
}

let stdIn = FileHandle.standardInput

var char: UInt8 = 0

let usage = """
Usage:
echod [-uh] [-k key] [port]
-u          use udp instead of tcp
-h          help
-k key      TLS/DTLS with TLS_PSK key
-w timeout  cancel raw udp connection silently
            when receiveloop is idle for more
            then timeout seconds (default 10)

port        port number (default 7)
            if 0, port will be assign by the system
"""

let user = NSUserName()

// defaults
var nwparam = NWParameters.tcp
var isTCP = true
var type = "_echo._tcp"
var portNumber: UInt16 = 7
var psk: UnsafeMutablePointer<Int8>?
var wait = 10.0

var argc = CommandLine.argc

while case let option = getopt(CommandLine.argc, CommandLine.unsafeArgv, "uhk:w:"), option != -1 {
    let o = UnicodeScalar(CUnsignedChar(option))
    switch o {
    case "u":
        nwparam = NWParameters.udp
        type = "_echo._udp"
        isTCP = false
    case "k":
        psk = optarg
    case "h":
        print(usage)
        exit(0)
    case "w":
        guard let timeout = Double(String(cString: optarg)) else {
            print(usage)
            exit(0)
        }
        wait = timeout
    case "?":
        print(usage)
        exit(0)
    default:
        print(usage)
        exit(0)
    }
}
argc -= optind
if argc > 1 {
    print(usage)
    exit(0)
}
if argc == 1, let pn = UInt16(CommandLine.arguments[Int(optind)]) {
    portNumber = pn
}
guard let port = NWEndpoint.Port(rawValue: portNumber) else {
    print(usage)
    exit(0)
}
guard portNumber == 0 || portNumber > 1024 || user == "root" else {
    print(usage)
    print("Port:", portNumber, "requires special permission.")
    exit(0)
}


if let psk = psk {
    let dd = DispatchData(bytes: UnsafeRawBufferPointer(start: psk, count: strlen(psk)))
    let tlsOptions = NWProtocolTLS.Options()
    sec_protocol_options_add_pre_shared_key(tlsOptions.securityProtocolOptions, dd as __DispatchData, dd as __DispatchData)
    sec_protocol_options_add_tls_ciphersuite(tlsOptions.securityProtocolOptions, SSLCipherSuite(TLS_PSK_WITH_AES_128_GCM_SHA256))
    if isTCP {
        nwparam = NWParameters(tls: tlsOptions, tcp: NWProtocolTCP.Options())
    } else {
        nwparam = NWParameters(dtls: tlsOptions, udp: NWProtocolUDP.Options())
    }
}

let sq = DispatchQueue(label: "sq", qos: .default)
let dg = DispatchGroup()

var w: DispatchWorkItem?

func receive(conn: NWConnection) {
    
    if isTCP == false /*&& psk == nil*/ { // udp (conection-less) connection timeout
        w = DispatchWorkItem(block: { [weak conn] in
            guard let c = conn else { return }
            print(c.endpoint, "timeout, cancel()")
            c.cancel()
        })
        sq.asyncAfter(deadline: .now() + wait, execute: w!)
    }
    print(conn.endpoint, "start receive")
    conn.receive(minimumIncompleteLength: 0, maximumLength: Int(UInt16.max)) { [unowned conn] (d, c, f, e) in
        
        w?.cancel()
        w = nil
        print(conn.endpoint, "received", d)
        if let d = d, d.isEmpty == false {
            print(conn.endpoint, "send", d)
            conn.send(content: d, completion: .contentProcessed({ (e) in
                print(conn.endpoint, "sent")
                if let e = e {
                    print(conn.endpoint, "send error:", e, ", cancel()")
                    conn.cancel()
                } else {
                    receive(conn: conn)
                }
            }))
        } else {
            print(conn.endpoint, "nil or empty data received:", d, ", cancel()")
            conn.cancel()
        }
        if e != nil {
            print(conn.endpoint, "receive error:", e, ", cancel()")
            conn.cancel()
        }
    }
}

var listener: NWListener? = nil
let originalTerm = enableRawMode(fileHandle: stdIn)

var wi: DispatchWorkItem?
wi = DispatchWorkItem {
    print("Error: \(type) port \(portNumber) already in use")
    restoreRawMode(fileHandle: stdIn, originalTerm: originalTerm)
    wi = nil
    listener = nil
    exit(0)
}

do {
    listener = try NWListener(using: nwparam, on: port)
    if let l = listener {
        l.newConnectionHandler = { connection in
            connection.stateUpdateHandler = { [unowned connection] state in
                print("\(connection.endpoint)", state)
                switch state {
                case .ready:
                    receive(conn: connection)
                case .failed( _):
                    connection.cancel()
                default:
                    break
                }
            }
            connection.start(queue: sq)
        }
        l.service = NWListener.Service(name: "echo", type: type, domain: "local")
        l.stateUpdateHandler = { state in
            switch state {
            case .ready:
                let info =
                """
                Hello world!
                \(l.service!) is listening on port: \(l.port!)
                Press CTRL+D to exit
                
                """
                print(info)
                wi?.cancel()
            case .cancelled:
                print("cancelled")
                dg.leave()
            case .failed(let e):
                print("failed", e)
                l.cancel()
            case .waiting( _):
                print("waiting ...")
                sq.asyncAfter(deadline: .now() + 2.0, execute: wi!)
            case .setup:
                print("seting up ...")
            }
        }
        l.start(queue: sq)
    }
} catch let e {
    print(e)
    exit(0)
}

defer {
    dg.enter()
    listener?.cancel()
    print()
    print("By, by ...")
    dg.wait()
    wi = nil
    listener = nil
    // It would be also nice to disable raw input when exiting the app.
    restoreRawMode(fileHandle: stdIn, originalTerm: originalTerm)
}

while read(stdIn.fileDescriptor, &char, 1) == 1 {
    if char == 0x04 { // detect EOF (Ctrl+D)
        break
    }
    // don't echo stdin, just ignore the rest
}