1#!/usr/bin/env python3
2# Copyright (c) 2020 The Bitcoin Core developers
3# Distributed under the MIT software license, see the accompanying
4# file COPYING or http://www.opensource.org/licenses/mit-license.php.
5"""Test addr response caching"""
6
7import time
8
9from test_framework.messages import msg_getaddr
10from test_framework.p2p import (
11    P2PInterface,
12    p2p_lock
13)
14from test_framework.test_framework import BitcoinTestFramework
15from test_framework.util import (
16    assert_equal,
17)
18
19# As defined in net_processing.
20MAX_ADDR_TO_SEND = 1000
21MAX_PCT_ADDR_TO_SEND = 23
22
23class AddrReceiver(P2PInterface):
24
25    def __init__(self):
26        super().__init__()
27        self.received_addrs = None
28
29    def get_received_addrs(self):
30        with p2p_lock:
31            return self.received_addrs
32
33    def on_addr(self, message):
34        self.received_addrs = []
35        for addr in message.addrs:
36            self.received_addrs.append(addr.ip)
37
38    def addr_received(self):
39        return self.received_addrs is not None
40
41
42class AddrTest(BitcoinTestFramework):
43    def set_test_params(self):
44        self.setup_clean_chain = False
45        self.num_nodes = 1
46
47    def run_test(self):
48        self.log.info('Fill peer AddrMan with a lot of records')
49        for i in range(10000):
50            first_octet = i >> 8
51            second_octet = i % 256
52            a = "{}.{}.1.1".format(first_octet, second_octet)
53            self.nodes[0].addpeeraddress(a, 8333)
54
55        # Need to make sure we hit MAX_ADDR_TO_SEND records in the addr response later because
56        # only a fraction of all known addresses can be cached and returned.
57        assert(len(self.nodes[0].getnodeaddresses(0)) > int(MAX_ADDR_TO_SEND / (MAX_PCT_ADDR_TO_SEND / 100)))
58
59        responses = []
60        self.log.info('Send many addr requests within short time to receive same response')
61        N = 5
62        cur_mock_time = int(time.time())
63        for i in range(N):
64            addr_receiver = self.nodes[0].add_p2p_connection(AddrReceiver())
65            addr_receiver.send_and_ping(msg_getaddr())
66            # Trigger response
67            cur_mock_time += 5 * 60
68            self.nodes[0].setmocktime(cur_mock_time)
69            addr_receiver.wait_until(addr_receiver.addr_received)
70            responses.append(addr_receiver.get_received_addrs())
71        for response in responses[1:]:
72            assert_equal(response, responses[0])
73        assert(len(response) == MAX_ADDR_TO_SEND)
74
75        cur_mock_time += 3 * 24 * 60 * 60
76        self.nodes[0].setmocktime(cur_mock_time)
77
78        self.log.info('After time passed, see a new response to addr request')
79        last_addr_receiver = self.nodes[0].add_p2p_connection(AddrReceiver())
80        last_addr_receiver.send_and_ping(msg_getaddr())
81        # Trigger response
82        cur_mock_time += 5 * 60
83        self.nodes[0].setmocktime(cur_mock_time)
84        last_addr_receiver.wait_until(last_addr_receiver.addr_received)
85        # new response is different
86        assert(set(responses[0]) != set(last_addr_receiver.get_received_addrs()))
87
88
89if __name__ == '__main__':
90    AddrTest().main()
91