punchd/tcp_nutss_responder.nim

import asyncdispatch, asyncnet, strformat
from net import IpAddress, Port, `$`, `==`
import message
import port_prediction
import puncher
import utils

export Puncher, Responder, PunchHoleError, cleanup, respond

type
  TcpNutssResponder* = ref object of Responder

  Request = object
    dstIp: IpAddress
    dstPorts: seq[Port]
    srcIp: IpAddress
    srcPorts: seq[Port]
    extraData: string

method 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()

method respond*(puncher: TcpNutssResponder, args: string): Future[AsyncSocket] {.async.} =
  let req = parseMessage[Request](args)
  let localIp = getPrimaryIPAddr(req.dstIp)
  try:
    let connectFuture = newFuture[AsyncSocket]("respond")
    let portRange = predictPortRange(req.dstPorts)
    for dstPort in portRange:
      asyncCheck connect(localIp, req.srcPorts[0], req.dstIp, req.dstPorts[0],
                         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)