punchd/tcp_nutss_responder.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.} =
  while puncher.attempts.len() != 0:
    await puncher.attempts.pop().deleteFirewallRules() # FIXME: needed?

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)