## Nim-Libp2p
## Copyright (c) 2021 Status Research & Development GmbH
## Licensed under either of
##  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
##  * MIT license ([LICENSE-MIT](LICENSE-MIT))
## at your option.
## This file may not be copied, modified, or distributed except according to
## those terms.

{.push raises: [Defect].}

import std/[tables,
            strutils,
            strformat,
            sequtils,
            uri,
            parseutils]

import pkg/[chronos,
            chronos/apps/http/httptable,
            chronos/streams/asyncstream,
            chronos/streams/tlsstream,
            chronicles,
            httputils,
            stew/byteutils,
            stew/base64,
            stew/base10,
            nimcrypto/sha]

import ./utils, ./frame, ./session, /types, ./http, ./extensions/extutils

export utils, session, frame, types, http

logScope:
  topics = "ws-server"

type
  WSServer* = ref object of WebSocket
    protocols: seq[string]
    factories: seq[ExtFactory]

func toException(e: string): ref WebSocketError =
  (ref WebSocketError)(msg: e)

func toException(e: cstring): ref WebSocketError =
  (ref WebSocketError)(msg: $e)

func contains(extensions: openArray[Ext], extName: string): bool =
  for ext in extensions:
    if ext.name == extName:
      return true

proc getFactory(factories: openArray[ExtFactory], extName: string): ExtFactoryProc =
  for n in factories:
    if n.name == extName:
      return n.factory

proc selectExt(isServer: bool,
  extensions: var seq[Ext],
  factories: openArray[ExtFactory],
  exts: openArray[string]): string {.raises: [Defect, WSExtError].} =

  var extList: seq[AppExt]
  var response = ""
  for ext in exts:
    # each of "Sec-WebSocket-Extensions" can have multiple
    # extensions or fallback extension
    if not parseExt(ext, extList):
      raise newException(WSExtError, "extension syntax error: " & ext)

  for i, ext in extList:
    if extensions.contains(ext.name):
      # don't accept this fallback if prev ext
      # configuration already accepted
      trace "extension fallback not accepted", ext=ext.name
      continue

    # now look for right factory
    let factory = factories.getFactory(ext.name)
    if factory.isNil:
      # no factory? it's ok, just skip it
      trace "no extension factory", ext=ext.name
      continue

    let extRes = factory(isServer, ext.params)
    if extRes.isErr:
      # cannot create extension because of
      # wrong/incompatible params? skip or fallback
      trace "skip extension", ext=ext.name, msg=extRes.error
      continue

    let ext = extRes.get()
    doAssert(not ext.isNil)
    if i > 0:
      # add separator if more than one exts
      response.add ", "
    response.add ext.toHttpOptions

    # finally, accept the extension
    trace "extension accepted", ext=ext.name
    extensions.add ext

  # HTTP response for "Sec-WebSocket-Extensions"
  response

proc connect*(
  _: type WebSocket,
  uri: Uri,
  protocols: seq[string] = @[],
  factories: seq[ExtFactory] = @[],
  flags: set[TLSFlags] = {},
  version = WSDefaultVersion,
  frameSize = WSDefaultFrameSize,
  onPing: ControlCb = nil,
  onPong: ControlCb = nil,
  onClose: CloseCb = nil,
  rng: Rng = nil): Future[WSSession] {.async.} =
  ## create a new websockets client
  ##

  var rng = if isNil(rng): newRng() else: rng
  var key = Base64Pad.encode(genWebSecKey(rng))
  var uri = uri
  let client = case uri.scheme:
    of "wss":
      uri.scheme = "https"
      await TlsHttpClient.connect(uri.hostname, uri.port.parseInt(), tlsFlags = flags)
    of "ws":
      uri.scheme = "http"
      await HttpClient.connect(uri.hostname, uri.port.parseInt())
    else:
      raise newException(WSWrongUriSchemeError,
        "uri scheme has to be 'ws' or 'wss'")

  let headerData = [
    ("Connection", "Upgrade"),
    ("Upgrade", "websocket"),
    ("Cache-Control", "no-cache"),
    ("Sec-WebSocket-Version", $version),
    ("Sec-WebSocket-Key", key),
    ("Host", uri.hostname & ":" & uri.port)]

  var headers = HttpTable.init(headerData)
  if protocols.len > 0:
    headers.add("Sec-WebSocket-Protocol", protocols.join(", "))

  var extOffer = ""
  for i, f in factories:
    if i > 0:
      extOffer.add ", "
    extOffer.add f.clientOffer

  if extOffer.len > 0:
    headers.add("Sec-WebSocket-Extensions", extOffer)

  let response = try:
     await client.request(uri, headers = headers)
  except CatchableError as exc:
    trace "Websocket failed during handshake", exc = exc.msg
    await client.close()
    raise exc

  if response.code != Http101.toInt():
    raise newException(WSFailedUpgradeError,
          &"Server did not reply with a websocket upgrade: " &
          &"Header code: {response.code} Header reason: {response.reason} " &
          &"Address: {client.address}")

  let proto = response.headers.getString("Sec-WebSocket-Protocol")
  if proto.len > 0 and protocols.len > 0:
    if proto notin protocols:
      raise newException(WSFailedUpgradeError,
        &"Invalid protocol returned {proto}!")

  var extensions: seq[Ext]
  let exts = response.headers.getList("Sec-WebSocket-Extensions")
  discard selectExt(false, extensions, factories, exts)

  # Client data should be masked.
  let session = WSSession(
    stream: client.stream,
    readyState: ReadyState.Open,
    masked: true,
    extensions: system.move(extensions),
    rng: rng,
    frameSize: frameSize,
    onPing: onPing,
    onPong: onPong,
    onClose: onClose)

  for ext in session.extensions:
    ext.session = session

  return session

