logos-messaging-nim/tests/waku_peer_exchange/test_protocol.nim

{.used.}

import
  std/[options, sequtils, tables],
  testutils/unittests,
  chronos,
  chronicles,
  stew/shims/net,
  libp2p/[switch, peerId, crypto/crypto, multistream, muxers/muxer],
  eth/[keys, p2p/discoveryv5/enr]

import
  ../../../waku/[
    waku_node,
    node/peer_manager,
    waku_discv5,
    waku_peer_exchange,
    waku_peer_exchange/rpc,
    waku_peer_exchange/rpc_codec,
    waku_peer_exchange/protocol,
    node/peer_manager,
    waku_relay/protocol,
    waku_relay,
    waku_core,
    waku_core/message/codec
  ],
  ../testlib/[wakucore, wakunode, simple_mock, assertions],
  ./utils.nim

suite "Waku Peer Exchange":
  # Some of this tests use node.wakuPeerExchange instead of just a standalone PeerExchange.
  # This is because attempts to connect the switches for two standalones PeerExchanges failed.
  # TODO: Try to make the tests work with standalone PeerExchanges

  suite "request":
    asyncTest "Retrieve and provide peer exchange peers from discv5":
      ## Given (copied from test_waku_discv5.nim)
      let
        # todo: px flag
        flags =
          CapabilitiesBitfield.init(
            lightpush = false, filter = false, store = false, relay = true
          )
        bindIp = parseIpAddress("0.0.0.0")
        extIp = parseIpAddress("127.0.0.1")

        nodeKey1 = generateSecp256k1Key()
        nodeTcpPort1 = Port(64010)
        nodeUdpPort1 = Port(9000)
        node1 =
          newTestWakuNode(
            nodeKey1,
            bindIp,
            nodeTcpPort1,
            some(extIp),
            wakuFlags = some(flags),
            discv5UdpPort = some(nodeUdpPort1),
          )

        nodeKey2 = generateSecp256k1Key()
        nodeTcpPort2 = Port(64012)
        nodeUdpPort2 = Port(9002)
        node2 =
          newTestWakuNode(
            nodeKey2,
            bindIp,
            nodeTcpPort2,
            some(extIp),
            wakuFlags = some(flags),
            discv5UdpPort = some(nodeUdpPort2),
          )

        nodeKey3 = generateSecp256k1Key()
        nodeTcpPort3 = Port(64014)
        nodeUdpPort3 = Port(9004)
        node3 =
          newTestWakuNode(
            nodeKey3,
            bindIp,
            nodeTcpPort3,
            some(extIp),
            wakuFlags = some(flags),
            discv5UdpPort = some(nodeUdpPort3),
          )

      # discv5
      let
        conf1 =
          WakuDiscoveryV5Config(
            discv5Config: none(DiscoveryConfig),
            address: bindIp,
            port: nodeUdpPort1,
            privateKey: keys.PrivateKey(nodeKey1.skkey),
            bootstrapRecords: @[],
            autoupdateRecord: true,
          )

      let
        disc1 =
          WakuDiscoveryV5.new(
            node1.rng, conf1, some(node1.enr), some(node1.peerManager)
          )

      let
        conf2 =
          WakuDiscoveryV5Config(
            discv5Config: none(DiscoveryConfig),
            address: bindIp,
            port: nodeUdpPort2,
            privateKey: keys.PrivateKey(nodeKey2.skkey),
            bootstrapRecords: @[disc1.protocol.getRecord()],
            autoupdateRecord: true,
          )

      let
        disc2 =
          WakuDiscoveryV5.new(
            node2.rng, conf2, some(node2.enr), some(node2.peerManager)
          )

      await allFutures(node1.start(), node2.start(), node3.start())
      let resultDisc1StartRes = await disc1.start()
      assert resultDisc1StartRes.isOk(), resultDisc1StartRes.error
      let resultDisc2StartRes = await disc2.start()
      assert resultDisc2StartRes.isOk(), resultDisc2StartRes.error

      ## When
      var attempts = 10
      while (disc1.protocol.nodesDiscovered < 1 or disc2.protocol.nodesDiscovered < 1) and
          attempts > 0:
        await sleepAsync(1.seconds)
        attempts -= 1

      # node2 can be connected, so will be returned by peer exchange
      require (
        await node1.peerManager.connectRelay(node2.switch.peerInfo.toRemotePeerInfo())
      )

      # Mount peer exchange
      await node1.mountPeerExchange()
      await node3.mountPeerExchange()

      let
        dialResponse =
          await node3.dialForPeerExchange(node1.switch.peerInfo.toRemotePeerInfo())
      let response = dialResponse.get()

      ## Then
      check:
        response.get().peerInfos.len == 1
        response.get().peerInfos[0].enr == disc2.protocol.localNode.record.raw

      await allFutures(
        [node1.stop(), node2.stop(), node3.stop(), disc1.stop(), disc2.stop()]
      )

    asyncTest "Request returns some discovered peers":
      let
        node1 =
          newTestWakuNode(generateSecp256k1Key(), parseIpAddress("0.0.0.0"), Port(0))
        node2 =
          newTestWakuNode(generateSecp256k1Key(), parseIpAddress("0.0.0.0"), Port(0))

      # Start and mount peer exchange
      await allFutures([node1.start(), node2.start()])
      await allFutures([node1.mountPeerExchange(), node2.mountPeerExchange()])

      # Create connection
      let
        connOpt =
          await node2.peerManager.dialPeer(
            node1.switch.peerInfo.toRemotePeerInfo(), WakuPeerExchangeCodec
          )
      require:
        connOpt.isSome

      # Create some enr and add to peer exchange (simulating disv5)
      var enr1, enr2 = enr.Record()
      check enr1.fromUri(
        "enr:-Iu4QGNuTvNRulF3A4Kb9YHiIXLr0z_CpvWkWjWKU-o95zUPR_In02AWek4nsSk7G_-YDcaT4bDRPzt5JIWvFqkXSNcBgmlkgnY0gmlwhE0WsGeJc2VjcDI1NmsxoQKp9VzU2FAh7fwOwSpg1M_Ekz4zzl0Fpbg6po2ZwgVwQYN0Y3CC6mCFd2FrdTIB"
      )
      check enr2.fromUri(
        "enr:-Iu4QGJllOWlviPIh_SGR-VVm55nhnBIU5L-s3ran7ARz_4oDdtJPtUs3Bc5aqZHCiPQX6qzNYF2ARHER0JPX97TFbEBgmlkgnY0gmlwhE0WsGeJc2VjcDI1NmsxoQP3ULycvday4EkvtVu0VqbBdmOkbfVLJx8fPe0lE_dRkIN0Y3CC6mCFd2FrdTIB"
      )

      # Mock that we have discovered these enrs
      node1.wakuPeerExchange.enrCache.add(enr1)
      node1.wakuPeerExchange.enrCache.add(enr2)

      # Request 2 peer from px. Test all request variants
      let response1 = await node2.wakuPeerExchange.request(2)
      let
        response2 =
          await node2.wakuPeerExchange.request(2, node1.peerInfo.toRemotePeerInfo())
      let response3 = await node2.wakuPeerExchange.request(2, connOpt.get())

      # Check the response or dont even continue
      require:
        response1.isOk
        response2.isOk
        response3.isOk

      check:
        response1.get().peerInfos.len == 2
        response2.get().peerInfos.len == 2
        response3.get().peerInfos.len == 2

        # Since it can return duplicates test that at least one of the enrs is in the response
        response1.get().peerInfos.anyIt(it.enr == enr1.raw) or
          response1.get().peerInfos.anyIt(it.enr == enr2.raw)
        response2.get().peerInfos.anyIt(it.enr == enr1.raw) or
          response2.get().peerInfos.anyIt(it.enr == enr2.raw)
        response3.get().peerInfos.anyIt(it.enr == enr1.raw) or
          response3.get().peerInfos.anyIt(it.enr == enr2.raw)

    asyncTest "Request fails gracefully":
      let
        node1 =
          newTestWakuNode(generateSecp256k1Key(), parseIpAddress("0.0.0.0"), Port(0))
        node2 =
          newTestWakuNode(generateSecp256k1Key(), parseIpAddress("0.0.0.0"), Port(0))

      # Start and mount peer exchange
      await allFutures([node1.start(), node2.start()])
      await allFutures([node1.mountPeerExchange(), node2.mountPeerExchange()])

      # Create connection
      let
        connOpt =
          await node2.peerManager.dialPeer(
            node1.switch.peerInfo.toRemotePeerInfo(), WakuPeerExchangeCodec
          )
      require connOpt.isSome

      # Force closing the connection to simulate a failed peer
      await connOpt.get().close()

      # Request 2 peer from px
      let response = await node1.wakuPeerExchange.request(2, connOpt.get())

      # Check that it failed gracefully
      check:
        response.isErr

    asyncTest "Request 0 peers, with 0 peers in PeerExchange":
      # Given a disconnected PeerExchange
      let
        switch = newTestSwitch()
        peerManager = PeerManager.new(switch)
        peerExchange = WakuPeerExchange.new(peerManager)

      # When requesting 0 peers
      let response = await peerExchange.request(0)

      # Then the response should be an error
      check:
        response.isErr
        response.error == "peer_not_found_failure"

    asyncTest "Request 0 peers, with 1 peer in PeerExchange":
      # Given two valid nodes with PeerExchange
      let
        node1 =
          newTestWakuNode(generateSecp256k1Key(), parseIpAddress("0.0.0.0"), Port(0))
        node2 =
          newTestWakuNode(generateSecp256k1Key(), parseIpAddress("0.0.0.0"), Port(0))

      # Start and mount peer exchange
      await allFutures([node1.start(), node2.start()])
      await allFutures([node1.mountPeerExchange(), node2.mountPeerExchange()])

      # Connect the nodes
      let
        dialResponse =
          await node2.peerManager.dialPeer(
            node1.switch.peerInfo.toRemotePeerInfo(), WakuPeerExchangeCodec
          )
      assert dialResponse.isSome

      # Mock that we have discovered one enr
      var record = enr.Record()
      check record.fromUri(
        "enr:-Iu4QGNuTvNRulF3A4Kb9YHiIXLr0z_CpvWkWjWKU-o95zUPR_In02AWek4nsSk7G_-YDcaT4bDRPzt5JIWvFqkXSNcBgmlkgnY0gmlwhE0WsGeJc2VjcDI1NmsxoQKp9VzU2FAh7fwOwSpg1M_Ekz4zzl0Fpbg6po2ZwgVwQYN0Y3CC6mCFd2FrdTIB"
      )
      node1.wakuPeerExchange.enrCache.add(record)

      # When requesting 0 peers
      let response = await node1.wakuPeerExchange.request(0)

      # Then the response should be empty
      assertResultOk(response)
      check response.get().peerInfos.len == 0

    asyncTest "Request with invalid peer info":
      # Given two valid nodes with PeerExchange
      let
        node1 =
          newTestWakuNode(generateSecp256k1Key(), parseIpAddress("0.0.0.0"), Port(0))
        node2 =
          newTestWakuNode(generateSecp256k1Key(), parseIpAddress("0.0.0.0"), Port(0))

      # Start and mount peer exchange
      await allFutures([node1.start(), node2.start()])
      await allFutures([node1.mountPeerExchange(), node2.mountPeerExchange()])

      # Mock that we have discovered one enr
      var record = enr.Record()
      check record.fromUri(
        "enr:-Iu4QGNuTvNRulF3A4Kb9YHiIXLr0z_CpvWkWjWKU-o95zUPR_In02AWek4nsSk7G_-YDcaT4bDRPzt5JIWvFqkXSNcBgmlkgnY0gmlwhE0WsGeJc2VjcDI1NmsxoQKp9VzU2FAh7fwOwSpg1M_Ekz4zzl0Fpbg6po2ZwgVwQYN0Y3CC6mCFd2FrdTIB"
      )
      node1.wakuPeerExchange.enrCache.add(record)

      # When making any request with an invalid peer info
      var remotePeerInfo2 = node2.peerInfo.toRemotePeerInfo()
      remotePeerInfo2.peerId.data.add(255.byte)
      let response = await node1.wakuPeerExchange.request(1, remotePeerInfo2)

      # Then the response should be an error
      check:
        response.isErr
        response.error == "dial_failure"

    asyncTest "Connections are closed after response is sent":
      # Create 3 nodes
      let
        nodes =
          toSeq(0..<3).mapIt(
            newTestWakuNode(generateSecp256k1Key(), parseIpAddress("0.0.0.0"), Port(0))
          )

      await allFutures(nodes.mapIt(it.start()))
      await allFutures(nodes.mapIt(it.mountPeerExchange()))

      # Multiple nodes request to node 0
      for i in 1..<3:
        let
          resp =
            await nodes[i].wakuPeerExchange.request(
              2, nodes[0].switch.peerInfo.toRemotePeerInfo()
            )
        require resp.isOk

      # Wait for streams to be closed
      await sleepAsync(1.seconds)

      # Check that all streams are closed for px
      check:
        nodes[0].peerManager.getNumStreams(WakuPeerExchangeCodec) == (0, 0)
        nodes[1].peerManager.getNumStreams(WakuPeerExchangeCodec) == (0, 0)
        nodes[2].peerManager.getNumStreams(WakuPeerExchangeCodec) == (0, 0)

  suite "Protocol Handler":
    asyncTest "Works as expected":
      let
        node1 =
          newTestWakuNode(generateSecp256k1Key(), parseIpAddress("0.0.0.0"), Port(0))
        node2 =
          newTestWakuNode(generateSecp256k1Key(), parseIpAddress("0.0.0.0"), Port(0))

      # Start and mount peer exchange
      await allFutures([node1.start(), node2.start()])
      await allFutures([node1.mountPeerExchange(), node2.mountPeerExchange()])

      # Mock that we have discovered these enrs
      var enr1 = enr.Record()
      check enr1.fromUri(
        "enr:-Iu4QGNuTvNRulF3A4Kb9YHiIXLr0z_CpvWkWjWKU-o95zUPR_In02AWek4nsSk7G_-YDcaT4bDRPzt5JIWvFqkXSNcBgmlkgnY0gmlwhE0WsGeJc2VjcDI1NmsxoQKp9VzU2FAh7fwOwSpg1M_Ekz4zzl0Fpbg6po2ZwgVwQYN0Y3CC6mCFd2FrdTIB"
      )
      node1.wakuPeerExchange.enrCache.add(enr1)

      # Create connection
      let
        connOpt =
          await node2.peerManager.dialPeer(
            node1.switch.peerInfo.toRemotePeerInfo(), WakuPeerExchangeCodec
          )
      require connOpt.isSome
      let conn = connOpt.get()

      # Send bytes so that they directly hit the handler
      let rpc = PeerExchangeRpc(request: PeerExchangeRequest(numPeers: 1))

      var buffer: seq[byte]
      await conn.writeLP(rpc.encode().buffer)
      buffer = await conn.readLp(MaxRpcSize.int)

      # Decode the response
      let decodedBuff = PeerExchangeRpc.decode(buffer)
      require decodedBuff.isOk

      # Check we got back the enr we mocked
      check:
        decodedBuff.get().response.peerInfos.len == 1
        decodedBuff.get().response.peerInfos[0].enr == enr1.raw