net: add CVectorWriter and CNetMsgMaker

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.
This commit is contained in:
Cory Fields 2016-11-10 17:05:23 -05:00
parent b7695c2275
commit 2ec935dcaa
5 changed files with 178 additions and 0 deletions

View File

@ -111,6 +111,7 @@ BITCOIN_CORE_H = \
net.h \ net.h \
netaddress.h \ netaddress.h \
netbase.h \ netbase.h \
netmessagemaker.h \
noui.h \ noui.h \
policy/fees.h \ policy/fees.h \
policy/policy.h \ policy/policy.h \

View File

@ -101,6 +101,20 @@ class CTransaction;
class CNodeStats; class CNodeStats;
class CClientUIInterface; 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 class CConnman
{ {
public: public:

36
src/netmessagemaker.h Normal file
View File

@ -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

View File

@ -69,6 +69,75 @@ OverrideStream<S> WithOrVersion(S* s, int nVersionFlag)
return OverrideStream<S>(s, s->GetType(), s->GetVersion() | nVersionFlag); return OverrideStream<S>(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<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;
};
/** Double ended buffer combining vector and stream-like interfaces. /** Double ended buffer combining vector and stream-like interfaces.
* *
* >> and << read and write unformatted data using the above serialization templates. * >> and << read and write unformatted data using the above serialization templates.

View File

@ -15,6 +15,64 @@ using namespace boost::assign; // bring 'operator+=()' into scope
BOOST_FIXTURE_TEST_SUITE(streams_tests, BasicTestingSetup) 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_CASE(streams_serializedata_xor) BOOST_AUTO_TEST_CASE(streams_serializedata_xor)
{ {
std::vector<char> in; std::vector<char> in;