From 383260c95ed2635b4601540c1899066f159fe2bd Mon Sep 17 00:00:00 2001
From: Christian Ulrich <christian@ulrich.earth>
Date: Sun, 25 Oct 2020 21:57:09 +0100
Subject: [PATCH] first implementation of UdpPuncher (untested)

---
 examples/app/app.nim |   2 +-
 punchd.nim           |   2 +
 udp.nim              | 111 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 114 insertions(+), 1 deletion(-)
 create mode 100644 udp.nim

diff --git a/examples/app/app.nim b/examples/app/app.nim
index ca22888..91df1ab 100644
--- a/examples/app/app.nim
+++ b/examples/app/app.nim
@@ -239,7 +239,7 @@ proc runApp(serverHostname: string, serverPort: Port, peerId: string,
                                                 srcPort)
     asyncCheck handleServerMessages(serverConn)
     let sock = await punchHole(punchdConn, serverConn, peerId, otherPeerId,
-                               "tcp-nutss")
+                               "udp")
     echo "connected!"
     await sock.send("ping")
     let msg = await sock.recv(4)
diff --git a/punchd.nim b/punchd.nim
index 816574c..9bbea6d 100644
--- a/punchd.nim
+++ b/punchd.nim
@@ -8,6 +8,7 @@ import options
 import tables
 import tcp_syni
 import tcp_nutss
+import udp
 
 type
   Punchd = ref object
@@ -151,6 +152,7 @@ proc main() =
   let punchd = Punchd(unixSocket: unixSocket)
   punchd.punchers["tcp-syni"] = initTcpSyniPuncher()
   punchd.punchers["tcp-nutss"] = initTcpNutssPuncher()
+  punchd.punchers["udp"] = initUdpPuncher()
   asyncCheck handleUsers(punchd)
   try:
     runForever()
diff --git a/udp.nim b/udp.nim
new file mode 100644
index 0000000..bcfc52b
--- /dev/null
+++ b/udp.nim
@@ -0,0 +1,111 @@
+from nativesockets import Protocol, SockType
+from net import IpAddress, Port, `$`, `==`
+import asyncdispatch
+import asyncnet
+import message
+import options
+import port_prediction
+import puncher
+import strformat
+import utils
+
+export puncher
+
+type
+  UdpPuncher* = ref object of Puncher
+
+  InitiateRequest = object
+    srcIp: IpAddress
+    srcPorts: seq[Port]
+    dstIp: IpAddress
+    dstPorts: seq[Port]
+
+  RespondRequest = object
+    dstIp: IpAddress
+    dstPorts: seq[Port]
+    srcIp: IpAddress
+    srcPorts: seq[Port]
+    extraArgs: string
+
+proc doInitiate(srcIp: IpAddress, srcPort: Port, dstIp: IpAddress,
+                dstPort: Port, initiateFuture: Future[AsyncSocket]) {.async.} =
+  let sock = newAsyncSocket(sockType = SOCK_DGRAM, protocol = IPPROTO_UDP)
+  sock.bindAddr(srcPort, $srcIp)
+  await sock.connect($dstIp, dstPort)
+  await sock.send("SYN")
+  let recvFuture = sock.recv(6) # "SYNACK"
+  await recvFuture or sleepAsync(Timeout)
+  if recvFuture.finished():
+    let msg = recvFuture.read()
+    echo &"connection {srcIp}:{srcPort.int} -> {dstIp}:{dstPort.int} succeeded: ", msg
+    await sock.send("ACK")
+    initiateFuture.complete(sock)
+  else:
+    echo &"connection {srcIp}:{srcPort.int} -> {dstIp}:{dstPort.int} timed out"
+
+proc doRespond(srcIp: IpAddress, srcPort: Port, dstIp: IpAddress,
+               dstPort: Port, respondFuture: Future[AsyncSocket]) {.async.} =
+  let sock = newAsyncSocket(sockType = SOCK_DGRAM, protocol = IPPROTO_UDP)
+  sock.bindAddr(srcPort, $srcIp)
+  await sock.connect($dstIp, dstPort)
+  await sock.send("SYNACK")
+  let recvFuture = sock.recv(3) # "ACK"
+  await recvFuture or sleepAsync(Timeout)
+  if recvFuture.finished():
+    let msg = recvFuture.read()
+    echo &"connection {srcIp}:{srcPort.int} -> {dstIp}:{dstPort.int} succeeded: ", msg
+    respondFuture.complete(sock)
+  else:
+    echo &"connection {srcIp}:{srcPort.int} -> {dstIp}:{dstPort.int} timed out"
+
+proc initUdpPuncher*(): UdpPuncher =
+  UdpPuncher()
+
+method parseInitiateRequest*(puncher: UdpPuncher, args: string): Attempt =
+  let parsed = parseMessage[InitiateRequest](args)
+  let localIp = getPrimaryIPAddr(parsed.dstIp)
+  let predictedDstPorts = predictPortRange(parsed.dstPorts)
+  Attempt(protocol: IPPROTO_UDP, srcIp: localIp, srcPort: parsed.srcPorts[0],
+          dstIp: parsed.dstIp, dstPorts: predictedDstPorts,
+          acceptFuture: none(Future[AsyncSocket]))
+
+method parseRespondRequest*(puncher: UdpPuncher, args: string): Attempt =
+  let parsed = parseMessage[RespondRequest](args)
+  let localIp = getPrimaryIPAddr(parsed.dstIp)
+  let predictedDstPorts = predictPortRange(parsed.dstPorts)
+  Attempt(protocol: IPPROTO_UDP, srcIp: localIp, srcPort: parsed.srcPorts[0],
+          dstIp: parsed.dstIp, dstPorts: predictedDstPorts,
+          acceptFuture: none(Future[AsyncSocket]))
+
+method initiate*(puncher: UdpPuncher, attempt: Attempt,
+                 progress: PunchProgressCb): Future[AsyncSocket] {.async.} =
+  assert(attempt.acceptFuture.isNone(), "expected attempt without acceptFuture")
+  try:
+    let initiateFuture = newFuture[AsyncSocket]("initiate")
+    for dstPort in attempt.dstPorts:
+      asyncCheck doInitiate(attempt.srcIp, attempt.srcPort, attempt.dstIp,
+                            dstPort, initiateFuture)
+    await progress("")
+    await initiateFuture or sleepAsync(Timeout)
+    if initiateFuture.finished():
+      result = initiateFuture.read()
+    else:
+      raise newException(PunchHoleError, "timeout")
+  except OSError as e:
+    raise newException(PunchHoleError, e.msg)
+
+method respond*(puncher: UdpPuncher, attempt: Attempt):
+               Future[AsyncSocket] {.async.} =
+  assert(attempt.acceptFuture.isNone(), "expected attempt without acceptFuture")
+  try:
+    let respondFuture = newFuture[AsyncSocket]("respond")
+    for dstPort in attempt.dstPorts:
+      asyncCheck doRespond(attempt.srcIp, attempt.srcPort, attempt.dstIp,
+                           dstPort, respondFuture)
+    await respondFuture or sleepAsync(Timeout)
+    if respondFuture.finished():
+      result = respondFuture.read()
+    else:
+      raise newException(PunchHoleError, "timeout")
+  except OSError as e:
+    raise newException(PunchHoleError, e.msg)