use inheritance (Puncher <- Initiator/Responder <- puncher implementations) to simplify logic in punchd.nim

punchd.nim

@@ -1,8 +1,8 @@
 import asyncdispatch, asyncnet, os, strformat, strutils
 from nativesockets import Domain, SockType, Protocol
-from net import IpAddress, Port, `$`
 import asyncutils
 import message
+import tables
 import tcp_syni_initiator
 import tcp_syni_responder
 import tcp_nutss_initiator
@@ -14,40 +14,11 @@ from nativesockets import setSockOptInt
 type
   Punchd = ref object
     unixSocket: AsyncSocket
-    tcpSyniInitiator: TcpSyniInitiator
+    initiators: Table[string, Initiator]
-    tcpSyniResponder: TcpSyniResponder
+    responders: Table[string, Responder]
-    tcpNutssInitiator: TcpNutssInitiator
-    tcpNutssResponder: TcpNutssResponder
 
   Sigint = object of CatchableError
 
-  # Requests
-  InitiateTcpSyni = object
-    srcIp: IpAddress
-    srcPorts: seq[Port]
-    dstIp: IpAddress
-    dstPorts: seq[Port]
-
-  RespondTcpSyni = object
-    dstIp: IpAddress
-    dstPorts: seq[Port]
-    srcIp: IpAddress
-    srcPorts: seq[Port]
-    seqNums: seq[uint32]
-
-  InitiateTcpNutss = object
-    srcIp: IpAddress
-    srcPorts: seq[Port]
-    dstIp: IpAddress
-    dstPorts: seq[Port]
-
-  RespondTcpNutss = object
-    dstIp: IpAddress
-    dstPorts: seq[Port]
-    srcIp: IpAddress
-    srcPorts: seq[Port]
-    extraData: string
-
 const PunchdSocket = "/tmp/punchd.socket"
 
 proc handleSigint() {.noconv.} =
