all repos — wsabi @ c64953fdfcd1e69b9900fb8da4218ab8fc81bb14

websocket proxy that sends stats to statsd

src/wsabi.nim (view raw)

 1
 2
 3
 4
 5
 6
 7
 8
 9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 
import
  ws,
  asyncdispatch,
  asynchttpserver,
  strformat,
  asyncnet,
  wsabipkg/args,
  strutils,
  json,
  uri


var
  server = newAsyncHttpServer()
  remote: WebSocket
  client: WebSocket
  clients: seq[WebSocket]

proc connectRemote(host: string) {.async.} =
  remote = await newWebSocket(host, protocol = "xmpp")
  let
    (address, port) = remote.tcpSocket.getPeerAddr()
  echo fmt"connected to {address}:{port.int}"

proc commRemoteClient() {.async.} =
  ## Fetch from remote and send to client
  while remote.readyState == Open:
    var data = await remote.receiveStrPacket()
    echo "from remote: ", data
    await client.send(data)

proc localServer(req: Request) {.async, gcsafe.} =
  ## Listen on localhost:PORT/ws
  if req.url.path == "/ws":
    var newreq = req

    # TODO: loop through for all remote hosts
    let bits = parseUri(remoteHost[0])
    newreq.headers["host"] = @[fmt"{bits.hostname}:{bits.port}"]
    newreq.headers["hostname"] = @[bits.hostname]

    client = await newWebSocket(newreq, protocol = "xmpp")

    try:
      while client.readyState == Open:
        # TODO: don't hardcode these replacements
        var
          packet = await client.receiveStrPacket()
          repacket = packet.replace("127.0.0.1", by=bits.hostname)
        echo "from local: " & repacket

        await remote.send(repacket)
    except WebSocketError:
      echo "client closed socket: ", getCurrentExceptionMsg()

  ## Health check endpoint 
  if req.url.path == "/check-health":
    if req.reqMethod == HttpPost:
      type 
        HealthCheck = object
          reqType: string
          host: string
          port: int

      let 
        hcJson = parseJson(req.body)
        hc = to(hcJson, HealthCheck)

      try:
        echo &"trying to reach {hc.host}:{hc.port}"
        discard await asyncnet.dial(hc.host, Port(hc.port))
        let r = %*{"status": "OK"}
        await req.respond(
          Http200, $r,
          newHttpHeaders([("Content-Type", "application/json")])
        )
      except OSError:
        echo "unable to reach specified host:port pair"
        let r = %*{
          "status": "error",
          "msg": "unable to reach specified host:port pair"
        }
        await req.respond(
          Http400, $r, 
          newHttpHeaders([("Content-Type", "application/json")])
        )


when isMainModule:
  parseArgs()
  echo "connecting to remote host..."
  waitFor connectRemote(remoteHost[0])
  echo fmt"local server running at ws://127.0.0.1:{localPort}/ws"
  asyncCheck server.serve(Port(localPort), localServer)
  asyncCheck commRemoteClient()

  runForever()