From 5730a43703f7e5a5ca26245ba3b55fbdd027d0b6 Mon Sep 17 00:00:00 2001
From: Martin Zumsande <mzumsande@gmail.com>
Date: Mon, 31 May 2021 22:52:08 +0200
Subject: [PATCH] test: Add functional test for AddrFetch connections

Co-authored-by: Amiti Uttarwar <amiti@uttarwar.org>
---
 test/functional/p2p_addrfetch.py | 62 ++++++++++++++++++++++++++++++++
 test/functional/test_runner.py   |  1 +
 2 files changed, 63 insertions(+)
 create mode 100755 test/functional/p2p_addrfetch.py

diff --git a/test/functional/p2p_addrfetch.py b/test/functional/p2p_addrfetch.py
new file mode 100755
index 0000000000..66ee1544a9
--- /dev/null
+++ b/test/functional/p2p_addrfetch.py
@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+# Copyright (c) 2021 The Bitcoin Core developers
+# Distributed under the MIT software license, see the accompanying
+# file COPYING or http://www.opensource.org/licenses/mit-license.php.
+"""
+Test p2p addr-fetch connections
+"""
+
+import time
+
+from test_framework.messages import msg_addr, CAddress, NODE_NETWORK, NODE_WITNESS
+from test_framework.p2p import P2PInterface, p2p_lock
+from test_framework.test_framework import BitcoinTestFramework
+from test_framework.util import assert_equal
+
+ADDR = CAddress()
+ADDR.time = int(time.time())
+ADDR.nServices = NODE_NETWORK | NODE_WITNESS
+ADDR.ip = "192.0.0.8"
+ADDR.port = 18444
+
+
+class P2PAddrFetch(BitcoinTestFramework):
+
+    def set_test_params(self):
+        self.setup_clean_chain = True
+        self.num_nodes = 1
+
+    def run_test(self):
+        node = self.nodes[0]
+        self.log.info("Connect to an addr-fetch peer")
+        peer = node.add_outbound_p2p_connection(P2PInterface(), p2p_idx=0, connection_type="addr-fetch")
+        info = node.getpeerinfo()
+        assert_equal(len(info), 1)
+        assert_equal(info[0]['connection_type'], 'addr-fetch')
+
+        self.log.info("Check that we send getaddr but don't try to sync headers with the addr-fetch peer")
+        peer.sync_send_with_ping()
+        with p2p_lock:
+            assert peer.message_count['getaddr'] == 1
+            assert peer.message_count['getheaders'] == 0
+
+        self.log.info("Check that answering the getaddr with a single address does not lead to disconnect")
+        # This prevents disconnecting on self-announcements
+        msg = msg_addr()
+        msg.addrs = [ADDR]
+        peer.send_and_ping(msg)
+        assert_equal(len(node.getpeerinfo()), 1)
+
+        self.log.info("Check that answering with larger addr messages leads to disconnect")
+        msg.addrs = [ADDR] * 2
+        peer.send_message(msg)
+        peer.wait_for_disconnect(timeout=5)
+
+        self.log.info("Check timeout for addr-fetch peer that does not send addrs")
+        peer = node.add_outbound_p2p_connection(P2PInterface(), p2p_idx=1, connection_type="addr-fetch")
+        node.setmocktime(int(time.time()) + 301)  # Timeout: 5 minutes
+        peer.wait_for_disconnect(timeout=5)
+
+
+if __name__ == '__main__':
+    P2PAddrFetch().main()
diff --git a/test/functional/test_runner.py b/test/functional/test_runner.py
index 49f269f8b4..4dcd92708c 100755
--- a/test/functional/test_runner.py
+++ b/test/functional/test_runner.py
@@ -182,6 +182,7 @@ BASE_SCRIPTS = [
     'p2p_addr_relay.py',
     'p2p_getaddr_caching.py',
     'p2p_getdata.py',
+    'p2p_addrfetch.py',
     'rpc_net.py',
     'wallet_keypool.py --legacy-wallet',
     'wallet_keypool.py --descriptors',