From 006ec733615aeb3167e300ef8a064af66e56e1fa Mon Sep 17 00:00:00 2001
From: Christian Ulrich <christian@ulrich.earth>
Date: Thu, 27 Aug 2020 00:03:14 +0200
Subject: [PATCH] add timeouts

---
 tcp_syni.nim | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/tcp_syni.nim b/tcp_syni.nim
index 3e2fefc..49cbd88 100644
--- a/tcp_syni.nim
+++ b/tcp_syni.nim
@@ -1,4 +1,4 @@
-import asyncdispatch, asyncnet, strformat
+import asyncfutures, asyncdispatch, asyncnet, strformat
 from net import IpAddress, Port, `$`, `==`, getPrimaryIPAddr, toSockAddr
 from nativesockets import SockAddr, Sockaddr_storage, SockLen, setSockOptInt
 import asyncutils
@@ -9,6 +9,8 @@ import raw_socket
 var IPPROTO_IP {.importc: "IPPROTO_IP", header: "<netinet/in.h>".}: cint
 var IP_TTL {.importc: "IP_TTL", header: "<netinet/in.h>".}: cint
 
+const Timeout = 3000
+
 type
   TcpSyniPuncher* = ref object
     srcIp: IpAddress
@@ -157,22 +159,28 @@ proc connect*(puncher: TcpSyniPuncher,
   asyncCheck puncher.captureAndResendAck(captureAckFd, injectAckFd)
   await puncher.addFirewallRules()
   try:
-    result = await puncher.connectParallel()
+    let connectParallelFuture = puncher.connectParallel()
+    await connectParallelFuture or sleepAsync(Timeout)
     await puncher.cleanup()
+    if connectParallelFuture.finished():
+      result = connectParallelFuture.read()
+    else:
+      raise newException(PunchHoleError, "timeout")
   except OSError as e:
     raise newException(PunchHoleError, e.msg)
 
 proc prepareAccept(puncher: TcpSyniPuncher) {.async.} =
-  # FIXME: timeouts
   for dstPort in puncher.dstPorts:
     try:
       let sock = newAsyncSocket()
       sock.setSockOpt(OptReuseAddr, true)
       sock.getFd.setSockOptInt(IPPROTO_IP, IP_TTL, 2)
       sock.bindAddr(puncher.srcPort, $(puncher.srcIp))
-      await sock.connect($(puncher.dstIp), dstPort)
-      echo "connected during accept phase"
-      sock.close()
+      let connectFuture = sock.connect($(puncher.dstIp), dstPort)
+      await connectFuture or sleepAsync(Timeout)
+      if connectFuture.finished():
+        echo "connected during accept phase"
+        sock.close()
     except OSError:
       discard
 
@@ -180,7 +188,6 @@ proc accept*(puncher: TcpSyniPuncher): Future[AsyncSocket] {.async.} =
   await puncher.prepareAccept()
   await puncher.addFirewallRules()
   try:
-    # FIXME: timeout
     let rawFd = setupTcpInjectingSocket()
     for dstPort in puncher.dstPorts:
       for seqNum in puncher.seqNums:
@@ -200,8 +207,13 @@ proc accept*(puncher: TcpSyniPuncher): Future[AsyncSocket] {.async.} =
     sock.bindAddr(puncher.srcPort, $(puncher.srcIp))
     sock.listen()
     echo &"accepting connections from {puncher.dstIp}:{puncher.dstPorts[0].int}"
-    result = await sock.accept()
+    let acceptFuture = sock.accept()
+    await acceptFuture or sleepAsync(Timeout)
     await puncher.cleanup()
+    if acceptFuture.finished():
+      result = acceptFuture.read()
+    else:
+      raise newException(PunchHoleError, "timeout")
   except OSError as e:
     echo &"accepting connections from {puncher.dstIP}:{puncher.dstPorts[0].int} failed: ", e.msg
     await puncher.cleanup()