Last active
February 2, 2021 17:53
-
-
Save op183/776ef9b5ee0a77cd4dd759ff470e7b11 to your computer and use it in GitHub Desktop.
echod echo tcp/udp server with TLS-PSK support, using apple Network.framework
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
// | |
// 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 | |
} | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment