diff --git a/src/protocol/PackageWrapper.cc b/src/protocol/PackageWrapper.cc index ce1152aa43..a8fae11813 100644 --- a/src/protocol/PackageWrapper.cc +++ b/src/protocol/PackageWrapper.cc @@ -39,7 +39,7 @@ int PackageWrapper::encode(struct iovec vectors[], int max) } cnt += ret; - this->message = this->next_out(this->message); + this->set_message(this->next_out(this->message)); if (!this->message) return cnt; @@ -57,7 +57,7 @@ int PackageWrapper::append(const void *buf, size_t *size) if (ret > 0) { - this->message = this->next_in(this->message); + this->set_message(this->next_in(this->message)); if (this->message) { this->renew(); diff --git a/src/protocol/ProtocolMessage.h b/src/protocol/ProtocolMessage.h index bb54e16d00..1f845ec580 100644 --- a/src/protocol/ProtocolMessage.h +++ b/src/protocol/ProtocolMessage.h @@ -114,7 +114,7 @@ class ProtocolMessage : public CommMessageOut, public CommMessageIn this->size_limit = message.size_limit; this->attachment = message.attachment; message.attachment = NULL; - this->wrapper = message.wrapper; + this->wrapper = NULL; } ProtocolMessage& operator = (ProtocolMessage&& message) @@ -125,7 +125,6 @@ class ProtocolMessage : public CommMessageOut, public CommMessageIn delete this->attachment; this->attachment = message.attachment; message.attachment = NULL; - this->wrapper = message.wrapper; } return *this; @@ -153,22 +152,28 @@ class ProtocolWrapper : public ProtocolMessage return this->message->inner(); } +protected: + void set_message(ProtocolMessage *message) + { + this->message = message; + if (message) + message->wrapper = this; + } + protected: ProtocolMessage *message; public: ProtocolWrapper(ProtocolMessage *message) { - message->wrapper = this; - this->message = message; + this->set_message(message); } public: ProtocolWrapper(ProtocolWrapper&& wrapper) : ProtocolMessage(std::move(wrapper)) { - wrapper.message->wrapper = this; - this->message = wrapper.message; + this->set_message(wrapper.message); wrapper.message = NULL; } @@ -177,8 +182,7 @@ class ProtocolWrapper : public ProtocolMessage if (&wrapper != this) { *(ProtocolMessage *)this = std::move(wrapper); - wrapper.message->wrapper = this; - this->message = wrapper.message; + this->set_message(wrapper.message); wrapper.message = NULL; }