From b32bf5f1f12f0e8097cdf4c7baa97efaac9195d2 Mon Sep 17 00:00:00 2001
From: Reece <me@reece.sx>
Date: Tue, 18 Apr 2023 02:09:16 +0100
Subject: [PATCH] [*] Logger mitigations

Closes #47
---
 Include/Aurora/Logging/ILogger.hpp |  2 ++
 Source/Logging/AuLogger.cpp        | 41 ++++++++++++++++++++++++------
 Source/Logging/AuLogger.hpp        |  7 ++++-
 3 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/Include/Aurora/Logging/ILogger.hpp b/Include/Aurora/Logging/ILogger.hpp
index 6e16462d..d5441ca0 100644
--- a/Include/Aurora/Logging/ILogger.hpp
+++ b/Include/Aurora/Logging/ILogger.hpp
@@ -33,6 +33,8 @@ namespace Aurora::Logging
         virtual void PushFilter(AuUInt8 level, bool shouldFilter) = 0;
         virtual void PopFilter() = 0;
 
+        virtual bool ExchangeMitigationState(bool bMitigationsEnabled) = 0;
+
     #if defined(_AUHAS_FMT)
         template<typename ... T>
         inline void WriteLinef(AuUInt8 level, const AuString &tag, const AuString &msg, T&& ... args)
diff --git a/Source/Logging/AuLogger.cpp b/Source/Logging/AuLogger.cpp
index 7304e32a..4cecebb6 100644
--- a/Source/Logging/AuLogger.cpp
+++ b/Source/Logging/AuLogger.cpp
@@ -37,13 +37,15 @@ namespace Aurora::Logging
     {
         try
         {
-            if (msg.line.find('\n') == AuString::npos) [[likely]]
+            if (msg.line.find('\n') == AuString::npos &&
+                ((!this->bEnableMitigations) || 
+                 (msg.line.find('\r') == AuString::npos &&
+                  msg.line.find('\33') == AuString::npos))) [[likely]]
             {
                 if (WriteNow(uLevel, msg))
                 {
-                    AddToPushQueue(uLevel, msg);
+                    AddToPushQueueConst(uLevel, msg);
                 }
-                
             }
             else [[unlikely]]
             {
@@ -83,7 +85,30 @@ namespace Aurora::Logging
         WriteLines(uLevel, msg);
     }
 
-    void Logger::AddToPushQueue(AuUInt8 uLevel, const ConsoleMessage &msg)
+    bool Logger::ExchangeMitigationState(bool bMitigationsEnabled)
+    {
+        return AuExchange(this->bEnableMitigations, bMitigationsEnabled);
+    }
+
+    void Logger::AddToPushQueue(AuUInt8 uLevel, /*const*/ ConsoleMessage &msg)
+    {
+        if (this->bEnableMitigations)
+        {
+            if (msg.line.find('\r') != AuString::npos) [[unlikely]]
+            {
+                AuReplaceAll(msg.line, "\r", "\\r");
+            }
+
+            if (msg.line.find('\33') != AuString::npos) [[unlikely]]
+            {
+                AuReplaceAll(msg.line, "\33", "\\e");
+            }
+        }
+
+        AddToPushQueueConst(uLevel, msg);
+    }
+
+    void Logger::AddToPushQueueConst(AuUInt8 uLevel, const ConsoleMessage &msg)
     {
         {
             AU_LOCK_GUARD(gTaskSpin);
@@ -91,14 +116,14 @@ namespace Aurora::Logging
             auto nice = gLogTasks.size() + 1;
             if (gLogTasks.capacity() < nice)
             {
-                gLogTasks.reserve(nice*10);
+                gLogTasks.reserve(nice * 10);
             }
 
-            while (!AuTryInsert(gLogTasks, AuMakeTuple(this, uLevel, msg)))
+            while (!AuTryInsert(gLogTasks, AuMakeTuple(this, uLevel, AuConstReference(msg))))
             {
                 SysPushErrorMem("Push failed - trying again");
                 spin.Unlock();
-                AuThreading::Sleep(100);
+                AuThreading::ContextYield();
                 spin.Lock();
             }
         }
@@ -115,7 +140,7 @@ namespace Aurora::Logging
             while (!AuTryInsert(filters, AuMakeTuple(uLevel, shouldFilter)))
             {
                 SysPushErrorMem("Push failed - trying again. wont be able to handle pop - wont syspanic yet");
-                AuThreading::Sleep(100);
+                AuThreading::ContextYield();
             }
 
             AuMemset(this->shouldFilter, 0, sizeof(this->shouldFilter));
diff --git a/Source/Logging/AuLogger.hpp b/Source/Logging/AuLogger.hpp
index 0f44da29..1480a377 100644
--- a/Source/Logging/AuLogger.hpp
+++ b/Source/Logging/AuLogger.hpp
@@ -21,7 +21,11 @@ namespace Aurora::Logging
         void PushFilter(AuUInt8 level, bool shouldFilter) override;
         void PopFilter() override;
 
-        void AddToPushQueue(AuUInt8 level, const ConsoleMessage &msg);
+        bool ExchangeMitigationState(bool bMitigationsEnabled) override;
+
+
+        void AddToPushQueueConst(AuUInt8 level, const ConsoleMessage &msg);
+        void AddToPushQueue(AuUInt8 level, ConsoleMessage &msg);
         bool WriteNow(AuUInt8 level, const ConsoleMessage &msg);
         void WriteLater(AuUInt8 level, const ConsoleMessage &msg);
         void WriteLines(AuUInt8 level, const ConsoleMessage &msg);
@@ -32,6 +36,7 @@ namespace Aurora::Logging
         AuList<AuSPtr<IBasicSink>> sinks;
         AuList<AuTuple<AuUInt8, bool>> filters; // std::vector > std::stack performance. dont ask me why or for the benchmarks, just trust me bro
         AuUInt8 shouldFilter[0xFF];
+        bool bEnableMitigations { true };
     };
 
     AuThreadPrimitives::SpinLock &GetLock();