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

Code before:
    WRONG_GARBAGE = 3
    SEND_NO_AAD = 4
    SEND_NON_EMPTY_VERSION_PACKET = 5


class TestEncryptedP2PState(EncryptedP2PState):
    """ Modify v2 P2P protocol functions for testing scenarios listed in `TestType`"""
    def __init__(self, test_type):
        super().__init__(initiating=True, net='regtest')
        self.test_type = test_type
        if test_type == TestType.EARLY_KEY_RESPONSE:
            self.magic_sent = False  # set to True after first 4 bytes of ellswift which match network magic is sent.
            self.can_data_be_received = False  # variable used to assert if data is received on recvbuf.

    def generate_keypair_and_garbage(self):
        """Generate > MAX_GARBAGE_LEN garbage bytes, MAX_GARBAGE_LEN//2 garbage bytes, tampered garbage bytes
        when TestType = (EXCESS_GARBAGE, WRONG_GARBAGE_TERMINATOR, WRONG_GARBAGE)"""
        self.privkey_ours, self.ellswift_ours = ellswift_create()

        if self.test_type == TestType.EXCESS_GARBAGE:
            # send > 4095 bytes garbage
            garbage_len = MAX_GARBAGE_LEN + random.randrange(1, 10)
        elif self.test_type == TestType.WRONG_GARBAGE_TERMINATOR:
            garbage_len = random.randrange(MAX_GARBAGE_LEN//2)
        else:
            garbage_len = random.randrange(MAX_GARBAGE_LEN + 1)

        self.sent_garbage = random.randbytes(garbage_len)
        logger.debug(f"sending {garbage_len} bytes of garbage data")

        if self.test_type == TestType.WRONG_GARBAGE:
            # assume that garbage bytes sent to TestNode were tampered with
            return self.ellswift_ours + random_bitflip(self.sent_garbage)
        else:
            return self.ellswift_ours + self.sent_garbage

    def initiate_v2_handshake(self):
        """Send ellswift and garbage bytes in 2 parts when TestType = (EARLY_KEY_RESPONSE)"""
        if self.test_type == TestType.EARLY_KEY_RESPONSE:
            # Here, the 64 bytes ellswift is assumed to have it's first 4 bytes match network magic bytes.
            # It is sent in 2 phases:
            # 1. when `magic_sent` = False, send first 4 bytes of ellswift (matches network magic bytes)
            # 2. when `magic_sent` = True, send remaining 60 bytes of ellswift
            if not self.magic_sent:
                self.generate_keypair_and_garbage()
                self.magic_sent = True
                return b"\xfa\xbf\xb5\xda"
            else:
                # `can_data_be_received` is a variable used to assert if data is received on recvbuf.
                # 1. v2 TestNode shouldn't respond back if we send V1_PREFIX and data shouldn't be received on recvbuf.
                # This state is represented using `can_data_be_received` = False.

Review comment (suggestion):
doesn't the `+1` make the `garbage_len` greater than `4095` which is what is being covered in `TestType.EXCESS_GARBAGE`

Code after:
    SEND_NON_EMPTY_VERSION_PACKET = 5


class EarlyKeyResponseState(EncryptedP2PState):
    """ Modify v2 P2P protocol functions for testing EARLY_KEY_RESPONSE scenario"""
    def __init__(self, initiating, net):
        super().__init__(initiating=initiating, net=net)
        self.can_data_be_received = False  # variable used to assert if data is received on recvbuf.

    def initiate_v2_handshake(self):
        """Send ellswift and garbage bytes in 2 parts when TestType = (EARLY_KEY_RESPONSE)"""
        self.generate_keypair_and_garbage()
        return b""


class ExcessGarbageState(EncryptedP2PState):
    """Generate > MAX_GARBAGE_LEN garbage bytes"""
    def generate_keypair_and_garbage(self):
        garbage_len = MAX_GARBAGE_LEN + random.randrange(1, MAX_GARBAGE_LEN + 1)
        return super().generate_keypair_and_garbage(garbage_len)


class WrongGarbageTerminatorState(EncryptedP2PState):
    """Add option for sending wrong garbage terminator"""
    def generate_keypair_and_garbage(self):
        garbage_len = random.randrange(MAX_GARBAGE_LEN//2)
        return super().generate_keypair_and_garbage(garbage_len)

    def complete_handshake(self, response):
        length, handshake_bytes = super().complete_handshake(response)
        # first 16 bytes returned by complete_handshake() is the garbage terminator
        wrong_garbage_terminator = random_bitflip(handshake_bytes[:16])
        return length, wrong_garbage_terminator + handshake_bytes[16:]


class WrongGarbageState(EncryptedP2PState):
    """Generate tampered garbage bytes"""
    def generate_keypair_and_garbage(self):
        garbage_len = random.randrange(1, MAX_GARBAGE_LEN)
        ellswift_garbage_bytes = super().generate_keypair_and_garbage(garbage_len)
        # assume that garbage bytes sent to TestNode were tampered with
        return ellswift_garbage_bytes[:64] + random_bitflip(ellswift_garbage_bytes[64:])


class NoAADState(EncryptedP2PState):
    """Add option for not filling first encrypted packet after garbage terminator with AAD"""
    def generate_keypair_and_garbage(self):
        garbage_len = random.randrange(1, MAX_GARBAGE_LEN)
        return super().generate_keypair_and_garbage(garbage_len)

    def complete_handshake(self, response):