diff --git a/src/net.cpp b/src/net.cpp index ee9d71786..210258c25 100644 --- a/src/net.cpp +++ b/src/net.cpp @@ -571,18 +571,13 @@ bool CNode::ReceiveMsgBytes(const char *pch, unsigned int nBytes, bool& complete nRecvBytes += nBytes; while (nBytes > 0) { // absorb network data - int handled; - if (!m_deserializer->in_data) - handled = m_deserializer->readHeader(pch, nBytes); - else - handled = m_deserializer->readData(pch, nBytes); - + int handled = m_deserializer->Read(pch, nBytes); if (handled < 0) { m_deserializer->Reset(); return false; } - if (m_deserializer->in_data && m_deserializer->hdr.nMessageSize > MAX_PROTOCOL_MESSAGE_LENGTH) { + if (m_deserializer->OversizedMessageDetected()) { LogPrint(BCLog::NET, "Oversized message from peer=%i, disconnecting\n", GetId()); m_deserializer->Reset(); return false; @@ -591,13 +586,13 @@ bool CNode::ReceiveMsgBytes(const char *pch, unsigned int nBytes, bool& complete pch += handled; nBytes -= handled; - if (m_deserializer->complete()) { + if (m_deserializer->Complete()) { // decompose a transport agnostic CNetMessage from the deserializer CNetMessage msg = m_deserializer->GetMessage(Params().MessageStart(), nTimeMicros); //store received bytes per message command //to prevent a memory DOS, only allow valid commands - mapMsgCmdSize::iterator i = mapRecvBytesPerMsgCmd.find(m_deserializer->hdr.pchCommand); + mapMsgCmdSize::iterator i = mapRecvBytesPerMsgCmd.find(msg.m_command); if (i == mapRecvBytesPerMsgCmd.end()) i = mapRecvBytesPerMsgCmd.find(NET_MESSAGE_COMMAND_OTHER); assert(i != mapRecvBytesPerMsgCmd.end()); @@ -639,7 +634,7 @@ int CNode::GetSendVersion() const return nSendVersion; } -int TransportDeserializer::readHeader(const char *pch, unsigned int nBytes) +int V1TransportDeserializer::readHeader(const char *pch, unsigned int nBytes) { // copy data to temporary parsing buffer unsigned int nRemaining = 24 - nHdrPos; @@ -670,7 +665,7 @@ int TransportDeserializer::readHeader(const char *pch, unsigned int nBytes) return nCopy; } -int TransportDeserializer::readData(const char *pch, unsigned int nBytes) +int V1TransportDeserializer::readData(const char *pch, unsigned int nBytes) { unsigned int nRemaining = hdr.nMessageSize - nDataPos; unsigned int nCopy = std::min(nRemaining, nBytes); @@ -687,15 +682,15 @@ int TransportDeserializer::readData(const char *pch, unsigned int nBytes) return nCopy; } -const uint256& TransportDeserializer::GetMessageHash() const +const uint256& V1TransportDeserializer::GetMessageHash() const { - assert(complete()); + assert(Complete()); if (data_hash.IsNull()) hasher.Finalize(data_hash.begin()); return data_hash; } -CNetMessage TransportDeserializer::GetMessage(const CMessageHeader::MessageStartChars& message_start, int64_t time) { +CNetMessage V1TransportDeserializer::GetMessage(const CMessageHeader::MessageStartChars& message_start, int64_t time) { // decompose a single CNetMessage from the TransportDeserializer CNetMessage msg(std::move(vRecv)); @@ -2708,7 +2703,7 @@ CNode::CNode(NodeId idIn, ServiceFlags nLocalServicesIn, int nMyStartingHeightIn LogPrint(BCLog::NET, "Added connection peer=%d\n", id); } - m_deserializer = MakeUnique(TransportDeserializer(Params().MessageStart(), SER_NETWORK, INIT_PROTO_VERSION)); + m_deserializer = MakeUnique(V1TransportDeserializer(Params().MessageStart(), SER_NETWORK, INIT_PROTO_VERSION)); } CNode::~CNode() diff --git a/src/net.h b/src/net.h index 15996438c..16435a3aa 100644 --- a/src/net.h +++ b/src/net.h @@ -637,20 +637,40 @@ public: * transport protocol agnostic CNetMessage (command & payload) */ class TransportDeserializer { +public: + // prepare for next message + virtual void Reset() = 0; + // returns true if the current deserialization is complete + virtual bool Complete() const = 0; + // checks if the potential message in deserialization is oversized + virtual bool OversizedMessageDetected() const = 0; + // set the serialization context version + virtual void SetVersion(int version) = 0; + // read and deserialize data + virtual int Read(const char *data, unsigned int bytes) = 0; + // decomposes a message from the context + virtual CNetMessage GetMessage(const CMessageHeader::MessageStartChars& message_start, int64_t time) = 0; + virtual ~TransportDeserializer() {} +}; + +class V1TransportDeserializer : public TransportDeserializer +{ private: mutable CHash256 hasher; mutable uint256 data_hash; -public: bool in_data; // parsing header (false) or data (true) - CDataStream hdrbuf; // partially received header CMessageHeader hdr; // complete header - unsigned int nHdrPos; - CDataStream vRecv; // received message data + unsigned int nHdrPos; unsigned int nDataPos; - TransportDeserializer(const CMessageHeader::MessageStartChars& pchMessageStartIn, int nTypeIn, int nVersionIn) : hdrbuf(nTypeIn, nVersionIn), hdr(pchMessageStartIn), vRecv(nTypeIn, nVersionIn) { + const uint256& GetMessageHash() const; + int readHeader(const char *pch, unsigned int nBytes); + int readData(const char *pch, unsigned int nBytes); +public: + + V1TransportDeserializer(const CMessageHeader::MessageStartChars& pchMessageStartIn, int nTypeIn, int nVersionIn) : hdrbuf(nTypeIn, nVersionIn), hdr(pchMessageStartIn), vRecv(nTypeIn, nVersionIn) { Reset(); } @@ -664,25 +684,23 @@ public: data_hash.SetNull(); hasher.Reset(); } - - bool complete() const + bool Complete() const { if (!in_data) return false; return (hdr.nMessageSize == nDataPos); } - - const uint256& GetMessageHash() const; - void SetVersion(int nVersionIn) { hdrbuf.SetVersion(nVersionIn); vRecv.SetVersion(nVersionIn); } - - int readHeader(const char *pch, unsigned int nBytes); - int readData(const char *pch, unsigned int nBytes); - + bool OversizedMessageDetected() const { + return (in_data && hdr.nMessageSize > MAX_PROTOCOL_MESSAGE_LENGTH); + } + int Read(const char *pch, unsigned int nBytes) { + return in_data ? readData(pch, nBytes) : readHeader(pch, nBytes); + } CNetMessage GetMessage(const CMessageHeader::MessageStartChars& message_start, int64_t time); };