diff --git a/punchd.nim b/punchd.nim index cea9779..9944819 100644 --- a/punchd.nim +++ b/punchd.nim @@ -11,6 +11,12 @@ from nativesockets import setSockOptInt const PunchdSocket = "/tmp/punchd.socket" type + Punchd = ref object + unixSocket: AsyncSocket + punchers: seq[TcpSyniPuncher] + + Sigint = object of CatchableError + # Requests TcpSyniConnect = object srcIp: IpAddress @@ -25,12 +31,11 @@ type srcPorts: array[3, Port] seqNums: seq[uint32] - Sigint = object of CatchableError - proc handleSigint() {.noconv.} = raise newException(Sigint, "received SIGINT") -proc handleRequest(line: string, unixSock: AsyncSocket) {.async.} = +proc handleRequest(punchd: Punchd, line: string, + unixSock: AsyncSocket) {.async.} = var id: string var sock: AsyncSocket var puncher: TcpSyniPuncher @@ -48,12 +53,14 @@ proc handleRequest(line: string, unixSock: AsyncSocket) {.async.} = seqNumbers.join(",")].join("|") await unixSock.send(&"progress|{id}|{content}\n") puncher = initPuncher(req.srcPorts[0], req.dstIp, req.dstPorts) + punchd.punchers.add(puncher) sock = await puncher.connect(handleSeqNumbers) of "tcp-syni-accept": let req = parseMessage[TcpSyniAccept](args[2]) puncher = initPuncher(req.srcPorts[0], req.dstIp, req.dstPorts, req.seqNums) + punchd.punchers.add(puncher) sock = await puncher.accept() else: @@ -67,35 +74,36 @@ proc handleRequest(line: string, unixSock: AsyncSocket) {.async.} = except ValueError: unixSock.close -proc handleRequests(userSock: AsyncSocket) {.async.} = +proc handleRequests(punchd: Punchd, userSock: AsyncSocket) {.async.} = while true: if userSock.isClosed: break let line = await userSock.recvLine(maxLength = 400) if line.len == 0: break - asyncCheck handleRequest(line, userSock) + asyncCheck punchd.handleRequest(line, userSock) -proc handleUsers(sock: AsyncSocket) {.async.} = +proc handleUsers(punchd: Punchd) {.async.} = while true: - let user = await sock.accept() - asyncCheck handleRequests(user) + let user = await punchd.unixSocket.accept() + asyncCheck punchd.handleRequests(user) proc main() = setControlCHook(handleSigint) removeFile(PunchdSocket) - let unixSocket = newAsyncSocket(AF_UNIX, SOCK_STREAM, IPPROTO_IP) - unixSocket.bindUnix(PunchdSocket) + let punchd = Punchd(unixSocket: newAsyncSocket(AF_UNIX, SOCK_STREAM, IPPROTO_IP)) + punchd.unixSocket.bindUnix(PunchdSocket) setFilePermissions(PunchdSocket, {fpUserRead, fpUserWrite, fpGroupRead, fpGroupWrite, fpOthersRead, fpOthersWrite}) - unixSocket.listen() - asyncCheck handleUsers(unixSocket) + punchd.unixSocket.listen() + asyncCheck handleUsers(punchd) try: runForever() except Sigint: - # FIXME: need to cleaup the punchers (i.e. delete firewall rules) + for puncher in punchd.punchers: + waitFor puncher.cleanup() removeFile(PunchdSocket) when isMainModule: