From 22648fa7c0506458aeefa7a644dd38cd7d3a8505 Mon Sep 17 00:00:00 2001 From: Ivan Vandot Date: Fri, 24 Mar 2023 16:08:47 +0100 Subject: [PATCH] chore: small code improvements --- lodns.nimble | 2 +- src/lodns.nim | 23 +++++++---------------- src/lodnspkg/actions.nim | 8 +++++--- src/lodnspkg/server.nim | 22 ++++++++++++---------- 4 files changed, 25 insertions(+), 30 deletions(-) diff --git a/lodns.nimble b/lodns.nimble index 1a684a9..c982276 100644 --- a/lodns.nimble +++ b/lodns.nimble @@ -1,6 +1,6 @@ # Package -version = "0.1.5" +version = "0.1.6" author = "vandot" description = "Simple DNS server for local development" license = "MIT" diff --git a/src/lodns.nim b/src/lodns.nim index 14e3e37..6bb503a 100644 --- a/src/lodns.nim +++ b/src/lodns.nim @@ -23,7 +23,6 @@ proc writeHelp() = quit() proc main() = - var install, start, uninstall = false var ip :string port :int @@ -33,35 +32,27 @@ proc main() = of cmdArgument: case key of "install": - install = true + (ip, port) = systemProbe() + install(ip, port, tld) of "start": - start = true + (ip, port) = systemProbe() + serve(ip, port, tld) of "uninstall": - uninstall = true + uninstall(tld) else: echo "unknown argument: ", key + writeHelp() of cmdLongOption, cmdShortOption: case key of "v", "version": writeVersion() - quit() of "h", "help": writeHelp() else: echo "unknown option: ", key + writeHelp() of cmdEnd: discard - if install or start: - (ip, port) = systemProbe() - if install: - install(ip, port, tld) - if uninstall: - uninstall(tld) - if start: - serve(ip, port, tld) - else: - writeHelp() - when isMainModule: main() diff --git a/src/lodnspkg/actions.nim b/src/lodnspkg/actions.nim index 0974019..bfccc6b 100644 --- a/src/lodnspkg/actions.nim +++ b/src/lodnspkg/actions.nim @@ -16,7 +16,7 @@ proc systemProbe*(): (string, int) = active.stripLineEnd if systemd != "systemd" or active != "active": echo "linux initialization is supported only for systemd using systemd-resolved" - quit() + quit(1) var version = execProcess("systemd --version | head -1 | awk '{ print $2}'") version.stripLineEnd # IPv4 dummy address (RFC7600) @@ -40,6 +40,7 @@ proc install*(ip: string, port: int, tld: string) = var exitCode = execCmd("systemctl restart systemd-networkd.service") if exitCode != 0: echo "systemctl restart systemd-networkd.service failed with code " & $exitCode + quit(1) when defined macosx: if not dirExists("/etc/resolver"): @@ -52,7 +53,7 @@ proc install*(ip: string, port: int, tld: string) = var exitCode = execCmd(addCommand) if exitCode != 0: echo addCommand & " failed with code " & $exitCode - quit() + quit(1) proc uninstall*(tld: string) = when defined linux: @@ -64,6 +65,7 @@ proc uninstall*(tld: string) = exitCode = execCmd("systemctl restart systemd-networkd.service") if exitCode != 0: echo "systemctl restart systemd-networkd.service failed with code " & $exitCode + quit(1) when defined macosx: removeFile("/etc/resolver/" & tld) @@ -73,4 +75,4 @@ proc uninstall*(tld: string) = var exitCode = execCmd(removeCommand) if exitCode != 0: echo removeCommand & " failed with code " & $exitCode - quit() + quit(1) diff --git a/src/lodnspkg/server.nim b/src/lodnspkg/server.nim index 5d32642..71b26a3 100644 --- a/src/lodnspkg/server.nim +++ b/src/lodnspkg/server.nim @@ -1,34 +1,36 @@ -import std/[net, logging, asyncnet, strutils, asyncdispatch] +import std/[net, asyncnet, strutils, asyncdispatch, times] import pkg/[dnsprotocol] +template log(lvl: string, msg: varargs[string, `$`]) = + stdout.writeLine("[" & getDateStr() & "T" & getClockStr() & "] " & lvl.toUpper() & " " & @msg) + proc serve*(ip: string, port: int, tld: string) = - var logger = newConsoleLogger(fmtStr="[$datetime] $levelname ") const withSize = 500 var rr :ResourceRecord - var socket = newAsyncSocket(sockType = SOCK_DGRAM, protocol = IPPROTO_UDP) + let socket = newAsyncSocket(sockType = SOCK_DGRAM, protocol = IPPROTO_UDP) socket.bindAddr(address = ip, port = Port port) - logger.log(lvlInfo, "status=listening ip=", ip, " port=", port) + log("info", "status=listening ip=", ip, " port=", port) while true: - let request = waitfor socket.recvFrom(withSize) + let request = waitFor socket.recvFrom(withSize) let message = parseMessage(request.data) let domain = message.questions[0].qname case message.questions[0].qtype of QType.A: - logger.log(lvlInfo, "type=A domain=", domain) + log("info", "type=A domain=", domain) rr = initResourceRecord(domain, Type.A, Class.IN, 299'i32, 4'u16, RDataA(address: [127'u8, 0, 0, 1])) of QType.AAAA: - logger.log(lvlInfo, "type=AAAA domain=", domain) + log("info", "type=AAAA domain=", domain) rr = initResourceRecord(domain, Type.AAAA, Class.IN, 299'i32, 16'u16, RDataAAAA(address: [0'u8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1])) else: - logger.log(lvlError, "type=", message.questions[0].qtype, " domain=", domain, " error=unsupported-type") + log("error", "type=", message.questions[0].qtype, " domain=", domain, " error=unsupported-type") rr = initResourceRecord(domain, Type.NULL, Class.IN, 0'i32, 0'u16, RDataNULL()) if rsplit(domain, ".", maxsplit=2)[1] != tld: - logger.log(lvlError, "type=", message.questions[0].qtype, " domain=", domain, " error=unsupported-domain") + log("error", "type=", message.questions[0].qtype, " domain=", domain, " error=unsupported-domain") rr = initResourceRecord(domain, Type.NULL, Class.IN, 0'i32, 0'u16, RDataNULL()) let header = initHeader(message.header.id, QR.Response) let reply = initMessage(header, message.questions, @[rr]) let bmsg = toBinMsg(reply) - waitfor socket.sendTo(request.address, request.port, bmsg) + waitFor socket.sendTo(request.address, request.port, bmsg)