From e23457a647de7e0935a74c72271308fcc6d1416f Mon Sep 17 00:00:00 2001
From: Christian Ulrich <christian@ulrich.earth>
Date: Sat, 22 Aug 2020 12:34:12 +0200
Subject: [PATCH] cleanup on SIGINT

---
 punchd.nim | 34 +++++++++++++++++++++-------------
 1 file changed, 21 insertions(+), 13 deletions(-)

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: