net: add CVectorWriter and CNetMsgMaker

theuni · Fuzzbawls · commit b79e4168d857 · 2020-09-02T00:42:01.000-07:00
CVectorWriter is useful for overwriting or appending an existing byte vector.

CNetMsgMaker is a shortcut for creating messages on-the-fly which are suitable
for pushing to CConnman.
diff --git a/src/Makefile.am b/src/Makefile.am
@@ -212,6 +212,7 @@ BITCOIN_CORE_H = \
   net.h \
   netaddress.h \
   netbase.h \
+  netmessagemaker.h \
   noui.h \
   policy/fees.h \
   policy/policy.h \
diff --git a/src/Makefile.test.include b/src/Makefile.test.include
@@ -88,6 +88,7 @@ BITCOIN_TESTS =\
   test/sigopcount_tests.cpp \
   test/skiplist_tests.cpp \
   test/sync_tests.cpp \
+  test/streams_tests.cpp \
   test/timedata_tests.cpp \
   test/torcontrol_tests.cpp \
   test/transaction_tests.cpp \
diff --git a/src/net.h b/src/net.h
@@ -106,6 +106,20 @@ class CTransaction;
 class CNodeStats;
 class CClientUIInterface;
 
+struct CSerializedNetMsg
+{
+    CSerializedNetMsg() = default;
+    CSerializedNetMsg(CSerializedNetMsg&&) = default;
+    CSerializedNetMsg& operator=(CSerializedNetMsg&&) = default;
+    // No copying, only moves.
+    CSerializedNetMsg(const CSerializedNetMsg& msg) = delete;
+    CSerializedNetMsg& operator=(const CSerializedNetMsg&) = delete;
+
+    std::vector<unsigned char> data;
+    std::string command;
+};
+
+
 class CConnman
 {
 public:
diff --git a/src/netmessagemaker.h b/src/netmessagemaker.h
@@ -0,0 +1,36 @@
+// Copyright (c) 2009-2010 Satoshi Nakamoto
+// Copyright (c) 2009-2016 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_NETMESSAGEMAKER_H
+#define BITCOIN_NETMESSAGEMAKER_H
+
+#include "net.h"
+#include "serialize.h"
+
+class CNetMsgMaker
+{
+public:
+    CNetMsgMaker(int nVersionIn) : nVersion(nVersionIn){}
+
+    template <typename... Args>
+    CSerializedNetMsg Make(int nFlags, std::string sCommand, Args&&... args)
+    {
+        CSerializedNetMsg msg;
+        msg.command = std::move(sCommand);
+        CVectorWriter{ SER_NETWORK, nFlags | nVersion, msg.data, 0, std::forward<Args>(args)... };
+        return msg;
+    }
+
+    template <typename... Args>
+    CSerializedNetMsg Make(std::string sCommand, Args&&... args)
+    {
+        return Make(0, std::move(sCommand), std::forward<Args>(args)...);
+    }
+
+private:
+    const int nVersion;
+};
+
+#endif // BITCOIN_NETMESSAGEMAKER_H
diff --git a/src/streams.h b/src/streams.h
@@ -296,6 +296,75 @@ class CBaseDataStream
     }
 };
 
