diff --git a/src/Makefile.am b/src/Makefile.am index 5a5e3abcfa..69ae02f582 100644 --- a/src/Makefile.am +++ b/src/Makefile.am @@ -111,6 +111,7 @@ BITCOIN_CORE_H = \ net.h \ netaddress.h \ netbase.h \ + netmessagemaker.h \ noui.h \ policy/fees.h \ policy/policy.h \ diff --git a/src/net.h b/src/net.h index 7a32e21f4a..9ea056c6ee 100644 --- a/src/net.h +++ b/src/net.h @@ -101,6 +101,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 data; + std::string command; +}; + + class CConnman { public: diff --git a/src/netmessagemaker.h b/src/netmessagemaker.h new file mode 100644 index 0000000000..7167434a19 --- /dev/null +++ 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 + 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)... }; + return msg; + } + + template + CSerializedNetMsg Make(std::string sCommand, Args&&... args) + { + return Make(0, std::move(sCommand), std::forward(args)...); + } + +private: + const int nVersion; +}; + +#endif // BITCOIN_NETMESSAGEMAKER_H diff --git a/src/streams.h b/src/streams.h index c3e7c9e9e4..b508784238 100644 --- a/src/streams.h +++ b/src/streams.h @@ -69,6 +69,75 @@ OverrideStream WithOrVersion(S* s, int nVersionFlag) return OverrideStream(s, s->GetType(), s->GetVersion() | nVersionFlag); } +/* 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& 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 + CVectorWriter(int nTypeIn, int nVersionIn, std::vector& vchDataIn, size_t nPosIn, Args&&... args) : CVectorWriter(nTypeIn, nVersionIn, vchDataIn, nPosIn) + { + ::SerializeMany(*this, std::forward(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(pch), nOverwrite); + } + if (nOverwrite < nSize) { + vchData.insert(vchData.end(), reinterpret_cast(pch) + nOverwrite, reinterpret_cast(pch) + nSize); + } + nPos += nSize; + } + template + 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& vchData; + size_t nPos; +}; + /** Double ended buffer combining vector and stream-like interfaces. * * >> and << read and write unformatted data using the above serialization templates. diff --git a/src/test/streams_tests.cpp b/src/test/streams_tests.cpp index 34f501e867..8b715ce93e 100644 --- a/src/test/streams_tests.cpp +++ b/src/test/streams_tests.cpp @@ -15,6 +15,64 @@ 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 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{{1, 2}})); + CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 0, a, b); + BOOST_CHECK((vch == std::vector{{1, 2}})); + vch.clear(); + + CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 2, a, b); + BOOST_CHECK((vch == std::vector{{0, 0, 1, 2}})); + CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 2, a, b); + BOOST_CHECK((vch == std::vector{{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{{0, 0, 1, 2, 0}})); + CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 2, a, b); + BOOST_CHECK((vch == std::vector{{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{{0, 0, 0, 1, 2}})); + CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 3, a, b); + BOOST_CHECK((vch == std::vector{{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{{0, 0, 0, 0, 1, 2}})); + CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 4, a, b); + BOOST_CHECK((vch == std::vector{{0, 0, 0, 0, 1, 2}})); + vch.clear(); + + CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 0, FLATDATA(bytes)); + BOOST_CHECK((vch == std::vector{{3, 4, 5, 6}})); + CVectorWriter(SER_NETWORK, INIT_PROTO_VERSION, vch, 0, FLATDATA(bytes)); + BOOST_CHECK((vch == std::vector{{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{{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{{8, 8, 1, 3, 4, 5, 6, 2}})); + vch.clear(); +} + BOOST_AUTO_TEST_CASE(streams_serializedata_xor) { std::vector in;