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

Code before:
        assert (self.addrv2_test)
        self.handle_addr_message(message)

    def on_addr(self, message):
        assert (not self.addrv2_test)
        self.handle_addr_message(message)


class AddrSelfAnnouncementTest(BitcoinTestFramework):
    def set_test_params(self):
        self.num_nodes = 1
        self.extra_args = [[f"-externalip={IP_TO_ANNOUNCE}"]]

    def run_test(self):
        # populate addrman to have some addresses for a GETADDR response
        for i in range(50):
            a = f"{1 + i}.{i}.1.1"
            self.nodes[0].addpeeraddress(a, 8333)

        self.self_announcement_test(outbound=False, addrv2=False, assert_on_connection_open=self.inbound_connection_open_assertions)
        self.self_announcement_test(outbound=False, addrv2=True, assert_on_connection_open=self.inbound_connection_open_assertions)
        self.self_announcement_test(outbound=True, addrv2=False, assert_on_connection_open=self.outbound_connection_open_assertions)
        self.self_announcement_test(outbound=True, addrv2=True, assert_on_connection_open=self.outbound_connection_open_assertions)

    def inbound_connection_open_assertions(self, addr_receiver):
        # We expect one self-announcement and multiple other addresses in
        # response to a GETADDR in a single addr / addrv2 message.
        assert_equal(addr_receiver.self_announcements_received, 1)
        assert_equal(addr_receiver.addr_messages_received, 1)
        assert_greater_than(addr_receiver.addresses_received, 1)

    def outbound_connection_open_assertions(self, addr_receiver):
        # We expect only the self-announcement.
        assert_equal(addr_receiver.self_announcements_received, 1)
        assert_equal(addr_receiver.addr_messages_received, 1)
        assert_equal(addr_receiver.addresses_received, 1)

    def self_announcement_test(self, outbound, addrv2, assert_on_connection_open):
        connection_type = "outbound" if outbound else "inbound"
        addr_version = "addrv2" if addrv2 else "addrv1"
        self.log.info(f"Test that the node does an address self-announcement to {connection_type} connections ({addr_version})")

        # We only self-announce after initial block download is done
        assert (not self.nodes[0].getblockchaininfo()["initialblockdownload"])

        netinfo = self.nodes[0].getnetworkinfo()
        port = netinfo["localaddresses"][0]["port"]
        self.nodes[0].setmocktime(int(time.time()))

        expected = CAddress()
        expected.nServices = int(netinfo["localservices"], 16)

Review comment (bug):
note: I first thought that there might be a race condition (node receives verack first, and getaddr with some delay, sends out self-advertisment in between) but this isn't actually true because in this case we won't self-advertise due to #21528 (`SetupAddressRelay` not called yet) . So this is fine for now (whether the node behavior of mixing these is good is another issue)

Code after:
        assert (not self.addrv2_test)
        self.handle_addr_message(message)


class AddrSelfAnnouncementTest(BitcoinTestFramework):
    def set_test_params(self):
        self.num_nodes = 1
        self.extra_args = [[f"-externalip={IP_TO_ANNOUNCE}"]]

    def run_test(self):
        # populate addrman to have some addresses for a GETADDR response
        for i in range(50):
            a = f"{1 + i}.{i}.1.1"
            self.nodes[0].addpeeraddress(a, 8333)

        self.self_announcement_test(outbound=False, addrv2=False)
        self.self_announcement_test(outbound=False, addrv2=True)
        self.self_announcement_test(outbound=True, addrv2=False)
        self.self_announcement_test(outbound=True, addrv2=True)

    def inbound_connection_open_assertions(self, addr_receiver):
        # We expect one self-announcement and multiple other addresses in
        # response to a GETADDR in a single addr / addrv2 message.
        assert_equal(addr_receiver.self_announcements_received, 1)
        assert_equal(addr_receiver.addr_messages_received, 1)
        assert_greater_than(addr_receiver.addresses_received, 1)

    def outbound_connection_open_assertions(self, addr_receiver):
        # We expect only the self-announcement.
        assert_equal(addr_receiver.self_announcements_received, 1)
        assert_equal(addr_receiver.addr_messages_received, 1)
        assert_equal(addr_receiver.addresses_received, 1)

    def self_announcement_test(self, *, outbound, addrv2):
        connection_type = "outbound" if outbound else "inbound"
        addr_version = "addrv2" if addrv2 else "addrv1"
        self.log.info(f"Test that the node does an address self-announcement to {connection_type} connections ({addr_version})")

        # We only self-announce after initial block download is done
        assert (not self.nodes[0].getblockchaininfo()["initialblockdownload"])

        netinfo = self.nodes[0].getnetworkinfo()
        port = netinfo["localaddresses"][0]["port"]
        self.nodes[0].setmocktime(int(time.time()))

        expected = CAddress()
        expected.nServices = int(netinfo["localservices"], 16)
        expected.ip = IP_TO_ANNOUNCE
        expected.port = port
        expected.time = self.nodes[0].mocktime