From 1c52e0f370c4d676af442a7ea7f020d02fed80bc Mon Sep 17 00:00:00 2001
From: hiimgoodpack <hiimgoodpack@hiimgoodpack.hiimgoodpack>
Date: Wed, 7 Apr 2021 20:18:21 -0400
Subject: [PATCH] Add better escaping

---
 brainlet.h | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/brainlet.h b/brainlet.h
index ce33560..af670c3 100644
--- a/brainlet.h
+++ b/brainlet.h
@@ -73,7 +73,8 @@ namespace Brainlet {
 			State state = NOT_CONNECTED;
 
 			static void onReceive(Client* client, std::string buffer);
-			void encode(std::string& string); ///< Encodes a string to a JSON value
+			void escapeJSON(std::string& string); ///< Escapes a string to a JSON value
+			void escapeOutput(std::string& string); ///< Escapes a string to output safely
 
 		public:
 			std::function<void(Channel)> onNewChannel = nullptr;
@@ -112,7 +113,7 @@ namespace Brainlet {
 namespace Brainlet {
 	Client::Client(const std::string& domain) : domain(domain) {}
 
-	void Client::encode(std::string& string) {
+	void Client::escapeJSON(std::string& string) {
 		std::string result;
 		result.reserve(string.size());
 
@@ -131,6 +132,20 @@ namespace Brainlet {
 
 		string = std::move(result);
 	}
+	void Client::escapeOutput(std::string& string) {
+		std::string result;
+		result.reserve(string.size());
+
+		std::for_each(string.begin(), string.end(), [&](const char character) {
+			switch (character) {
+				case '\b':
+				case '\x1f': result.push_back('?'); break;
+				default: result.push_back(character); break;
+			}
+		});
+
+		string = std::move(result);
+	}
 	void Client::processNextEvent() {
 		socket->run_until(
 			// Stop in 250 microseconds
@@ -172,7 +187,9 @@ namespace Brainlet {
 		bool didError = value.get("error", true).asBool();
 		if (didError) {
 			// TODO: Use exception
-			std::cerr << "Error while logging in: " << value.get("message", "Server did not specify error").asString() << "\n";
+			std::string error = value.get("message", "Server did not specify error").asString();
+			escapeOutput(error);
+			std::cerr << "Error while logging in: " << error << "\n";
 			return;
 		}
 
@@ -210,6 +227,7 @@ namespace Brainlet {
 			case NOT_CONNECTED: throw std::logic_error("Buffer was received while client is not connected");
 			case WAIT_HELLO: {
 				if (opcode != HELLO) {
+					escapeOutput(buffer);
 					std::cerr << "Failed authentication: Expected HELLO packet, received: " << buffer << "\n";
 					client->socket->close();
 					return;
@@ -225,6 +243,7 @@ namespace Brainlet {
 			}
 			case WAIT_ACK: {
 				if (opcode != YOO_ACK) {
+					escapeOutput(buffer);
 					std::cerr << "Failed authentication: Expected YOO_ACK packet, received: " << buffer << "\n";
 					client->socket->close();
 					return;
@@ -243,6 +262,7 @@ namespace Brainlet {
 					if (client->onNewChannel)
 						client->onNewChannel(channelObject);
 				}
+				break;
 			}
 			case READY: {
 				switch (opcode) {
@@ -260,15 +280,17 @@ namespace Brainlet {
 						break;
 					}
 					default: {
+						escapeOutput(buffer);
 						std::cerr << "Unknown opcode " << opcode << ". Buffer " << buffer << "\n";
 						break;
 					}
 				}
+				break;
 			}
 		}
 	}
 	void Client::sendMessage(std::string message, std::string channelId) {
-		encode(message);
+		escapeJSON(message);
 
 		std::ostringstream bufferToSend;
 		bufferToSend << ACTION_CREATE_MESSAGE << "@" << R"({"content":")" << message << R"(","channel":{"_id":")" << channelId << R"("}})";