Repository: bitcoin/bitcoin
File: test/functional/p2p_v2_misbehaving.py (Python)

Code before:
            self.v2_state = ExcessGarbageState(initiating=True, net='regtest')
        elif self.test_type == TestType.WRONG_GARBAGE_TERMINATOR:
            self.v2_state = WrongGarbageTerminatorState(initiating=True, net='regtest')
        elif self.test_type == TestType.WRONG_GARBAGE:
            self.v2_state = WrongGarbageState(initiating=True, net='regtest')
        elif self.test_type == TestType.SEND_NO_AAD:
            self.v2_state = NoAADState(initiating=True, net='regtest')
        elif TestType.SEND_NON_EMPTY_VERSION_PACKET:
            self.v2_state = NonEmptyVersionPacketState(initiating=True, net='regtest')
        super().connection_made(transport)

    def data_received(self, t):
        if self.test_type == TestType.EARLY_KEY_RESPONSE:
            # check that data can be received on recvbuf only when mismatch from V1_PREFIX happens
            assert self.v2_state.can_data_be_received
        else:
            super().data_received(t)


class EncryptedP2PMisbehaving(BitcoinTestFramework):
    def set_test_params(self):
        self.num_nodes = 1
        self.extra_args = [["-v2transport=1", "-peertimeout=3"]]

    def run_test(self):
        self.test_earlykeyresponse()
        self.test_v2disconnection()

    def test_earlykeyresponse(self):
        self.log.info('Sending ellswift bytes in parts to ensure that response from responder is received only when')
        self.log.info('ellswift bytes have a mismatch from the 16 bytes(network magic followed by "version\\x00\\x00\\x00\\x00\\x00")')
        node0 = self.nodes[0]
        self.log.info('Sending first 4 bytes of ellswift which match network magic')
        self.log.info('If a response is received, assertion failure would happen in our custom data_received() function')
        peer1 = node0.add_p2p_connection(MisbehavingV2Peer(TestType.EARLY_KEY_RESPONSE), wait_for_verack=False, send_version=False, supports_v2_p2p=True, wait_for_v2_handshake=False)
        peer1.send_raw_message(MAGIC_BYTES['regtest'])
        self.log.info('Sending remaining ellswift and garbage which are different from V1_PREFIX. Since a response is')
        self.log.info('expected now, our custom data_received() function wouldn\'t result in assertion failure')
        peer1.v2_state.can_data_be_received = True
        peer1.send_raw_message(peer1.v2_state.ellswift_ours[4:] + peer1.v2_state.sent_garbage)
        with node0.assert_debug_log(['V2 handshake timeout peer=0']):
            peer1.wait_for_disconnect(timeout=5)
        self.log.info('successful disconnection since modified ellswift was sent as response')

    def test_v2disconnection(self):
        # test v2 disconnection scenarios
        node0 = self.nodes[0]
        expected_debug_message = [
            [],  # EARLY_KEY_RESPONSE
            ["V2 transport error: missing garbage terminator, peer=1"],  # EXCESS_GARBAGE
            ["V2 handshake timeout peer=2"],  # WRONG_GARBAGE_TERMINATOR

Review comment (refactor):
Nit: You can use the existing `bumpmocktime` helper?

Code after:
        elif self.test_type == TestType.EXCESS_GARBAGE:
            self.v2_state = ExcessGarbageState(initiating=True, net='regtest')
        elif self.test_type == TestType.WRONG_GARBAGE_TERMINATOR:
            self.v2_state = WrongGarbageTerminatorState(initiating=True, net='regtest')
        elif self.test_type == TestType.WRONG_GARBAGE:
            self.v2_state = WrongGarbageState(initiating=True, net='regtest')
        elif self.test_type == TestType.SEND_NO_AAD:
            self.v2_state = NoAADState(initiating=True, net='regtest')
        elif TestType.SEND_NON_EMPTY_VERSION_PACKET:
            self.v2_state = NonEmptyVersionPacketState(initiating=True, net='regtest')
        super().connection_made(transport)

    def data_received(self, t):
        if self.test_type == TestType.EARLY_KEY_RESPONSE:
            # check that data can be received on recvbuf only when mismatch from V1_PREFIX happens
            assert self.v2_state.can_data_be_received
        else:
            super().data_received(t)


class EncryptedP2PMisbehaving(BitcoinTestFramework):
    def set_test_params(self):
        self.num_nodes = 1
        self.extra_args = [["-v2transport=1", "-peertimeout=3"]]

    def run_test(self):
        self.test_earlykeyresponse()
        self.test_v2disconnection()

    def test_earlykeyresponse(self):
        self.log.info('Sending ellswift bytes in parts to ensure that response from responder is received only when')
        self.log.info('ellswift bytes have a mismatch from the 16 bytes(network magic followed by "version\\x00\\x00\\x00\\x00\\x00")')
        node0 = self.nodes[0]
        node0.setmocktime(int(time.time()))
        self.log.info('Sending first 4 bytes of ellswift which match network magic')
        self.log.info('If a response is received, assertion failure would happen in our custom data_received() function')
        peer1 = node0.add_p2p_connection(MisbehavingV2Peer(TestType.EARLY_KEY_RESPONSE), wait_for_verack=False, send_version=False, supports_v2_p2p=True, wait_for_v2_handshake=False)
        peer1.send_raw_message(MAGIC_BYTES['regtest'])
        self.log.info('Sending remaining ellswift and garbage which are different from V1_PREFIX. Since a response is')
        self.log.info('expected now, our custom data_received() function wouldn\'t result in assertion failure')
        peer1.v2_state.can_data_be_received = True
        self.wait_until(lambda: peer1.v2_state.ellswift_ours)
        peer1.send_raw_message(peer1.v2_state.ellswift_ours[4:] + peer1.v2_state.sent_garbage)
        node0.bumpmocktime(3)
        # Ensure that the bytes sent after 4 bytes network magic are actually received.
        self.wait_until(lambda: node0.getpeerinfo()[-1]["bytesrecv"] > 4)
        with node0.assert_debug_log(['V2 handshake timeout peer=0']):
            node0.bumpmocktime(1)  # `InactivityCheck()` triggers now
            peer1.wait_for_disconnect(timeout=1)
        self.log.info('successful disconnection since modified ellswift was sent as response')