#!/usr/bin/env python3
# Copyright (c) 2019-2021 The Bitcoin Core developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.

""" Test node eviction logic

When the number of peers has reached the limit of maximum connections,
the next connecting inbound peer will trigger the eviction mechanism.
We cannot currently test the parts of the eviction logic that are based on
address/netgroup since in the current framework, all peers are connecting from
the same local address. See Issue #14210 for more info.
Therefore, this test is limited to the remaining protection criteria.
"""
import time

from test_framework.blocktools import (
    create_block,
    create_coinbase,
)
from test_framework.messages import (
    msg_pong,
    msg_tx,
)
from test_framework.p2p import (
    P2PDataStore,
    P2PInterface,
)
from test_framework.test_framework import BitcoinTestFramework
from test_framework.util import assert_equal
from test_framework.wallet import MiniWallet


class SlowP2PDataStore(P2PDataStore):
    def on_ping(self, message):
        time.sleep(0.1)
        self.send_message(msg_pong(message.nonce))


class SlowP2PInterface(P2PInterface):
    def on_ping(self, message):
        time.sleep(0.1)
        self.send_message(msg_pong(message.nonce))


class P2PEvict(BitcoinTestFramework):
    def set_test_params(self):
        self.num_nodes = 1
        # The choice of maxconnections=32 results in a maximum of 21 inbound connections
        # (32 - 10 outbound - 1 feeler). 20 inbound peers are protected from eviction:
        # 4 by netgroup, 4 that sent us blocks, 4 that sent us transactions and 8 via lowest ping time
        self.extra_args = [['-maxconnections=32']]

    def run_test(self):
        protected_peers = set()  # peers that we expect to be protected from eviction
        current_peer = -1
        node = self.nodes[0]
        self.wallet = MiniWallet(node)

        self.log.info("Create 4 peers and protect them from eviction by sending us a block")
        for _ in range(4):
            block_peer = node.add_p2p_connection(SlowP2PDataStore())
            current_peer += 1
            block_peer.sync_with_ping()
            best_block = node.getbestblockhash()
            tip = int(best_block, 16)
            best_block_time = node.getblock(best_block)['time']
            block = create_block(tip, create_coinbase(node.getblockcount() + 1), best_block_time + 1)
            block.solve()
            block_peer.send_blocks_and_test([block], node, success=True)
            protected_peers.add(current_peer)

        self.log.info("Create 5 slow-pinging peers, making them eviction candidates")
        for _ in range(5):
            node.add_p2p_connection(SlowP2PInterface())
            current_peer += 1

        self.log.info("Create 4 peers and protect them from eviction by sending us a tx")
        for i in range(4):
            txpeer = node.add_p2p_connection(SlowP2PInterface())
            current_peer += 1
            txpeer.sync_with_ping()

            tx = self.wallet.create_self_transfer()['tx']
            txpeer.send_message(msg_tx(tx))
            protected_peers.add(current_peer)

        self.log.info("Create 8 peers and protect them from eviction by having faster pings")
        for _ in range(8):
            fastpeer = node.add_p2p_connection(P2PInterface())
            current_peer += 1
            self.wait_until(lambda: "ping" in fastpeer.last_message, timeout=10)

        # Make sure by asking the node what the actual min pings are
        peerinfo = node.getpeerinfo()
        pings = {}
        for i in range(len(peerinfo)):
            pings[i] = peerinfo[i]['minping'] if 'minping' in peerinfo[i] else 1000000
        sorted_pings = sorted(pings.items(), key=lambda x: x[1])

        # Usually the 8 fast peers are protected. In rare case of unreliable pings,
        # one of the slower peers might have a faster min ping though.
        for i in range(8):
            protected_peers.add(sorted_pings[i][0])

        self.log.info("Create peer that triggers the eviction mechanism")
        node.add_p2p_connection(SlowP2PInterface())

        # One of the non-protected peers must be evicted. We can't be sure which one because
        # 4 peers are protected via netgroup, which is identical for all peers,
        # and the eviction mechanism doesn't preserve the order of identical elements.
        evicted_peers = []
        for i in range(len(node.p2ps)):
            if not node.p2ps[i].is_connected:
                evicted_peers.append(i)

        self.log.info("Test that one peer was evicted")
        self.log.debug("{} evicted peer: {}".format(len(evicted_peers), set(evicted_peers)))
        assert_equal(len(evicted_peers), 1)

        self.log.info("Test that no peer expected to be protected was evicted")
        self.log.debug("{} protected peers: {}".format(len(protected_peers), protected_peers))
        assert evicted_peers[0] not in protected_peers

        self.log.info("Test that whitebind inbounds get extra eviction power")
        # Allow 10 inbound connections, set whitebind and forceinbound
        self.restart_node(0, extra_args=['-maxconnections=21', '-whitebind=127.0.0.1:30201', '-whitebind=forceinbound@127.0.0.1:30202'])
        self.log.debug("Fill connections with unprivileged peers")
        for i in range(10):
            node.add_p2p_connection(P2PInterface())

        # Create a peer that expects to be rejected
        # FIXME: "multiprocess, i686, DEBUG" CI task has a reliable timeout issue for v2transport
        class RejectedPeer(P2PInterface):
            def connection_lost(self, exc):
                return

        allowed_peers = []

        self.log.debug("Generic inbound gets rejected when full")
        with node.assert_debug_log(["failed to find an eviction candidate - connection dropped (full)"]):
            node.add_p2p_connection(RejectedPeer(), wait_for_verack=False, supports_v2_p2p=False)

        self.log.debug("Default whitebind inbound gets rejected, even when full")
        with node.assert_debug_log(["failed to find an eviction candidate - connection dropped (full)"]):
            node.add_p2p_connection(RejectedPeer(), wait_for_verack=False, supports_v2_p2p=False, dstport=30201)

        self.log.debug("ForceInbound whitebind inbound gets connected, even when full")
        allowed_peers.append(node.add_p2p_connection(P2PInterface(), dstport=30202))

        peerinfo = node.getpeerinfo()
        assert_equal(len(peerinfo), 10)
        for peer in peerinfo:
            if "30202" in peer["addrbind"]:
                assert peer["forced_inbound"]
            else:
                assert not peer["forced_inbound"]

        self.log.debug("Generic inbound gets rejected when whitebind peer is filling inbound slot")
        with node.assert_debug_log(["failed to find an eviction candidate - connection dropped (full)"]):
            node.add_p2p_connection(RejectedPeer(), supports_v2_p2p=False, wait_for_verack=False)

if __name__ == '__main__':
    P2PEvict(__file__).main()