proc connect*(
  _: type WebSocket,
  address: TransportAddress,
  path: string,
  protocols: seq[string] = @[],
  factories: seq[ExtFactory] = @[],
  secure = false,
  flags: set[TLSFlags] = {},
  version = WSDefaultVersion,
  frameSize = WSDefaultFrameSize,
  onPing: ControlCb = nil,
  onPong: ControlCb = nil,
  onClose: CloseCb = nil,
  rng: Rng = nil): Future[WSSession] {.async.} =
  ## Create a new websockets client
  ## using a string path
  ##

  var uri = if secure:
      &"wss://"
    else:
      &"ws://"

  uri &= address.host & ":" & $address.port
  if path.startsWith("/"):
    uri.add path
  else:
    uri.add &"/{path}"

  return await WebSocket.connect(
    uri = parseUri(uri),
    protocols = protocols,
    factories = factories,
    flags = flags,
    version = version,
    frameSize = frameSize,
    onPing = onPing,
    onPong = onPong,
    onClose = onClose)

proc connect*(
  _: type WebSocket,
  host: string,
  port: Port,
  path: string,
  protocols: seq[string] = @[],
  factories: seq[ExtFactory] = @[],
  secure = false,
  flags: set[TLSFlags] = {},
  version = WSDefaultVersion,
  frameSize = WSDefaultFrameSize,
  onPing: ControlCb = nil,
  onPong: ControlCb = nil,
  onClose: CloseCb = nil,
  rng: Rng = nil): Future[WSSession] {.async.} =

  return await WebSocket.connect(
    address = initTAddress(host, port),
    path = path,
    protocols = protocols,
    factories = factories,
    secure = secure,
    flags = flags,
    version = version,
    frameSize = frameSize,
    onPing = onPing,
    onPong = onPong,
    onClose = onClose,
    rng = rng)

proc handleRequest*(
  ws: WSServer,
  request: HttpRequest,
  version: uint = WSDefaultVersion): Future[WSSession]
  {.
    async,
    raises: [
      Defect,
      WSHandshakeError,
      WSProtoMismatchError]
  .} =
  ## Creates a new socket from a request.
  ##

  if not request.headers.contains("Sec-WebSocket-Version"):
    raise newException(WSHandshakeError, "Missing version header")

  ws.version = Base10.decode(
    uint,
    request.headers.getString("Sec-WebSocket-Version"))
    .tryGet() # this method throws

  if ws.version != version:
    await request.stream.writer.sendError(Http426)
    trace "Websocket version not supported", version = ws.version

    raise newException(WSVersionError,
      &"Websocket version not supported, Version: {version}")

  ws.key = request.headers.getString("Sec-WebSocket-Key").strip()
  let wantProtos = if request.headers.contains("Sec-WebSocket-Protocol"):
        request.headers.getList("Sec-WebSocket-Protocol")
     else:
       @[""]

  let protos = wantProtos.filterIt(
    it in ws.protocols
  )

  let
    cKey = ws.key & WSGuid
    acceptKey = Base64Pad.encode(
      sha1.digest(cKey.toOpenArray(0, cKey.high)).data)

  var headers = HttpTable.init([
    ("Connection", "Upgrade"),
    ("Upgrade", "websocket"),
    ("Sec-WebSocket-Accept", acceptKey)])

  let protocol = if protos.len > 0: protos[0] else: ""
  if protocol.len > 0:
    headers.add("Sec-WebSocket-Protocol", protocol) # send back the first matching proto
  else:
    trace "Didn't match any protocol", supported = ws.protocols, requested = wantProtos

  # it is possible to have multiple "Sec-WebSocket-Extensions"
  let exts = request.headers.getList("Sec-WebSocket-Extensions")
  let extResp = selectExt(true, ws.extensions, ws.factories, exts)
  if extResp.len > 0:
    # send back any accepted extensions
    headers.add("Sec-WebSocket-Extensions", extResp)

  try:
    await request.sendResponse(Http101, headers = headers)
  except CancelledError as exc:
    raise exc
  except CatchableError as exc:
    raise newException(WSHandshakeError,
        "Failed to sent handshake response. Error: " & exc.msg)

  let session = WSSession(
    readyState: ReadyState.Open,
    stream: request.stream,
    proto: protocol,
    extensions: system.move(ws.extensions),
    masked: false,
    rng: ws.rng,
    frameSize: ws.frameSize,
    onPing: ws.onPing,
    onPong: ws.onPong,
    onClose: ws.onClose)

  for ext in session.extensions:
    ext.session = session

  return session

proc new*(
  _: typedesc[WSServer],
  protos: openArray[string] = [""],
  factories: openArray[ExtFactory] = [],
  frameSize = WSDefaultFrameSize,
  onPing: ControlCb = nil,
  onPong: ControlCb = nil,
  onClose: CloseCb = nil,
  rng: Rng = nil): WSServer =

  return WSServer(
    protocols: @protos,
    masked: false,
    rng: if isNil(rng): newRng() else: rng,
    frameSize: frameSize,
    factories: @factories,
    onPing: onPing,
    onPong: onPong,
    onClose: onClose)