@@ -69,46 +40,13 @@ proc handleRequest(punchd: Punchd, line: string,
 
     case args[0]:
     of "initiate":
-      case args[2]:
+      proc progress(extraArgs: string) {.async.} =
-      of "tcp-syni":
+        let msg = &"progress|{id}|{args[2]}|{args[3]}|{extraArgs}\n"
-        let req = parseMessage[InitiateTcpSyni](args[3])
+        await sendToClient(unixSock, msg)
-        proc progress(seqNumbers: seq[uint32]) {.async.} =
+      sock = await punchd.initiators[args[2]].initiate(args[3], progress)
-          echo "progress! seqNumbers: ", seqNumbers
-          let content = @["tcp-syni", $req.srcIp, req.srcPorts.join(","),
-                          $req.dstIp, req.dstPorts.join(","),
-                          seqNumbers.join(",")].join("|")
-          await sendToClient(unixSock, &"progress|{id}|{content}\n")
-        sock = await punchd.tcpSyniInitiator.initiate(req.srcPorts[0], req.dstIp,
-                                                      req.dstPorts, progress)
-
-      of "tcp-nutss":
-        let req = parseMessage[InitiateTcpNutss](args[3])
-        proc progress() {.async.} =
-          echo "progress!"
-          let content = @["tcp-nutss", $req.srcIp, req.srcPorts.join(","),
-                          $req.dstIp, req.dstPorts.join(","), ""].join("|")
-          await sendToClient(unixSock, &"progress|{id}|{content}\n")
-        sock = await punchd.tcpNutssInitiator.initiate(req.srcPorts[0], req.dstIp,
-                                                       req.dstPorts, progress)
-
-      else:
-        raise newException(ValueError, "invalid request")
-
 
     of "respond":
-      case args[2]:
+      sock = await punchd.responders[args[2]].respond(args[3])
-      of "tcp-syni":
-        let req = parseMessage[RespondTcpSyni](args[3])
-        sock = await punchd.tcpSyniResponder.respond(req.srcPorts[0], req.dstIp,
-                                                     req.dstPorts, req.seqNums)
-
-      of "tcp-nutss":
-        let req = parseMessage[RespondTcpNutss](args[3])
-        sock = await punchd.tcpNutssResponder.respond(req.srcPorts[0],
-                                                      req.dstIp, req.dstPorts)
-
-      else:
-        raise newException(ValueError, "invalid request")
 
     else:
       raise newException(ValueError, "invalid request")
@@ -118,7 +56,7 @@ proc handleRequest(punchd: Punchd, line: string,
 
   except PunchHoleError as e:
     await sendToClient(unixSock, &"error|{id}|{e.msg}\n")
-  except ValueError:
+  except KeyError, ValueError:
     unixSock.close
 
 proc handleRequests(punchd: Punchd, userSock: AsyncSocket) {.async.} =
@@ -145,20 +83,20 @@ proc main() =
   setFilePermissions(PunchdSocket,
                      {fpUserRead, fpUserWrite, fpGroupRead, fpGroupWrite,
                       fpOthersRead, fpOthersWrite})
-  let punchd = Punchd(unixSocket: unixSocket,
+  let punchd = Punchd(unixSocket: unixSocket)
-                      tcpSyniInitiator: initTcpSyniInitiator(),
+  punchd.initiators["tcp-syni"] = initTcpSyniInitiator()
-                      tcpSyniResponder: initTcpSyniResponder(),
+  punchd.initiators["tcp-nutss"] = initTcpNutssInitiator()
-                      tcpNutssInitiator: initTcpNutssInitiator(),
+  punchd.responders["tcp-syni"] = initTcpSyniResponder()
-                      tcpNutssResponder: initTcpNutssResponder())
+  punchd.responders["tcp-nutss"] = initTcpNutssResponder()
   asyncCheck handleUsers(punchd)
   try:
     runForever()
 
   except Sigint:
-    waitFor punchd.tcpSyniInitiator.cleanup()
+    for i in punchd.initiators.values:
-    waitFor punchd.tcpSyniResponder.cleanup()
+      waitFor i.cleanup()
-    waitFor punchd.tcpNutssInitiator.cleanup()
+    for r in punchd.responders.values:
-    waitFor punchd.tcpNutssResponder.cleanup()
+      waitFor r.cleanup()
     punchd.unixSocket.close()
     removeFile(PunchdSocket)
 

puncher.nim

@@ -1,16 +1,26 @@
-import asyncdispatch, strformat
+import asyncdispatch, asyncnet, strformat
-from net import IpAddress, Port, `$`
+from net import IpAddress, Port, `$`, `==`
 from sequtils import any
 import asyncutils
 
 type
-  Attempt = tuple | object
+  Attempt* = ref object of RootObj
+    srcIp*: IpAddress
+    srcPort*: Port
+    dstIp*: IpAddress
+    dstPorts*: seq[Port]
 
-  Puncher*[T: Attempt] = ref object
+  Puncher* = ref object of RootObj
-    attempts*: seq[T]
+    attempts*: seq[Attempt]
+
+  Initiator* = ref object of Puncher
+
+  Responder* = ref object of Puncher
 
   PunchHoleError* = object of ValueError
 
+  PunchProgressCb* = proc(extraArgs: string) {.async.}
+
 const Timeout* = 3000
 
 proc findAttempt*(puncher: Puncher, srcIp: IpAddress, srcPort: Port,
@@ -22,12 +32,26 @@ proc findAttempt*(puncher: Puncher, srcIp: IpAddress, srcPort: Port,
       return index
   return -1
 
-proc findAttemptsByLocalAddr*(puncher: Puncher[Attempt], address: IpAddress,
+proc findAttemptsByLocalAddr*(puncher: Puncher, address: IpAddress,
                               port: Port): seq[Attempt] =
   for attempt in puncher.attempts:
     if attempt.srcIp == address and attempt.srcPort == port:
       result.add(attempt)
 
+method cleanup*(puncher: Puncher): Future[void] {.base, async.} =
+  block: # workaround for https://github.com/nim-lang/Nim/issues/12530
+    raise newException(CatchableError, "Method without implementation override")
+
+method initiate*(puncher: Initiator, args: string, progress: PunchProgressCb):
+                Future[AsyncSocket] {.base, async.} =
+  block: # workaround for https://github.com/nim-lang/Nim/issues/12530
+    raise newException(CatchableError, "Method without implementation override")
+
+method respond*(puncher: Responder, args: string):
+               Future[AsyncSocket] {.base, async.} =
+  block: # workaround for https://github.com/nim-lang/Nim/issues/12530
+    raise newException(CatchableError, "Method without implementation override")
+
 proc makeFirewallRule(srcIp: IpAddress, srcPort: Port,
                       dstIp: IpAddress, dstPort: Port): string =
   # FIXME: use & instead of fmt?

tcp_nutss_initiator.nim

@@ -3,33 +3,34 @@ from net import IpAddress, Port, `$`, `==`, parseIpAddress
 from random import randomize, rand
 from sequtils import any
 import ip_packet
+import message
 import port_prediction
 import puncher
 import raw_socket
 import utils
 
-export PunchHoleError
+export Puncher, Initiator, PunchProgressCb, PunchHoleError, cleanup, initiate
 
 type
-  Attempt = object
+  TNIAttempt = ref object of Attempt
-    srcIp: IpAddress
-    srcPort: Port
-    dstIp: IpAddress
-    dstPorts: seq[Port]
     future: Future[AsyncSocket]
 
-  TcpNutssInitiator* = Puncher[Attempt]
+  TcpNutssInitiator* = ref object of Initiator
 
-  PunchProgressCb* = proc() {.async.}
+  Request = object
+    srcIp: IpAddress
+    srcPorts: seq[Port]
+    dstIp: IpAddress
+    dstPorts: seq[Port]
 
-proc cleanup*(puncher: TcpNutssInitiator) {.async.} =
+method cleanup*(puncher: TcpNutssInitiator) {.async.} =
   discard
 
 proc initTcpNutssInitiator*(): TcpNutssInitiator =
   randomize()
   TcpNutssInitiator()
 
-proc injectSynPackets(attempt: Attempt) {.async.} =
+proc injectSynPackets(attempt: TNIAttempt) {.async.} =
   let injectFd = setupTcpInjectingSocket()
   for dstPort in attempt.dstPorts:
     let synOut = IpPacket(protocol: tcp, ipAddrSrc: attempt.srcIp,
@@ -58,33 +59,34 @@ proc accept(puncher: TcpNutssInitiator, ip: IpAddress, port: Port) {.async.} =
         peer.close()
         continue
       else:
-        let attempt = puncher.attempts[i]
+        let attempt = TNIAttempt(puncher.attempts[i])
         attempt.future.complete(peer)
     let attempts = puncher.findAttemptsByLocalAddr(ip, port)
     if attempts.len() <= 1:
       break
   sock.close()
 
-proc initiate*(puncher: TcpNutssInitiator, srcPort: Port, dstIp: IpAddress,
+method initiate*(puncher: TcpNutssInitiator, args: string,
-               dstPorts: seq[Port], progressCb: PunchProgressCb):
+                 progressCb: PunchProgressCb): Future[AsyncSocket] {.async.} =
-               Future[AsyncSocket] {.async.} =
+  let req = parseMessage[Request](args)
-  let localIp = getPrimaryIPAddr(dstIp)
+  let localIp = getPrimaryIPAddr(req.dstIp)
-  let existingAttempts = puncher.findAttemptsByLocalAddr(localIp, srcPort)
+  let existingAttempts = puncher.findAttemptsByLocalAddr(localIp, req.srcPorts[0])
   if existingAttempts.len() == 0:
-    echo &"initiating connection to {dstIp}:{dstPorts[0].int}"
+    echo &"initiating connection to {req.dstIp}:{req.dstPorts[0].int}"
-    asyncCheck puncher.accept(localIp, srcPort)
+    asyncCheck puncher.accept(localIp, req.srcPorts[0])
   else:
     for a in existingAttempts:
-      if a.dstIp == dstIp and
+      if a.dstIp == req.dstIp and
-         a.dstPorts.any(proc (p: Port): bool = p in dstPorts):
+         a.dstPorts.any(proc (p: Port): bool = p in req.dstPorts):
         raise newException(PunchHoleError, "hole punching for given parameters already active")
   try:
-    let attempt = Attempt(srcIp: localIp, srcPort: srcPort, dstIp: dstIp,
+    let attempt = TNIAttempt(srcIp: localIp, srcPort: req.srcPorts[0],
-                          dstPorts: predictPortRange(dstPorts),
+                             dstIp: req.dstIp,
-                          future: newFuture[AsyncSocket]("initiate"))
+                             dstPorts: predictPortRange(req.dstPorts),
+                             future: newFuture[AsyncSocket]("initiate"))
     puncher.attempts.add(attempt)
     await attempt.injectSynPackets()
-    await progressCb()
+    await progressCb("")
     await attempt.future or sleepAsync(Timeout)
     puncher.attempts.del(puncher.attempts.find(attempt))
     if attempt.future.finished():

tcp_nutss_responder.nim

@@ -1,21 +1,23 @@
 import asyncdispatch, asyncnet, strformat
 from net import IpAddress, Port, `$`, `==`
+import message
 import port_prediction
 import puncher
 import utils
 
-export PunchHoleError
+export Puncher, Responder, PunchHoleError, cleanup, respond
 
 type
-  Attempt = object
+  TcpNutssResponder* = ref object of Responder
-    srcIp: IpAddress
+
-    srcPort: Port
+  Request = object
     dstIp: IpAddress
     dstPorts: seq[Port]
+    srcIp: IpAddress
+    srcPorts: seq[Port]
+    extraData: string
 
-  TcpNutssResponder* = Puncher[Attempt]
+method cleanup*(puncher: TcpNutssResponder) {.async.} =
-
-proc cleanup*(puncher: TcpNutssResponder) {.async.} =
   discard
 
 proc initTcpNutssResponder*(): TcpNutssResponder =
@@ -34,14 +36,15 @@ proc connect(srcIp: IpAddress, srcPort: Port, dstIp: IpAddress, dstPort: Port,
     echo &"connection {srcIP}:{srcPort.int} -> {dstIp}:{dstPort.int} failed: ", e.msg
     sock.close()
 
-proc respond*(puncher: TcpNutssResponder, srcPort: Port, dstIp: IpAddress,
+method respond*(puncher: TcpNutssResponder, args: string): Future[AsyncSocket] {.async.} =
-              dstPorts: seq[Port]): Future[AsyncSocket] {.async.} =
+  let req = parseMessage[Request](args)
-  let localIp = getPrimaryIPAddr(dstIp)
+  let localIp = getPrimaryIPAddr(req.dstIp)
   try:
     let connectFuture = newFuture[AsyncSocket]("respond")
-    let portRange = predictPortRange(dstPorts)
+    let portRange = predictPortRange(req.dstPorts)
     for dstPort in portRange:
-      asyncCheck connect(localIp, srcPort, dstIp, dstPort, connectFuture)
+      asyncCheck connect(localIp, req.srcPorts[0], req.dstIp, req.dstPorts[0],
+                         connectFuture)
     await connectFuture or sleepAsync(Timeout)
     if connectFuture.finished():
       result = connectFuture.read()

tcp_syni_initiator.nim

@@ -1,30 +1,30 @@
 import asyncdispatch, asyncnet, strformat
 from net import IpAddress, Port, `$`, `==`
 from nativesockets import setSockOptInt
+from strutils import join
 import ip_packet
+import message
 import network_interface
 import port_prediction
 import puncher
 import raw_socket
 import utils
 
-export PunchHoleError
+export Puncher, Initiator, PunchProgressCb, PunchHoleError, cleanup, initiate
 
 type
-  Attempt = object
+  TcpSyniInitiator* = ref object of Initiator
+
+  Request = object
     srcIp: IpAddress
-    srcPort: Port
+    srcPorts: seq[Port]
     dstIp: IpAddress
     dstPorts: seq[Port]
 
-  TcpSyniInitiator* = Puncher[Attempt]
-
-  PunchProgressCb* = proc(seqNums: seq[uint32]) {.async.}
-
 var IPPROTO_IP {.importc: "IPPROTO_IP", header: "<netinet/in.h>".}: cint
 var IP_TTL {.importc: "IP_TTL", header: "<netinet/in.h>".}: cint
 
-proc cleanup*(puncher: TcpSyniInitiator) {.async.} =
+method cleanup*(puncher: TcpSyniInitiator): Future[void] {.async.} =
   while puncher.attempts.len() != 0:
     await puncher.attempts.pop().deleteFirewallRules()
 
@@ -51,7 +51,7 @@ proc captureSeqNumbers(attempt: Attempt, cb: PunchProgressCb) {.async.} =
           seqNums.add(parsed.tcpSeqNumber)
           break
   closeSocket(captureFd)
-  await cb(seqNums)
+  await cb(seqNums.join(","))
 
 proc captureAndResendAck(attempt: Attempt) {.async.} =
   let iface = getNetworkInterface(attempt.srcIp)
@@ -92,14 +92,14 @@ proc connect(srcIp: IpAddress, srcPort: Port, dstIp: IpAddress, dstPort: Port,
     echo &"connection {srcIP}:{srcPort.int} -> {dstIp}:{dstPort.int} failed: ", e.msg
     sock.close()
 
-proc initiate*(puncher: TcpSyniInitiator, srcPort: Port, dstIp: IpAddress,
+method initiate*(puncher: TcpSyniInitiator, args: string,
-               dstPorts: seq[Port],
+                 progressCb: PunchProgressCb): Future[AsyncSocket] {.async.} =
-               progressCb: PunchProgressCb): Future[AsyncSocket] {.async.} =
+  let req = parseMessage[Request](args)
-  let localIp = getPrimaryIPAddr(dstIp)
+  let localIp = getPrimaryIPAddr(req.dstIp)
-  if puncher.findAttempt(localIp, srcPort, dstIp, dstPorts) != -1:
+  if puncher.findAttempt(localIp, req.srcPorts[0], req.dstIp, req.dstPorts) != -1:
     raise newException(PunchHoleError, "hole punching for given parameters already active")
-  let attempt = Attempt(srcIp: localIp, srcPort: srcPort, dstIp: dstIp,
+  let attempt = Attempt(srcIp: localIp, srcPort: req.srcPorts[0], dstIp: req.dstIp,
-                        dstPorts: predictPortRange(dstPorts))
+                        dstPorts: predictPortRange(req.dstPorts))
   puncher.attempts.add(attempt)
   await attempt.addFirewallRules()
   asyncCheck attempt.captureSeqNumbers(progressCb)

tcp_syni_responder.nim

@@ -3,25 +3,29 @@ from net import IpAddress, Port, `$`, `==`, parseIpAddress
 from random import randomize, rand
 from sequtils import any
 import ip_packet
+import message
 import port_prediction
 import puncher
 import raw_socket
 import utils
 
-export PunchHoleError
+export Puncher, Responder, PunchHoleError, cleanup, respond
 
 type
-  Attempt = object
+  TSRAttempt = ref object of Attempt
-    srcIp: IpAddress
-    srcPort: Port
-    dstIp: IpAddress
-    dstPorts: seq[Port]
     seqNums: seq[uint32]
     future: Future[AsyncSocket]
 
-  TcpSyniResponder* = Puncher[Attempt]
+  TcpSyniResponder* = ref object of Responder
 
-proc cleanup*(puncher: TcpSyniResponder) {.async.} =
+  Request = object
+    dstIp: IpAddress
+    dstPorts: seq[Port]
+    srcIp: IpAddress
+    srcPorts: seq[Port]
+    seqNums: seq[uint32]
+
+method cleanup*(puncher: TcpSyniResponder) {.async.} =
   while puncher.attempts.len() != 0:
     await puncher.attempts.pop().deleteFirewallRules()
 
@@ -29,7 +33,7 @@ proc initTcpSyniResponder*(): TcpSyniResponder =
   randomize()
   TcpSyniResponder()
 
-proc injectSynPackets(attempt: Attempt) {.async.} =
+proc injectSynPackets(attempt: TSRAttempt) {.async.} =
   let injectFd = setupTcpInjectingSocket()
   for dstPort in attempt.dstPorts:
     let synOut = IpPacket(protocol: tcp, ipAddrSrc: attempt.srcIp,
@@ -69,7 +73,7 @@ proc accept(puncher: TcpSyniResponder, srcIp: IpAddress,
         peer.close()
         continue
       else:
-        let attempt = puncher.attempts[i]
+        let attempt = TSRAttempt(puncher.attempts[i])
         attempt.future.complete(peer)
     let attempts = puncher.findAttemptsByLocalAddr(srcIp, srcPort)
     # FIXME: should attempts have timestamps, so we can decide here which ones to delete?
@@ -77,24 +81,25 @@ proc accept(puncher: TcpSyniResponder, srcIp: IpAddress,
       break
   sock.close()
 
-proc respond*(puncher: TcpSyniResponder, srcPort: Port, dstIp: IpAddress,
+method respond*(puncher: TcpSyniResponder, args: string):
-              dstPorts: seq[Port],
+               Future[AsyncSocket] {.async.} =
-              seqNums: seq[uint32]): Future[AsyncSocket] {.async.} =
+  let req = parseMessage[Request](args)
-  let localIp = getPrimaryIPAddr(dstIp)
+  let localIp = getPrimaryIPAddr(req.dstIp)
-  let existingAttempts = puncher.findAttemptsByLocalAddr(localIp, srcPort)
+  let existingAttempts = puncher.findAttemptsByLocalAddr(localIp, req.srcPorts[0])
   if existingAttempts.len() == 0:
-    echo &"accepting connections from {dstIp}:{dstPorts[0].int}"
+    echo &"accepting connections from {req.dstIp}:{req.dstPorts[0].int}"
-    asyncCheck puncher.accept(localIp, srcPort)
+    asyncCheck puncher.accept(localIp, req.srcPorts[0])
   else:
     for a in existingAttempts:
-      if a.dstIp == dstIp and
+      if a.dstIp == req.dstIp and
-         a.dstPorts.any(proc (p: Port): bool = p in dstPorts):
+         a.dstPorts.any(proc (p: Port): bool = p in req.dstPorts):
         raise newException(PunchHoleError, "hole punching for given parameters already active")
   try:
-    let attempt = Attempt(srcIp: localIp, srcPort: srcPort, dstIp: dstIp,
+    let attempt = TSRAttempt(srcIp: localIp, srcPort: req.srcPorts[0],
-                          dstPorts: predictPortRange(dstPorts),
+                             dstIp: req.dstIp,
-                          seqNums: seqNums,
+                             dstPorts: predictPortRange(req.dstPorts),
-                          future: newFuture[AsyncSocket]("accept"))
+                             seqNums: req.seqNums,
+                             future: newFuture[AsyncSocket]("respond"))
     puncher.attempts.add(attempt)
     await attempt.addFirewallRules() # FIXME: needed?
     await attempt.injectSynPackets()