From f8bd13f85ae5404adef23a52719d804a5c36b1e8 Mon Sep 17 00:00:00 2001
From: Vasil Dimov <vd@FreeBSD.org>
Date: Fri, 23 Apr 2021 09:43:43 +0200
Subject: [PATCH] net: add new method Sock::Accept() that wraps accept()

This will help to increase `Sock` usage and make more code mockable.
---
 src/test/fuzz/util.cpp | 16 ++++++++++++++++
 src/test/fuzz/util.h   |  2 ++
 src/test/util/net.h    | 18 ++++++++++++++++++
 src/util/sock.cpp      | 27 +++++++++++++++++++++++++++
 src/util/sock.h        |  9 +++++++++
 5 files changed, 72 insertions(+)

diff --git a/src/test/fuzz/util.cpp b/src/test/fuzz/util.cpp
index ae5f7a379e0..dae7920eff2 100644
--- a/src/test/fuzz/util.cpp
+++ b/src/test/fuzz/util.cpp
@@ -10,6 +10,8 @@
 #include <util/time.h>
 #include <version.h>
 
+#include <memory>
+
 FuzzedSock::FuzzedSock(FuzzedDataProvider& fuzzed_data_provider)
     : m_fuzzed_data_provider{fuzzed_data_provider}
 {
@@ -155,6 +157,20 @@ int FuzzedSock::Connect(const sockaddr*, socklen_t) const
     return 0;
 }
 
+std::unique_ptr<Sock> FuzzedSock::Accept(sockaddr* addr, socklen_t* addr_len) const
+{
+    constexpr std::array accept_errnos{
+        ECONNABORTED,
+        EINTR,
+        ENOMEM,
+    };
+    if (m_fuzzed_data_provider.ConsumeBool()) {
+        SetFuzzedErrNo(m_fuzzed_data_provider, accept_errnos);
+        return std::unique_ptr<FuzzedSock>();
+    }
+    return std::make_unique<FuzzedSock>(m_fuzzed_data_provider);
+}
+
 int FuzzedSock::GetSockOpt(int level, int opt_name, void* opt_val, socklen_t* opt_len) const
 {
     constexpr std::array getsockopt_errnos{
diff --git a/src/test/fuzz/util.h b/src/test/fuzz/util.h
index 40aaeac63f2..fb42dcd0fc4 100644
--- a/src/test/fuzz/util.h
+++ b/src/test/fuzz/util.h
@@ -410,6 +410,8 @@ public:
 
     int Connect(const sockaddr*, socklen_t) const override;
 
+    std::unique_ptr<Sock> Accept(sockaddr* addr, socklen_t* addr_len) const override;
+
     int GetSockOpt(int level, int opt_name, void* opt_val, socklen_t* opt_len) const override;
 
     bool Wait(std::chrono::milliseconds timeout, Event requested, Event* occurred = nullptr) const override;
diff --git a/src/test/util/net.h b/src/test/util/net.h
index 2de6e712a25..3ef18400767 100644
--- a/src/test/util/net.h
+++ b/src/test/util/net.h
@@ -13,6 +13,7 @@
 #include <array>
 #include <cassert>
 #include <cstring>
+#include <memory>
 #include <string>
 
 struct ConnmanTestMsg : public CConnman {
@@ -126,6 +127,23 @@ public:
 
     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);
diff --git a/src/util/sock.cpp b/src/util/sock.cpp
index 1a4d67a65ee..2029d70a374 100644
--- a/src/util/sock.cpp
+++ b/src/util/sock.cpp
@@ -10,6 +10,7 @@
 #include <util/system.h>
 #include <util/time.h>
 
+#include <memory>
 #include <stdexcept>
 #include <string>
 
@@ -73,6 +74,32 @@ int Sock::Connect(const sockaddr* addr, socklen_t addr_len) const
     return connect(m_socket, addr, addr_len);
 }
 
+std::unique_ptr<Sock> Sock::Accept(sockaddr* addr, socklen_t* addr_len) const
+{
+#ifdef WIN32
+    static constexpr auto ERR = INVALID_SOCKET;
+#else
+    static constexpr auto ERR = SOCKET_ERROR;
+#endif
+
+    std::unique_ptr<Sock> sock;
+
+    const auto socket = accept(m_socket, addr, addr_len);
+    if (socket != ERR) {
+        try {
+            sock = std::make_unique<Sock>(socket);
+        } catch (const std::exception&) {
+#ifdef WIN32
+            closesocket(socket);
+#else
+            close(socket);
+#endif
+        }
+    }
+
+    return sock;
+}
+
 int Sock::GetSockOpt(int level, int opt_name, void* opt_val, socklen_t* opt_len) const
 {
     return getsockopt(m_socket, level, opt_name, static_cast<char*>(opt_val), opt_len);
diff --git a/src/util/sock.h b/src/util/sock.h
index 59cc8c0b1df..75104828578 100644
--- a/src/util/sock.h
+++ b/src/util/sock.h
@@ -10,6 +10,7 @@
 #include <util/time.h>
 
 #include <chrono>
+#include <memory>
 #include <string>
 
 /**
@@ -96,6 +97,14 @@ public:
      */
     [[nodiscard]] virtual int Connect(const sockaddr* addr, socklen_t addr_len) const;
 
+    /**
+     * accept(2) wrapper. Equivalent to `std::make_unique<Sock>(accept(this->Get(), addr, addr_len))`.
+     * Code that uses this wrapper can be unit tested if this method is overridden by a mock Sock
+     * implementation.
+     * The returned unique_ptr is empty if `accept()` failed in which case errno will be set.
+     */
+    [[nodiscard]] virtual std::unique_ptr<Sock> Accept(sockaddr* addr, socklen_t* addr_len) const;
+
     /**
      * getsockopt(2) wrapper. Equivalent to
      * `getsockopt(this->Get(), level, opt_name, opt_val, opt_len)`. Code that uses this