From 8b094f8c6f992f5a044b005f2f1a7797db2bb33b Mon Sep 17 00:00:00 2001
From: "Unknown W. Brackets" <checkins@unknownbrackets.org>
Date: Wed, 11 Apr 2018 22:53:41 -0700
Subject: [PATCH] http: Allow fragmented data for WebSocket.

---
 ext/native/net/websocket_server.cpp | 44 +++++++++++++++++++++++++++++
 ext/native/net/websocket_server.h   |  8 +++++-
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/ext/native/net/websocket_server.cpp b/ext/native/net/websocket_server.cpp
index 05b5279387..e71a6050f0 100644
--- a/ext/native/net/websocket_server.cpp
+++ b/ext/native/net/websocket_server.cpp
@@ -96,6 +96,7 @@ WebSocketServer *WebSocketServer::CreateAsUpgrade(const http::Request &request)
 
 void WebSocketServer::Send(const std::string &str) {
 	assert(open_);
+	assert(fragmentOpcode_ == -1);
 	SendHeader(true, (int)Opcode::TEXT, str.size());
 	// TODO: For long strings, this will block.  Possibly not ideal.
 	if (!out_->Push(str.c_str(), str.size())) {
@@ -107,6 +108,7 @@ void WebSocketServer::Send(const std::string &str) {
 
 void WebSocketServer::Send(const std::vector<uint8_t> &payload) {
 	assert(open_);
+	assert(fragmentOpcode_ == -1);
 	SendHeader(true, (int)Opcode::BINARY, payload.size());
 	// TODO: For long strings, this will block.  Possibly not ideal.
 	if (!out_->Push((const char *)&payload[0], payload.size())) {
@@ -116,6 +118,48 @@ void WebSocketServer::Send(const std::vector<uint8_t> &payload) {
 	}
 }
 
+void WebSocketServer::AddFragment(bool finish, const std::string &str) {
+	assert(open_);
+	if (fragmentOpcode_ == -1) {
+		SendHeader(finish, (int)Opcode::TEXT, str.size());
+		fragmentOpcode_ = (int)Opcode::TEXT;
+	} else if (fragmentOpcode_ == (int)Opcode::TEXT) {
+		SendHeader(finish, (int)Opcode::CONTINUE, str.size());
+	} else {
+		assert(fragmentOpcode_ == (int)Opcode::TEXT || fragmentOpcode_ == -1);
+	}
+	// TODO: For long strings, this will block.  Possibly not ideal.
+	if (!out_->Push(str.c_str(), str.size())) {
+		assert(false);
+		open_ = false;
+		closeReason_ = WebSocketClose::ABNORMAL;
+	}
+	if (finish) {
+		fragmentOpcode_ = -1;
+	}
+}
+
+void WebSocketServer::AddFragment(bool finish, const std::vector<uint8_t> &payload) {
+	assert(open_);
+	if (fragmentOpcode_ == -1) {
+		SendHeader(finish, (int)Opcode::BINARY, payload.size());
+		fragmentOpcode_ = (int)Opcode::BINARY;
+	} else if (fragmentOpcode_ == (int)Opcode::BINARY) {
+		SendHeader(finish, (int)Opcode::CONTINUE, payload.size());
+	} else {
+		assert(fragmentOpcode_ == (int)Opcode::BINARY || fragmentOpcode_ == -1);
+	}
+	// TODO: For long strings, this will block.  Possibly not ideal.
+	if (!out_->Push((const char *)&payload[0], payload.size())) {
+		assert(false);
+		open_ = false;
+		closeReason_ = WebSocketClose::ABNORMAL;
+	}
+	if (finish) {
+		fragmentOpcode_ = -1;
+	}
+}
+
 void WebSocketServer::Ping(const std::vector<uint8_t> &payload) {
 	assert(open_);
 	assert(payload.size() <= 125);
diff --git a/ext/native/net/websocket_server.h b/ext/native/net/websocket_server.h
index 2aed3254f7..dae09d3d68 100644
--- a/ext/native/net/websocket_server.h
+++ b/ext/native/net/websocket_server.h
@@ -30,9 +30,14 @@ class WebSocketServer {
 public:
 	static WebSocketServer *CreateAsUpgrade(const http::Request &request);
 
-	// TODO: Doesn't support fragmented data yet.
 	void Send(const std::string &str);
 	void Send(const std::vector<uint8_t> &payload);
+
+	// Call with finish = false to start and continue, then finally with finish = true to complete.
+	// Note: Fragmented data cannot be interleaved, per protocol.
+	void AddFragment(bool finish, const std::string &str);
+	void AddFragment(bool finish, const std::vector<uint8_t> &payload);
+
 	void Ping(const std::vector<uint8_t> &payload = {});
 	void Pong(const std::vector<uint8_t> &payload = {});
 	void Close(WebSocketClose reason = WebSocketClose::GOING_AWAY);
@@ -73,6 +78,7 @@ protected:
 
 	bool open_ = true;
 	bool sentClose_ = false;
+	int fragmentOpcode_ = -1;
 	size_t fd_ = 0;
 	InputSink *in_ = nullptr;
 	OutputSink *out_ = nullptr;