1#!/usr/bin/env python3
2# Copyright (c) 2017-2019 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 message sending before handshake completion.
6
7A node should never send anything other than VERSION/VERACK/REJECT until it's
8received a VERACK.
9
10This test connects to a node and sends it a few messages, trying to entice it
11into sending us something it shouldn't."""
12
13import time
14
15from test_framework.messages import msg_getaddr, msg_ping, msg_verack
16from test_framework.mininode import mininode_lock, P2PInterface
17from test_framework.test_framework import BitcoinTestFramework
18from test_framework.util import wait_until
19
20banscore = 10
21
22
23class CLazyNode(P2PInterface):
24    def __init__(self):
25        super().__init__()
26        self.unexpected_msg = False
27        self.ever_connected = False
28
29    def bad_message(self, message):
30        self.unexpected_msg = True
31        self.log.info("should not have received message: %s" % message.command)
32
33    def on_open(self):
34        self.ever_connected = True
35
36    def on_version(self, message): self.bad_message(message)
37    def on_verack(self, message): self.bad_message(message)
38    def on_reject(self, message): self.bad_message(message)
39    def on_inv(self, message): self.bad_message(message)
40    def on_addr(self, message): self.bad_message(message)
41    def on_getdata(self, message): self.bad_message(message)
42    def on_getblocks(self, message): self.bad_message(message)
43    def on_tx(self, message): self.bad_message(message)
44    def on_block(self, message): self.bad_message(message)
45    def on_getaddr(self, message): self.bad_message(message)
46    def on_headers(self, message): self.bad_message(message)
47    def on_getheaders(self, message): self.bad_message(message)
48    def on_ping(self, message): self.bad_message(message)
49    def on_mempool(self, message): self.bad_message(message)
50    def on_pong(self, message): self.bad_message(message)
51    def on_feefilter(self, message): self.bad_message(message)
52    def on_sendheaders(self, message): self.bad_message(message)
53    def on_sendcmpct(self, message): self.bad_message(message)
54    def on_cmpctblock(self, message): self.bad_message(message)
55    def on_getblocktxn(self, message): self.bad_message(message)
56    def on_blocktxn(self, message): self.bad_message(message)
57
58# Node that never sends a version. We'll use this to send a bunch of messages
59# anyway, and eventually get disconnected.
60class CNodeNoVersionBan(CLazyNode):
61    # send a bunch of veracks without sending a message. This should get us disconnected.
62    # NOTE: implementation-specific check here. Remove if bitcoind ban behavior changes
63    def on_open(self):
64        super().on_open()
65        for i in range(banscore):
66            self.send_message(msg_verack())
67
68    def on_reject(self, message): pass
69
70# Node that never sends a version. This one just sits idle and hopes to receive
71# any message (it shouldn't!)
72class CNodeNoVersionIdle(CLazyNode):
73    def __init__(self):
74        super().__init__()
75
76# Node that sends a version but not a verack.
77class CNodeNoVerackIdle(CLazyNode):
78    def __init__(self):
79        self.version_received = False
80        super().__init__()
81
82    def on_reject(self, message): pass
83    def on_verack(self, message): pass
84    # When version is received, don't reply with a verack. Instead, see if the
85    # node will give us a message that it shouldn't. This is not an exhaustive
86    # list!
87    def on_version(self, message):
88        self.version_received = True
89        self.send_message(msg_ping())
90        self.send_message(msg_getaddr())
91
92
93class P2PLeakTest(BitcoinTestFramework):
94    def set_test_params(self):
95        self.num_nodes = 1
96        self.extra_args = [['-banscore=' + str(banscore)]]
97
98    def run_test(self):
99        no_version_bannode = self.nodes[0].add_p2p_connection(CNodeNoVersionBan(), send_version=False, wait_for_verack=False)
100        no_version_idlenode = self.nodes[0].add_p2p_connection(CNodeNoVersionIdle(), send_version=False, wait_for_verack=False)
101        no_verack_idlenode = self.nodes[0].add_p2p_connection(CNodeNoVerackIdle(), wait_for_verack=False)
102
103        # Wait until we got the verack in response to the version. Though, don't wait for the other node to receive the
104        # verack, since we never sent one
105        no_verack_idlenode.wait_for_verack()
106
107        wait_until(lambda: no_version_bannode.ever_connected, timeout=10, lock=mininode_lock)
108        wait_until(lambda: no_version_idlenode.ever_connected, timeout=10, lock=mininode_lock)
109        wait_until(lambda: no_verack_idlenode.version_received, timeout=10, lock=mininode_lock)
110
111        # Mine a block and make sure that it's not sent to the connected nodes
112        self.nodes[0].generatetoaddress(1, self.nodes[0].get_deterministic_priv_key().address)
113
114        #Give the node enough time to possibly leak out a message
115        time.sleep(5)
116
117        #This node should have been banned
118        assert not no_version_bannode.is_connected
119
120        self.nodes[0].disconnect_p2ps()
121
122        # Wait until all connections are closed
123        wait_until(lambda: len(self.nodes[0].getpeerinfo()) == 0)
124
125        # Make sure no unexpected messages came in
126        assert no_version_bannode.unexpected_msg == False
127        assert no_version_idlenode.unexpected_msg == False
128        assert no_verack_idlenode.unexpected_msg == False
129
130
131if __name__ == '__main__':
132    P2PLeakTest().main()
133