From 2b1cd9eeca8fdf18866aec7791d115a384469581 Mon Sep 17 00:00:00 2001
From: Christian Ulrich <christian@ulrich.earth>
Date: Wed, 18 Nov 2020 09:31:09 +0100
Subject: [PATCH] try IP_TTL = 2 for ACK packets

---
 puncher.nim | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/puncher.nim b/puncher.nim
index d78e68c..a62406b 100644
--- a/puncher.nim
+++ b/puncher.nim
@@ -1,6 +1,11 @@
 import asyncdispatch, asyncnet, net, port_prediction, strformat
 
-from nativesockets import SockAddr, SockAddr_storage, SockLen
+from nativesockets import
+  SockAddr,
+  SockAddr_storage,
+  SockLen,
+  getSockOptInt,
+  setSockOptInt
 from sequtils import any
 
 type
@@ -17,6 +22,9 @@ type
 
   PunchHoleError* = object of ValueError
 
+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
 
 proc `==`(a, b: Attempt): bool =
@@ -31,7 +39,8 @@ proc initPuncher*(sock: AsyncSocket): Puncher =
   Puncher(sock: sock)
 
 proc punch(puncher: Puncher, peerIp: IpAddress, peerPort: Port,
-           peerProbedPorts: seq[Port], msg: string): Future[Port] {.async.} =
+           peerProbedPorts: seq[Port], lowTTL: bool, msg: string): Future[Port]
+          {.async.} =
   let punchFuture = newFuture[Port]("punch")
   let predictedDstPorts = predictPortRange(peerPort, peerProbedPorts)
   let (_, myPort) = puncher.sock.getLocalAddr()
@@ -45,11 +54,17 @@ proc punch(puncher: Puncher, peerIp: IpAddress, peerPort: Port,
   var peerAddr: Sockaddr_storage
   var peerSockLen: SockLen
   try:
+    var defaultTTL: int
+    if lowTTL:
+      defaultTTL = puncher.sock.getFd.getSockOptInt(IPPROTO_IP, IP_TTL)
+      puncher.sock.getFd.setSockOptInt(IPPROTO_IP, IP_TTL, 2)
     for dstPort in attempt.dstPorts:
       toSockAddr(attempt.dstIp, dstPort, peerAddr, peerSockLen)
       # TODO: replace asyncdispatch.sendTo with asyncnet.sendTo (Nim 1.4 required)
       await sendTo(puncher.sock.getFd().AsyncFD, msg.cstring, msg.len,
                    cast[ptr SockAddr](addr peerAddr), peerSockLen)
+    if lowTTL:
+      puncher.sock.getFd.setSockOptInt(IPPROTO_IP, IP_TTL, defaultTTL)
     await punchFuture or sleepAsync(Timeout)
     if punchFuture.finished():
       result = punchFuture.read()
@@ -60,11 +75,11 @@ proc punch(puncher: Puncher, peerIp: IpAddress, peerPort: Port,
 
 proc initiate*(puncher: Puncher, peerIp: IpAddress, peerPort: Port,
                peerProbedPorts: seq[Port]): Future[Port] =
-  punch(puncher, peerIp, peerPort, peerProbedPorts, "SYN")
+  punch(puncher, peerIp, peerPort, peerProbedPorts, true, "SYN")
   
 proc respond*(puncher: Puncher, peerIp: IpAddress, peerPort: Port,
               peerProbedPorts: seq[Port]): Future[Port] =
-  punch(puncher, peerIp, peerPort, peerProbedPorts, "ACK")
+  punch(puncher, peerIp, peerPort, peerProbedPorts, false, "ACK")
 
 proc handleMsg*(puncher: Puncher, msg: string, peerIp: IpAddress,
                 peerPort: Port) =