Update ProtocolWrapper.

sogou · Aug 2, 2024 · 8ad625b · 8ad625b
1 parent 47eeced
commit 8ad625b
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 28 deletions.
diff --git a/src/factory/MySQLTaskImpl.cc b/src/factory/MySQLTaskImpl.cc
@@ -100,8 +100,8 @@ class ComplexMySQLTask : public WFComplexClientTask<MySQLRequest, MySQLResponse>
 		MySSLWrapper(ProtocolMessage *msg, SSL *ssl) :
 			SSLWrapper(msg, ssl)
 		{ }
-		ProtocolMessage *get_msg() const { return this->msg; }
-		virtual ~MySSLWrapper() { delete this->msg; }
+		ProtocolMessage *get_msg() const { return this->message; }
+		virtual ~MySSLWrapper() { delete this->message; }
 	};
 
 private:

diff --git a/src/protocol/ProtocolMessage.h b/src/protocol/ProtocolMessage.h
@@ -81,6 +81,16 @@ class ProtocolMessage : public CommMessageOut, public CommMessageIn
 			return this->CommMessageIn::feedback(buf, size);
 	}
 
+	virtual void renew()
+	{
+		if (this->wrapper)
+			return this->wrapper->renew();
+		else
+			return this->CommMessageIn::renew();
+	}
+
+	virtual ProtocolMessage *inner() { return this; }
+
 protected:
 	size_t size_limit;
 
@@ -99,23 +109,23 @@ class ProtocolMessage : public CommMessageOut, public CommMessageIn
 	virtual ~ProtocolMessage() { delete this->attachment; }
 
 public:
-	ProtocolMessage(ProtocolMessage&& msg)
+	ProtocolMessage(ProtocolMessage&& message)
 	{
-		this->size_limit = msg.size_limit;
-		msg.size_limit = (size_t)-1;
-		this->attachment = msg.attachment;
-		msg.attachment = NULL;
+		this->size_limit = message.size_limit;
+		message.size_limit = (size_t)-1;
+		this->attachment = message.attachment;
+		message.attachment = NULL;
 	}
 
-	ProtocolMessage& operator = (ProtocolMessage&& msg)
+	ProtocolMessage& operator = (ProtocolMessage&& message)
 	{
-		if (&msg != this)
+		if (&message != this)
 		{
-			this->size_limit = msg.size_limit;
-			msg.size_limit = (size_t)-1;
+			this->size_limit = message.size_limit;
+			message.size_limit = (size_t)-1;
 			delete this->attachment;
-			this->attachment = msg.attachment;
-			msg.attachment = NULL;
+			this->attachment = message.attachment;
+			message.attachment = NULL;
 		}
 
 		return *this;
@@ -129,41 +139,47 @@ class ProtocolWrapper : public ProtocolMessage
 protected:
 	virtual int encode(struct iovec vectors[], int max)
 	{
-		return this->msg->encode(vectors, max);
+		return this->message->encode(vectors, max);
 	}
 
 	virtual int append(const void *buf, size_t *size)
 	{
-		return this->msg->append(buf, size);
+		return this->message->append(buf, size);
+	}
+
+protected:
+	virtual ProtocolMessage *inner()
+	{
+		return this->message->inner();
 	}
 
 protected:
-	ProtocolMessage *msg;
+	ProtocolMessage *message;
 
 public:
-	ProtocolWrapper(ProtocolMessage *msg)
+	ProtocolWrapper(ProtocolMessage *message)
 	{
-		msg->wrapper = this;
-		this->msg = msg;
+		message->wrapper = this;
+		this->message = message;
 	}
 
 public:
 	ProtocolWrapper(ProtocolWrapper&& wrapper) :
 		ProtocolMessage(std::move(wrapper))
 	{
-		wrapper.msg->wrapper = this;
-		this->msg = wrapper.msg;
-		wrapper.msg = NULL;
+		wrapper.message->wrapper = this;
+		this->message = wrapper.message;
+		wrapper.message = NULL;
 	}
 
 	ProtocolWrapper& operator = (ProtocolWrapper&& wrapper)
 	{
 		if (&wrapper != this)
 		{
 			*(ProtocolMessage *)this = std::move(wrapper);
-			wrapper.msg->wrapper = this;
-			this->msg = wrapper.msg;
-			wrapper.msg = NULL;
+			wrapper.message->wrapper = this;
+			this->message = wrapper.message;
+			wrapper.message = NULL;
 		}
 
 		return *this;

diff --git a/src/protocol/SSLWrapper.h b/src/protocol/SSLWrapper.h
@@ -61,8 +61,8 @@ class SSLWrapper : public ProtocolWrapper
 	SSL *ssl;
 
 public:
-	SSLWrapper(ProtocolMessage *msg, SSL *ssl) :
-		ProtocolWrapper(msg)
+	SSLWrapper(ProtocolMessage *message, SSL *ssl) :
+		ProtocolWrapper(message)
 	{
 		this->ssl = ssl;
 	}
@@ -78,7 +78,8 @@ class ServerSSLWrapper : public SSLWrapper
 	virtual int append(const void *buf, size_t *size);
 
 public:
-	ServerSSLWrapper(ProtocolMessage *msg, SSL *ssl) : SSLWrapper(msg, ssl)
+	ServerSSLWrapper(ProtocolMessage *message, SSL *ssl) :
+		SSLWrapper(message, ssl)
 	{
 	}