+/* Minimal stream for overwriting and/or appending to an existing byte vector
+ *
+ * The referenced vector will grow as necessary
+ */
+class CVectorWriter
+{
+ public:
+
+/*
+ * @param[in]  nTypeIn Serialization Type
+ * @param[in]  nVersionIn Serialization Version (including any flags)
+ * @param[in]  vchDataIn  Referenced byte vector to overwrite/append
+ * @param[in]  nPosIn Starting position. Vector index where writes should start. The vector will initially
+ *                    grow as necessary to  max(index, vec.size()). So to append, use vec.size().
+*/
+    CVectorWriter(int nTypeIn, int nVersionIn, std::vector<unsigned char>& vchDataIn, size_t nPosIn) : nType(nTypeIn), nVersion(nVersionIn), vchData(vchDataIn), nPos(nPosIn)
+    {
+        if(nPos > vchData.size())
+            vchData.resize(nPos);
+    }
+/*
+ * (other params same as above)
+ * @param[in]  args  A list of items to serialize starting at nPos.
+*/
+    template <typename... Args>
+    CVectorWriter(int nTypeIn, int nVersionIn, std::vector<unsigned char>& vchDataIn, size_t nPosIn, Args&&... args) : CVectorWriter(nTypeIn, nVersionIn, vchDataIn, nPosIn)
+    {
+        ::SerializeMany(*this, std::forward<Args>(args)...);
+    }
+    void write(const char* pch, size_t nSize)
+    {
+        assert(nPos <= vchData.size());
+        size_t nOverwrite = std::min(nSize, vchData.size() - nPos);
+        if (nOverwrite) {
+            memcpy(vchData.data() + nPos, reinterpret_cast<const unsigned char*>(pch), nOverwrite);
+        }
+        if (nOverwrite < nSize) {
+            vchData.insert(vchData.end(), reinterpret_cast<const unsigned char*>(pch) + nOverwrite, reinterpret_cast<const unsigned char*>(pch) + nSize);
+        }
+        nPos += nSize;
+    }
+    template<typename T>
+    CVectorWriter& operator<<(const T& obj)
+    {
+        // Serialize to this stream
+        ::Serialize(*this, obj);
+        return (*this);
+    }
+    int GetVersion() const
+    {
+        return nVersion;
+    }
+    int GetType() const
+    {
+        return nType;
+    }
+    void seek(size_t nSize)
+    {
+        nPos += nSize;
+        if(nPos > vchData.size())
+            vchData.resize(nPos);
+    }
+private:
+    const int nType;
+    const int nVersion;
+    std::vector<unsigned char>& vchData;
+    size_t nPos;
+};
+
 class CDataStream : public CBaseDataStream<CSerializeData>
 {
 public:
diff --git a/src/test/streams_tests.cpp b/src/test/streams_tests.cpp
@@ -0,0 +1,74 @@
+// Copyright (c) 2012-2015 The Bitcoin Core developers
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or http://www.opensource.org/licenses/mit-license.php.
+
+#include "streams.h"
+#include "test/test_pivx.h"
+
+#include <boost/assign/std/vector.hpp> // for 'operator+=()'
+#include <boost/assert.hpp>
+#include <boost/test/unit_test.hpp>
+
+using namespace boost::assign; // bring 'operator+=()' into scope
+
+BOOST_FIXTURE_TEST_SUITE(streams_tests, BasicTestingSetup)
+
+BOOST_AUTO_TEST_CASE(streams_vector_writer)
+{
+    unsigned char a(1);
+    unsigned char b(2);
+    unsigned char bytes[] = { 3, 4, 5, 6 };
+    std::vector<unsigned char> vch;
+
+    // Each test runs twice. Serializing a second time at the same starting
+    // point should yield the same results, even if the first test grew the
+    // vector.
+
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 0, a, b);
+    BOOST_CHECK((vch == std::vector<unsigned char>{{1, 2}}));
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 0, a, b);
+    BOOST_CHECK((vch == std::vector<unsigned char>{{1, 2}}));
+    vch.clear();
+
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 2, a, b);
+    BOOST_CHECK((vch == std::vector<unsigned char>{{0, 0, 1, 2}}));
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 2, a, b);
+    BOOST_CHECK((vch == std::vector<unsigned char>{{0, 0, 1, 2}}));
+    vch.clear();
+
+    vch.resize(5, 0);
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 2, a, b);
+    BOOST_CHECK((vch == std::vector<unsigned char>{{0, 0, 1, 2, 0}}));
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 2, a, b);
+    BOOST_CHECK((vch == std::vector<unsigned char>{{0, 0, 1, 2, 0}}));
+    vch.clear();
+
+    vch.resize(4, 0);
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 3, a, b);
+    BOOST_CHECK((vch == std::vector<unsigned char>{{0, 0, 0, 1, 2}}));
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 3, a, b);
+    BOOST_CHECK((vch == std::vector<unsigned char>{{0, 0, 0, 1, 2}}));
+    vch.clear();
+
+    vch.resize(4, 0);
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 4, a, b);
+    BOOST_CHECK((vch == std::vector<unsigned char>{{0, 0, 0, 0, 1, 2}}));
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 4, a, b);
+    BOOST_CHECK((vch == std::vector<unsigned char>{{0, 0, 0, 0, 1, 2}}));
+    vch.clear();
+
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 0, FLATDATA(bytes));
+    BOOST_CHECK((vch == std::vector<unsigned char>{{3, 4, 5, 6}}));
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 0, FLATDATA(bytes));
+    BOOST_CHECK((vch == std::vector<unsigned char>{{3, 4, 5, 6}}));
+    vch.clear();
+
+    vch.resize(4, 8);
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 2, a, FLATDATA(bytes), b);
+    BOOST_CHECK((vch == std::vector<unsigned char>{{8, 8, 1, 3, 4, 5, 6, 2}}));
+    CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 2, a, FLATDATA(bytes), b);
+    BOOST_CHECK((vch == std::vector<unsigned char>{{8, 8, 1, 3, 4, 5, 6, 2}}));
+    vch.clear();
+}
+
+BOOST_AUTO_TEST_SUITE_END()
