1 files changed, 122 insertions, 16 deletions

diff --git a/src/test/util/net.h b/src/test/util/net.h
index e25036be26..20c45058a1 100644
--- a/src/test/util/net.h
+++ b/src/test/util/net.h
@@ -1,26 +1,41 @@
-// Copyright (c) 2020 The Bitcoin Core developers
+// Copyright (c) 2020-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.
 
 #ifndef BITCOIN_TEST_UTIL_NET_H
 #define BITCOIN_TEST_UTIL_NET_H
 
+#include <compat.h>
+#include <netaddress.h>
 #include <net.h>
+#include <util/sock.h>
+
+#include <array>
+#include <cassert>
+#include <cstring>
+#include <memory>
+#include <string>
 
 struct ConnmanTestMsg : public CConnman {
     using CConnman::CConnman;
+
+    void SetPeerConnectTimeout(std::chrono::seconds timeout)
+    {
+        m_peer_connect_timeout = timeout;
+    }
+
     void AddTestNode(CNode& node)
     {
-        LOCK(cs_vNodes);
-        vNodes.push_back(&node);
+        LOCK(m_nodes_mutex);
+        m_nodes.push_back(&node);
     }
     void ClearTestNodes()
     {
-        LOCK(cs_vNodes);
-        for (CNode* node : vNodes) {
+        LOCK(m_nodes_mutex);
+        for (CNode* node : m_nodes) {
             delete node;
         }
-        vNodes.clear();
+        m_nodes.clear();
     }
 
     void ProcessMessagesOnce(CNode& node) { m_msgproc->ProcessMessages(&node, flagInterruptMsgProc); }
@@ -40,16 +55,16 @@ constexpr ServiceFlags ALL_SERVICE_FLAGS[]{
 };
 
 constexpr NetPermissionFlags ALL_NET_PERMISSION_FLAGS[]{
-    NetPermissionFlags::PF_NONE,
-    NetPermissionFlags::PF_BLOOMFILTER,
-    NetPermissionFlags::PF_RELAY,
-    NetPermissionFlags::PF_FORCERELAY,
-    NetPermissionFlags::PF_NOBAN,
-    NetPermissionFlags::PF_MEMPOOL,
-    NetPermissionFlags::PF_ADDR,
-    NetPermissionFlags::PF_DOWNLOAD,
-    NetPermissionFlags::PF_ISIMPLICIT,
-    NetPermissionFlags::PF_ALL,
+    NetPermissionFlags::None,
+    NetPermissionFlags::BloomFilter,
+    NetPermissionFlags::Relay,
+    NetPermissionFlags::ForceRelay,
+    NetPermissionFlags::NoBan,
+    NetPermissionFlags::Mempool,
+    NetPermissionFlags::Addr,
+    NetPermissionFlags::Download,
+    NetPermissionFlags::Implicit,
+    NetPermissionFlags::All,
 };
 
 constexpr ConnectionType ALL_CONNECTION_TYPES[]{
@@ -61,4 +76,95 @@ constexpr ConnectionType ALL_CONNECTION_TYPES[]{
     ConnectionType::ADDR_FETCH,
 };
 
+constexpr auto ALL_NETWORKS = std::array{
+    Network::NET_UNROUTABLE,
+    Network::NET_IPV4,
+    Network::NET_IPV6,
+    Network::NET_ONION,
+    Network::NET_I2P,
+    Network::NET_CJDNS,
+    Network::NET_INTERNAL,
+};
+
+/**
+ * A mocked Sock alternative that returns a statically contained data upon read and succeeds
+ * and ignores all writes. The data to be returned is given to the constructor and when it is
+ * exhausted an EOF is returned by further reads.
+ */
+class StaticContentsSock : public Sock
+{
+public:
+    explicit StaticContentsSock(const std::string& contents) : m_contents{contents}, m_consumed{0}
+    {
+        // Just a dummy number that is not INVALID_SOCKET.
+        m_socket = INVALID_SOCKET - 1;
+    }
+
+    ~StaticContentsSock() override { Reset(); }
+
+    StaticContentsSock& operator=(Sock&& other) override
+    {
+        assert(false && "Move of Sock into MockSock not allowed.");
+        return *this;
+    }
+
+    void Reset() override
+    {
+        m_socket = INVALID_SOCKET;
+    }
+
+    ssize_t Send(const void*, size_t len, int) const override { return len; }
+
+    ssize_t Recv(void* buf, size_t len, int flags) const override
+    {
+        const size_t consume_bytes{std::min(len, m_contents.size() - m_consumed)};
+        std::memcpy(buf, m_contents.data() + m_consumed, consume_bytes);
+        if ((flags & MSG_PEEK) == 0) {
+            m_consumed += consume_bytes;
+        }
+        return consume_bytes;
+    }
+
+    int Connect(const sockaddr*, socklen_t) const override { return 0; }
+
+    std::unique_ptr<Sock> Accept(sockaddr* addr, socklen_t* addr_len) const override
+    {
+        if (addr != nullptr) {
+            // Pretend all connections come from 5.5.5.5:6789
+            memset(addr, 0x00, *addr_len);
+            const socklen_t write_len = static_cast<socklen_t>(sizeof(sockaddr_in));
+            if (*addr_len >= write_len) {
+                *addr_len = write_len;
+                sockaddr_in* addr_in = reinterpret_cast<sockaddr_in*>(addr);
+                addr_in->sin_family = AF_INET;
+                memset(&addr_in->sin_addr, 0x05, sizeof(addr_in->sin_addr));
+                addr_in->sin_port = htons(6789);
+            }
+        }
+        return std::make_unique<StaticContentsSock>("");
+    };
+
+    int GetSockOpt(int level, int opt_name, void* opt_val, socklen_t* opt_len) const override
+    {
+        std::memset(opt_val, 0x0, *opt_len);
+        return 0;
+    }
+
+    bool Wait(std::chrono::milliseconds timeout,
+              Event requested,
+              Event* occurred = nullptr) const override
+    {
+        if (occurred != nullptr) {
+            *occurred = requested;
+        }
+        return true;
+    }
+
+private:
+    const std::string m_contents;
+    mutable size_t m_consumed;
+};
+
+std::vector<NodeEvictionCandidate> GetRandomNodeEvictionCandidates(int n_candidates, FastRandomContext& random_context);
+
 #endif // BITCOIN_TEST_UTIL_NET_H