punchd/punchd.nim

import asyncdispatch, asyncnet, os, strformat, strutils
from nativesockets import Domain, SockType, Protocol
from net import IpAddress, Port, `$`
import asyncutils
import message
import tcp_syni

from strutils import format, join
from nativesockets import setSockOptInt

const PunchdSocket = "/tmp/punchd.socket"

type
  Punchd = ref object
    unixSocket: AsyncSocket
    punchers: seq[TcpSyniPuncher]

  Sigint = object of CatchableError

  # Requests
  TcpSyniConnect = object
    srcIp: IpAddress
    srcPorts: array[3, Port]
    dstIp: IpAddress
    dstPorts: array[3, Port]

  TcpSyniAccept = object
    dstIp: IpAddress
    dstPorts: array[3, Port]
    srcIp: IpAddress
    srcPorts: array[3, Port]
    seqNums: seq[uint32]

proc handleSigint() {.noconv.} =
  raise newException(Sigint, "received SIGINT")

proc handleRequest(punchd: Punchd, line: string,
                   unixSock: AsyncSocket) {.async.} =
  var id: string
  var sock: AsyncSocket
  var puncher: TcpSyniPuncher
  try:
    let args = line.parseArgs(3)
    id = args[1]
    
    case args[0]:
    of "tcp-syni-connect":
      let req = parseMessage[TcpSyniConnect](args[2])
      proc handleSeqNumbers(seqNumbers: seq[uint32]) {.async.} =
        echo "progress! seqNumbers: ", seqNumbers
        let content = @["tcp-syni-accept", $req.srcIp, req.srcPorts.join(","),
                        $req.dstIp, req.dstPorts.join(","),
                        seqNumbers.join(",")].join("|")
        await unixSock.send(&"progress|{id}|{content}\n")
      puncher = initPuncher(req.srcPorts[0], req.dstIp, req.dstPorts)
      punchd.punchers.add(puncher)
      sock = await puncher.connect(handleSeqNumbers)
    
    of "tcp-syni-accept":
      let req = parseMessage[TcpSyniAccept](args[2])
      puncher = initPuncher(req.srcPorts[0], req.dstIp, req.dstPorts,
                            req.seqNums)
      punchd.punchers.add(puncher)
      sock = await puncher.accept()
    
    else:
      raise newException(ValueError, "invalid request")

    let unixFd = unixSock.getFd.AsyncFD
    await unixFd.asyncSendMsg(&"ok|{id}\n", @[fromFd(sock.getFd.AsyncFD)])

  except PunchHoleError as e:
    await unixSock.send(&"error|{id}|{e.msg}\n")
  except ValueError:
    unixSock.close

proc handleRequests(punchd: Punchd, userSock: AsyncSocket) {.async.} =
  while true:
    if userSock.isClosed:
      break
    let line = await userSock.recvLine(maxLength = 400)
    if line.len == 0:
      break
    asyncCheck punchd.handleRequest(line, userSock)

proc handleUsers(punchd: Punchd) {.async.} =
  while true:
    let user = await punchd.unixSocket.accept()
    asyncCheck punchd.handleRequests(user)

proc main() =
  setControlCHook(handleSigint)
  removeFile(PunchdSocket)
  let punchd = Punchd(unixSocket: newAsyncSocket(AF_UNIX, SOCK_STREAM, IPPROTO_IP))
  punchd.unixSocket.bindUnix(PunchdSocket)
  setFilePermissions(PunchdSocket,
                     {fpUserRead, fpUserWrite, fpGroupRead, fpGroupWrite,
                      fpOthersRead, fpOthersWrite})
  punchd.unixSocket.listen()
  asyncCheck handleUsers(punchd)
  try:
    runForever()

  except Sigint:
    for puncher in punchd.punchers:
      waitFor puncher.cleanup()
    punchd.unixSocket.close()
    removeFile(PunchdSocket)
    quit(0)

when isMainModule:
  main()