From 3f9cb6a3b1f73717ca189c3be81b44c70933166c Mon Sep 17 00:00:00 2001
From: Christian Ulrich <christian@ulrich.earth>
Date: Sat, 18 Jul 2020 15:04:06 +0200
Subject: [PATCH] only accept once and only use one injection socket

---
 tcp_syni.nim | 39 ++++++++++++++++++++-------------------
 1 file changed, 20 insertions(+), 19 deletions(-)

diff --git a/tcp_syni.nim b/tcp_syni.nim
index de9c676..a7e49e2 100644
--- a/tcp_syni.nim
+++ b/tcp_syni.nim
@@ -131,31 +131,34 @@ proc doConnect(srcIp: IpAddress, srcPort: Port, dstIp: IpAddress,
     echo &"connection {srcIP}:{srcPort.int} -> {dstIp}:{dstPort.int} failed: ", e.msg
     discard
 
-proc doAccept(srcIp: IpAddress, srcPort: Port, dstIp: IpAddress, dstPort: Port,
-              seqNums: seq[uint32], future: Future[AsyncSocket]) {.async.} =
-  try:
-    let sock = newAsyncSocket()
-    sock.setSockOpt(OptReuseAddr, true)
-    sock.getFd.setSockOptInt(IPPROTO_IP, IP_TTL, 2)
-    sock.bindAddr(srcPort, $srcIp)
-    await sock.connect($dstIp, dstPort)
-    echo "connected during accept phase"
-    sock.close()
-  except OSError:
-    discard
+proc doAccept(puncher: TcpSyniPuncher, future: Future[AsyncSocket]) {.async.} =
+  for dstPort in puncher.dstPorts:
+    # TODO: connect in parallel for better performance
+    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()
+    except OSError:
+      discard
   try:
     # FIXME: timeout
     let rawFd = setupTcpInjectingSocket()
-    asyncCheck injectSyns(rawFd, dstIp, dstPort, srcIp, srcPort, seqNums)
+    for dstPort in puncher.dstPorts:
+      asyncCheck injectSyns(rawFd, puncher.dstIp, dstPort, puncher.srcIp,
+                            puncher.srcPort, puncher.seqNums)
     let sock = newAsyncSocket()
     sock.setSockOpt(OptReuseAddr, true)
-    sock.setSockOpt(OptReusePort, true)
-    sock.bindAddr(srcPort, $srcIp)
+    sock.bindAddr(puncher.srcPort, $(puncher.srcIp))
     sock.listen()
+    echo &"accepting connections from {puncher.dstIp}:{puncher.dstPorts[0].int}"
     let connectedSock = await sock.accept()
     future.complete(connectedSock)
   except OSError as e:
-    echo &"accepting connection {dstIP}:{dstPort.int} -> {srcIp}:{srcPort.int} failed: ", e.msg
+    echo &"accepting connections from {puncher.dstIP}:{puncher.dstPorts[0].int} failed: ", e.msg
     discard
 
 proc connect*(puncher: TcpSyniPuncher,
@@ -170,6 +173,4 @@ proc connect*(puncher: TcpSyniPuncher,
 
 proc accept*(puncher: TcpSyniPuncher): Future[AsyncSocket] =
   result = newFuture[AsyncSocket]("tcp_syni.accept")
-  for dstPort in puncher.dstPorts:
-    asyncCheck doAccept(puncher.srcIp, puncher.srcPort, puncher.dstIp, dstPort,
-                        puncher.seqNums, result)
+  asyncCheck puncher.doAccept(result)