punchd/tcp_nutss_responder.nim

52 lines
1.5 KiB
Nim

import asyncdispatch, asyncnet, strformat
from net import IpAddress, Port, `$`, `==`
import port_prediction
import puncher
import utils
export PunchHoleError
type
Attempt = object
srcIp: IpAddress
srcPort: Port
dstIp: IpAddress
dstPorts: seq[Port]
TcpNutssResponder* = Puncher[Attempt]
proc cleanup*(puncher: TcpNutssResponder) {.async.} =
discard
proc initTcpNutssResponder*(): TcpNutssResponder =
TcpNutssResponder()
proc connect(srcIp: IpAddress, srcPort: Port, dstIp: IpAddress, dstPort: Port,
future: Future[AsyncSocket]) {.async.} =
let sock = newAsyncSocket()
sock.setSockOpt(OptReuseAddr, true)
echo &"connect {srcIp}:{srcPort} -> {dstIp}:{dstPort}"
sock.bindAddr(srcPort, $srcIp)
try:
await sock.connect($dstIp, dstPort)
future.complete(sock)
except OSError as e:
echo &"connection {srcIP}:{srcPort.int} -> {dstIp}:{dstPort.int} failed: ", e.msg
sock.close()
proc respond*(puncher: TcpNutssResponder, srcPort: Port, dstIp: IpAddress,
dstPorts: seq[Port]): Future[AsyncSocket] {.async.} =
let localIp = getPrimaryIPAddr(dstIp)
try:
let connectFuture = newFuture[AsyncSocket]("respond")
let portRange = predictPortRange(dstPorts)
for dstPort in portRange:
asyncCheck connect(localIp, srcPort, dstIp, dstPort, connectFuture)
await connectFuture or sleepAsync(Timeout)
if connectFuture.finished():
result = connectFuture.read()
else:
raise newException(PunchHoleError, "timeout")
except OSError as e:
raise newException(PunchHoleError, e.msg)