{.used.}

# Nim-Libp2p
# Copyright (c) 2023 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.

import bearssl, chronos, options
import ../libp2p
import ../libp2p/[protocols/connectivity/relay/relay,
                  protocols/connectivity/relay/messages,
                  protocols/connectivity/relay/utils,
                  protocols/connectivity/relay/client]
import ./helpers
import std/times
import stew/byteutils

proc createSwitch(r: Relay = nil, useYamux: bool = false): Switch =
  var builder = SwitchBuilder.new()
    .withRng(newRng())
    .withAddresses(@[ MultiAddress.init("/ip4/0.0.0.0/tcp/0").tryGet() ])
    .withTcpTransport()

  if useYamux:
    builder = builder.withYamux()
  else:
    builder = builder.withMplex()

  if r != nil:
    builder = builder.withCircuitRelay(r)

  return builder
    .withNoise()
    .build()

suite "Circuit Relay V2":

  suite "Reservation":
    asyncTeardown:
      await allFutures(src1.stop(), src2.stop(), rel.stop())
      checkTrackers()
    var
      ttl {.threadvar.}: int
      ldur {.threadvar.}: uint32
      ldata {.threadvar.}: uint64
      cl1 {.threadvar.}: RelayClient
      cl2 {.threadvar.}: RelayClient
      rv2 {.threadvar.}: Relay
      src1 {.threadvar.}: Switch
      src2 {.threadvar.}: Switch
      rel {.threadvar.}: Switch
      rsvp {.threadvar.}: Rsvp
      range {.threadvar.}: HSlice[times.DateTime, times.DateTime]

    asyncSetup:
      ttl = 3
      ldur = 60
      ldata = 2048
      cl1 = RelayClient.new()
      cl2 = RelayClient.new()
      rv2 = Relay.new(reservationTTL=initDuration(seconds=ttl),
                      limitDuration=ldur,
                      limitData=ldata,
                      maxCircuit=1)
      src1 = createSwitch(cl1)
      src2 = createSwitch(cl2)
      rel = createSwitch(rv2)

      await src1.start()
      await src2.start()
      await rel.start()
      await src1.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
      await src2.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
      rsvp = await cl1.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      range = now().utc + (ttl-3).seconds..now().utc + (ttl+3).seconds
      check:
        rsvp.expire.int64.fromUnix.utc in range
        rsvp.limitDuration == ldur
        rsvp.limitData == ldata

    asyncTest "Too many reservations":
      let conn = await cl2.switch.dial(rel.peerInfo.peerId, rel.peerInfo.addrs, RelayV2HopCodec)
      let pb = encode(HopMessage(msgType: HopMessageType.Reserve))
      await conn.writeLp(pb.buffer)
      let msg = HopMessage.decode(await conn.readLp(RelayMsgSize)).get()
      check:
        msg.msgType == HopMessageType.Status
        msg.status == Opt.some(StatusV2.ReservationRefused)

    asyncTest "Too many reservations + Reconnect":
      expect(ReservationError):
        discard await cl2.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      await rel.disconnect(src1.peerInfo.peerId)
      rsvp = await cl2.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      range = now().utc + (ttl-3).seconds..now().utc + (ttl+3).seconds
      check:
        rsvp.expire.int64.fromUnix.utc in range
        rsvp.limitDuration == ldur
        rsvp.limitData == ldata

    asyncTest "Reservation ttl expires":
      await sleepAsync(chronos.timer.seconds(ttl + 1))
      rsvp = await cl1.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      range = now().utc + (ttl-3).seconds..now().utc + (ttl+3).seconds
      check:
        rsvp.expire.int64.fromUnix.utc in range
        rsvp.limitDuration == ldur
        rsvp.limitData == ldata

    asyncTest "Reservation over relay":
      let
        rv2add = Relay.new()
        addrs = @[ MultiAddress.init($rel.peerInfo.addrs[0] & "/p2p/" &
                                     $rel.peerInfo.peerId & "/p2p-circuit").get() ]
      rv2add.setup(src2)
      await rv2add.start()
      src2.mount(rv2add)
      rv2.maxCircuit.inc()

      rsvp = await cl2.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      range = now().utc + (ttl-3).seconds..now().utc + (ttl+3).seconds
      check:
        rsvp.expire.int64.fromUnix.utc in range
        rsvp.limitDuration == ldur
        rsvp.limitData == ldata
      expect(ReservationError):
        discard await cl1.reserve(src2.peerInfo.peerId, addrs)

  for (useYamux, muxName) in [(false, "Mplex"), (true, "Yamux")]:
    suite "Circuit Relay V2 Connection using " & muxName:
      asyncTeardown:
        checkTrackers()
      var
        customProtoCodec {.threadvar.}: string
        proto {.threadvar.}: LPProtocol
        ttl {.threadvar.}: int
        ldur {.threadvar.}: uint32
        ldata {.threadvar.}: uint64
        srcCl {.threadvar.}: RelayClient
        dstCl {.threadvar.}: RelayClient
        rv2 {.threadvar.}: Relay
        src {.threadvar.}: Switch
        dst {.threadvar.}: Switch
        rel {.threadvar.}: Switch
        rsvp {.threadvar.}: Rsvp
        conn {.threadvar.}: Connection

      asyncSetup:
        customProtoCodec = "/test"
        proto = new LPProtocol
        proto.codec = customProtoCodec
        ttl = 60
        ldur = 120
        ldata = 16384
        srcCl = RelayClient.new()
        dstCl = RelayClient.new()
        src = createSwitch(srcCl, useYamux)
        dst = createSwitch(dstCl, useYamux)
        rel = createSwitch(nil, useYamux)

      asyncTest "Connection succeed":
        proto.handler = proc(conn: Connection, proto: string) {.async.} =
          check: "test1" == string.fromBytes(await conn.readLp(1024))
          await conn.writeLp("test2")
          check: "test3" == string.fromBytes(await conn.readLp(1024))
          await conn.writeLp("test4")
          await conn.close()
        rv2 = Relay.new(reservationTTL=initDuration(seconds=ttl),
                        limitDuration=ldur,
                        limitData=ldata)
        rv2.setup(rel)
        rel.mount(rv2)
        dst.mount(proto)

        await rel.start()
        await src.start()
        await dst.start()

        let addrs = MultiAddress.init($rel.peerInfo.addrs[0] & "/p2p/" &
                                  $rel.peerInfo.peerId & "/p2p-circuit").get()

        await src.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
        await dst.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)

        rsvp = await dstCl.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)

        conn = await src.dial(dst.peerInfo.peerId, @[ addrs ], customProtoCodec)
        await conn.writeLp("test1")
        check: "test2" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("test3")
        check: "test4" == string.fromBytes(await conn.readLp(1024))
        await allFutures(conn.close())
        await allFutures(src.stop(), dst.stop(), rel.stop())

      asyncTest "Connection duration exceeded":
        ldur = 3
        proto.handler = proc(conn: Connection, proto: string) {.async.} =
          check "wanna sleep?" == string.fromBytes(await conn.readLp(1024))
          await conn.writeLp("yeah!")
          check "go!" == string.fromBytes(await conn.readLp(1024))
          await sleepAsync(chronos.timer.seconds(ldur + 1))
          await conn.writeLp("that was a cool power nap")
          await conn.close()
        rv2 = Relay.new(reservationTTL=initDuration(seconds=ttl),
                        limitDuration=ldur,
                        limitData=ldata)
        rv2.setup(rel)
        rel.mount(rv2)
        dst.mount(proto)

        await rel.start()
        await src.start()
        await dst.start()

        let addrs = MultiAddress.init($rel.peerInfo.addrs[0] & "/p2p/" &
                                  $rel.peerInfo.peerId & "/p2p-circuit").get()

        await src.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
        await dst.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)

        rsvp = await dstCl.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
        conn = await src.dial(dst.peerInfo.peerId, @[ addrs ], customProtoCodec)
        await conn.writeLp("wanna sleep?")
        check: "yeah!" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("go!")
        expect(LPStreamEOFError):
          discard await conn.readLp(1024)
        await allFutures(conn.close())
        await allFutures(src.stop(), dst.stop(), rel.stop())

      asyncTest "Connection data exceeded":
        ldata = 1000
        proto.handler = proc(conn: Connection, proto: string) {.async.} =
          check "count me the better story you know" == string.fromBytes(await conn.readLp(1024))
          await conn.writeLp("do you expect a lorem ipsum or...?")
          check "surprise me!" == string.fromBytes(await conn.readLp(1024))
          await conn.writeLp("""Call me Ishmael. Some years ago--never mind how long
  precisely--having little or no money in my purse, and nothing
  particular to interest me on shore, I thought I would sail about a
  little and see the watery part of the world. It is a way I have of
  driving off the spleen and regulating the circulation. Whenever I
  find myself growing grim about the mouth; whenever it is a damp,
  drizzly November in my soul; whenever I find myself involuntarily
  pausing before coffin warehouses, and bringing up the rear of every
  funeral I meet; and especially whenever my hypos get such an upper
  hand of me, that it requires a strong moral principle to prevent me
  from deliberately stepping into the street, and methodically knocking
  people's hats off--then, I account it high time to get to sea as soon
  as I can. This is my substitute for pistol and ball. With a
  philosophical flourish Cato throws himself upon his sword; I quietly
  take to the ship.""")
        rv2 = Relay.new(reservationTTL=initDuration(seconds=ttl),
                        limitDuration=ldur,
                        limitData=ldata)
        rv2.setup(rel)
        rel.mount(rv2)
        dst.mount(proto)

        await rel.start()
        await src.start()
        await dst.start()

        let addrs = MultiAddress.init($rel.peerInfo.addrs[0] & "/p2p/" &
                                  $rel.peerInfo.peerId & "/p2p-circuit").get()

        await src.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
        await dst.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)

        rsvp = await dstCl.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
        conn = await src.dial(dst.peerInfo.peerId, @[ addrs ], customProtoCodec)
        await conn.writeLp("count me the better story you know")
        check: "do you expect a lorem ipsum or...?" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("surprise me!")
        expect(LPStreamEOFError):
          discard await conn.readLp(1024)
        await allFutures(conn.close())
        await allFutures(src.stop(), dst.stop(), rel.stop())

      asyncTest "Reservation ttl expire during connection":
        ttl = 3
        proto.handler = proc(conn: Connection, proto: string) {.async.} =
          check: "test1" == string.fromBytes(await conn.readLp(1024))
          await conn.writeLp("test2")
          check: "test3" == string.fromBytes(await conn.readLp(1024))
          await conn.writeLp("test4")
          await conn.close()
        rv2 = Relay.new(reservationTTL=initDuration(seconds=ttl),
                        limitDuration=ldur,
                        limitData=ldata)
        rv2.setup(rel)
        rel.mount(rv2)
        dst.mount(proto)

        await rel.start()
        await src.start()
        await dst.start()

        let addrs = MultiAddress.init($rel.peerInfo.addrs[0] & "/p2p/" &
                                  $rel.peerInfo.peerId & "/p2p-circuit").get()

        await src.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
        await dst.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)

        rsvp = await dstCl.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
        conn = await src.dial(dst.peerInfo.peerId, @[ addrs ], customProtoCodec)
        await conn.writeLp("test1")
        check: "test2" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("test3")
        check: "test4" == string.fromBytes(await conn.readLp(1024))
        await src.disconnect(rel.peerInfo.peerId)
        await sleepAsync(chronos.timer.seconds(ttl + 1))

        expect(DialFailedError):
          check: conn.atEof()
          await conn.close()
          await src.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
          conn = await src.dial(dst.peerInfo.peerId, @[ addrs ], customProtoCodec)
        await allFutures(conn.close())
        await allFutures(src.stop(), dst.stop(), rel.stop())

      asyncTest "Connection over relay":
        # src => rel => rel2 => dst
        # rel2 reserve rel
        # dst reserve rel2
        # src try to connect with dst
        proto.handler = proc(conn: Connection, proto: string) {.async.} =
          raise newException(CatchableError, "Should not be here")
        let
          rel2Cl = RelayClient.new(canHop = true)
          rel2 = createSwitch(rel2Cl, useYamux)
          rv2 = Relay.new()
        rv2.setup(rel)
        rel.mount(rv2)
        dst.mount(proto)
        await rel.start()
        await rel2.start()
        await src.start()
        await dst.start()

        let
          addrs = @[ MultiAddress.init($rel.peerInfo.addrs[0] & "/p2p/" &
                                       $rel.peerInfo.peerId & "/p2p-circuit/p2p/" &
                                       $rel2.peerInfo.peerId & "/p2p/" &
                                       $rel2.peerInfo.peerId & "/p2p-circuit").get() ]

        await src.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
        await rel2.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
        await dst.connect(rel2.peerInfo.peerId, rel2.peerInfo.addrs)

        rsvp = await rel2Cl.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
        let rsvp2 = await dstCl.reserve(rel2.peerInfo.peerId, rel2.peerInfo.addrs)

        expect(DialFailedError):
          conn = await src.dial(dst.peerInfo.peerId, addrs, customProtoCodec)
        if not conn.isNil():
          await allFutures(conn.close())
        await allFutures(src.stop(), dst.stop(), rel.stop(), rel2.stop())

      asyncTest "Connection using ClientRelay":
        var
          protoABC = new LPProtocol
          protoBCA = new LPProtocol
          protoCAB = new LPProtocol
        protoABC.codec = "/abctest"
        protoABC.handler = proc(conn: Connection, proto: string) {.async.} =
          check: "testABC1" == string.fromBytes(await conn.readLp(1024))
          await conn.writeLp("testABC2")
          check: "testABC3" == string.fromBytes(await conn.readLp(1024))
          await conn.writeLp("testABC4")
          await conn.close()
        protoBCA.codec = "/bcatest"
        protoBCA.handler = proc(conn: Connection, proto: string) {.async.} =
          check: "testBCA1" == string.fromBytes(await conn.readLp(1024))
          await conn.writeLp("testBCA2")
          check: "testBCA3" == string.fromBytes(await conn.readLp(1024))
          await conn.writeLp("testBCA4")
          await conn.close()
        protoCAB.codec = "/cabtest"
        protoCAB.handler = proc(conn: Connection, proto: string) {.async.} =
          check: "testCAB1" == string.fromBytes(await conn.readLp(1024))
          await conn.writeLp("testCAB2")
          check: "testCAB3" == string.fromBytes(await conn.readLp(1024))
          await conn.writeLp("testCAB4")
          await conn.close()

        let
          clientA = RelayClient.new(canHop = true)
          clientB = RelayClient.new(canHop = true)
          clientC = RelayClient.new(canHop = true)
          switchA = createSwitch(clientA, useYamux)
          switchB = createSwitch(clientB, useYamux)
          switchC = createSwitch(clientC, useYamux)

        switchA.mount(protoBCA)
        switchB.mount(protoCAB)
        switchC.mount(protoABC)

        await switchA.start()
        await switchB.start()
        await switchC.start()

        let
          addrsABC = MultiAddress.init($switchB.peerInfo.addrs[0] & "/p2p/" &
                                       $switchB.peerInfo.peerId & "/p2p-circuit").get()
          addrsBCA = MultiAddress.init($switchC.peerInfo.addrs[0] & "/p2p/" &
                                       $switchC.peerInfo.peerId & "/p2p-circuit").get()
          addrsCAB = MultiAddress.init($switchA.peerInfo.addrs[0] & "/p2p/" &
                                       $switchA.peerInfo.peerId & "/p2p-circuit").get()

        await switchA.connect(switchB.peerInfo.peerId, switchB.peerInfo.addrs)
        await switchB.connect(switchC.peerInfo.peerId, switchC.peerInfo.addrs)
        await switchC.connect(switchA.peerInfo.peerId, switchA.peerInfo.addrs)
        let rsvpABC = await clientA.reserve(switchC.peerInfo.peerId, switchC.peerInfo.addrs)
        let rsvpBCA = await clientB.reserve(switchA.peerInfo.peerId, switchA.peerInfo.addrs)
        let rsvpCAB = await clientC.reserve(switchB.peerInfo.peerId, switchB.peerInfo.addrs)
        let connABC = await switchA.dial(switchC.peerInfo.peerId, @[ addrsABC ], "/abctest")
        let connBCA = await switchB.dial(switchA.peerInfo.peerId, @[ addrsBCA ], "/bcatest")
        let connCAB = await switchC.dial(switchB.peerInfo.peerId, @[ addrsCAB ], "/cabtest")

        await connABC.writeLp("testABC1")
        await connBCA.writeLp("testBCA1")
        await connCAB.writeLp("testCAB1")
        check:
          "testABC2" == string.fromBytes(await connABC.readLp(1024))
          "testBCA2" == string.fromBytes(await connBCA.readLp(1024))
          "testCAB2" == string.fromBytes(await connCAB.readLp(1024))
        await connABC.writeLp("testABC3")
        await connBCA.writeLp("testBCA3")
        await connCAB.writeLp("testCAB3")
        check:
          "testABC4" == string.fromBytes(await connABC.readLp(1024))
          "testBCA4" == string.fromBytes(await connBCA.readLp(1024))
          "testCAB4" == string.fromBytes(await connCAB.readLp(1024))
        await allFutures(connABC.close(), connBCA.close(), connCAB.close())
        await allFutures(switchA.stop(), switchB.stop(), switchC.stop())