From 9ab2dcc83601bbcba9dbbb602c4671dc24828eb4 Mon Sep 17 00:00:00 2001
From: holmesyxli <holmesyxli@tencent.com>
Date: Sun, 7 Nov 2021 20:08:13 +0800
Subject: [PATCH] support random masking_key

---
 src/client/WFWebSocketClient.cc  | 12 +++++++++---
 src/client/WFWebSocketClient.h   |  2 +-
 src/factory/WFChannel.inl        | 21 ++++++++++++++++++---
 src/factory/WebSocketTaskImpl.cc |  2 ++
 4 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/src/client/WFWebSocketClient.cc b/src/client/WFWebSocketClient.cc
index 288a3158..ad431b96 100644
--- a/src/client/WFWebSocketClient.cc
+++ b/src/client/WFWebSocketClient.cc
@@ -23,9 +23,12 @@
 
 WFWebSocketTask *WebSocketClient::create_websocket_task(websocket_callback_t cb)
 {
-	return new ComplexWebSocketOutTask(this->channel,
-									   WFGlobal::get_scheduler(),
-									   std::move(cb));
+	WFWebSocketTask *task = new ComplexWebSocketOutTask(this->channel,
+														WFGlobal::get_scheduler(),
+														std::move(cb));
+
+	task->get_msg()->set_masking_key(this->channel->gen_masking_key());
+	return task;
 }
 
 int WebSocketClient::init(const std::string& url)
@@ -36,6 +39,7 @@ int WebSocketClient::init(const std::string& url)
 
 	this->channel = new ComplexWebSocketChannel(NULL,
 												WFGlobal::get_scheduler(),
+												this->params.random_masking_key,
 												std::move(process));
 	this->channel->set_idle_timeout(this->params.idle_timeout);
 	this->channel->set_size_limit(this->params.size_limit);
@@ -78,6 +82,7 @@ WFWebSocketTask *WebSocketClient::create_ping_task(websocket_callback_t cb)
 
 	protocol::WebSocketFrame *msg = ping_task->get_msg();
 	msg->set_opcode(WebSocketFramePing);
+	msg->set_masking_key(this->channel->gen_masking_key());
 
 	return ping_task;
 }
@@ -91,6 +96,7 @@ WFWebSocketTask *WebSocketClient::create_close_task(websocket_callback_t cb)
 
 	protocol::WebSocketFrame *msg = close_task->get_msg();
 	msg->set_opcode(WebSocketFrameConnectionClose);
+	msg->set_masking_key(this->channel->gen_masking_key());
 
 	return close_task;
 }
diff --git a/src/client/WFWebSocketClient.h b/src/client/WFWebSocketClient.h
index e0413572..9962003b 100644
--- a/src/client/WFWebSocketClient.h
+++ b/src/client/WFWebSocketClient.h
@@ -39,7 +39,7 @@ static constexpr struct WFWebSocketParams WEBSOCKET_PARAMS_DEFAULT =
 	.idle_timeout		=	WS_HANDSHAKE_TIMEOUT,
 	.ping_interval		=	-1,
 	.size_limit			=	(size_t)-1,
-	.random_masking_key	=	false,
+	.random_masking_key	=	true,
 };
 
 class WebSocketClient
diff --git a/src/factory/WFChannel.inl b/src/factory/WFChannel.inl
index eefdc05d..e4cbdbb4 100644
--- a/src/factory/WFChannel.inl
+++ b/src/factory/WFChannel.inl
@@ -17,6 +17,7 @@
           Xie Han (xiehan@sogou-inc.com)
 */
 
+#include <random>
 #include <functional>
 #include "TransRequest.h"
 #include "WFTask.h"
@@ -451,6 +452,12 @@ class ComplexWebSocketChannel : public WFComplexChannel<protocol::WebSocketFrame
 public:
 	void set_idle_timeout(int timeout) { this->idle_timeout = timeout; }
 	void set_size_limit(size_t size_limit) { this->size_limit = size_limit; }
+	uint32_t gen_masking_key()
+	{
+		if (random_masking_key == false)
+			return 0;
+		return this->gen();
+	}
 
 protected:
 	virtual CommMessageIn *message_in();
@@ -461,13 +468,20 @@ protected:
 private:
 	int idle_timeout;
 	size_t size_limit;
+	bool random_masking_key;
+	std::random_device rd;
+	std::mt19937 gen;
 
 public:
-	ComplexWebSocketChannel(CommSchedObject *object, CommScheduler *scheduler,
+	ComplexWebSocketChannel(CommSchedObject *object,
+							CommScheduler *scheduler,
+							bool random_masking_key,
 							websocket_process_t&& process) :
 		WFComplexChannel<protocol::WebSocketFrame>(object, scheduler,
-												   std::move(process))
+												   std::move(process)),
+		gen(rd())
 	{
+		this->random_masking_key = random_masking_key;
 	}
 };
 
@@ -478,7 +492,8 @@ protected:
 	virtual SubTask *done();
 
 public:
-	ComplexWebSocketOutTask(ComplexWebSocketChannel *channel, CommScheduler *scheduler,
+	ComplexWebSocketOutTask(ComplexWebSocketChannel *channel,
+							CommScheduler *scheduler,
 							websocket_callback_t&& cb) :
 		ComplexChannelOutTask<protocol::WebSocketFrame>(channel,
 														scheduler,
diff --git a/src/factory/WebSocketTaskImpl.cc b/src/factory/WebSocketTaskImpl.cc
index 72b4f231..dbb52930 100644
--- a/src/factory/WebSocketTaskImpl.cc
+++ b/src/factory/WebSocketTaskImpl.cc
@@ -80,6 +80,7 @@ SubTask *ComplexWebSocketInTask::done()
 		WebSocketFrame *msg = close_task->get_msg();
 		msg->set_opcode(WebSocketFrameConnectionClose);
 		msg->set_data(parser);
+		msg->set_masking_key(channel->gen_masking_key());
 		series->push_front(close_task);
 	}
 	else if (parser->opcode == WebSocketFramePing)
@@ -90,6 +91,7 @@ SubTask *ComplexWebSocketInTask::done()
 		WebSocketFrame *msg = pong_task->get_msg();
 		msg->set_opcode(WebSocketFramePong);
 		msg->set_data(parser);
+		msg->set_masking_key(channel->gen_masking_key());
 		series->push_front(pong_task);
 	}