From 5e8fdc997ada34e42160ef659831f9dc88b25c6c Mon Sep 17 00:00:00 2001
From: Bananymous <oskari.alaranta@bananymous.com>
Date: Sun, 17 Nov 2024 03:12:30 +0200
Subject: [PATCH] Kernel: Prevent userspace from overwriting CPU load memory

---
 kernel/include/kernel/Processor.h          |  1 +
 kernel/kernel/Device/FramebufferDevice.cpp | 71 +++++++++++++++++++---
 2 files changed, 64 insertions(+), 8 deletions(-)

diff --git a/kernel/include/kernel/Processor.h b/kernel/include/kernel/Processor.h
index 46a3a928..996f1350 100644
--- a/kernel/include/kernel/Processor.h
+++ b/kernel/include/kernel/Processor.h
@@ -60,6 +60,7 @@ namespace Kernel
 		static void wait_until_processors_ready();
 
 		static void toggle_should_print_cpu_load() { s_should_print_cpu_load = !s_should_print_cpu_load; }
+		static bool get_should_print_cpu_load() { return s_should_print_cpu_load; }
 
 		static ProcessorID bsb_id() { return s_bsb_id; }
 		static bool current_is_bsb() { return current_id() == bsb_id(); }
diff --git a/kernel/kernel/Device/FramebufferDevice.cpp b/kernel/kernel/Device/FramebufferDevice.cpp
index a5fac985..32a64076 100644
--- a/kernel/kernel/Device/FramebufferDevice.cpp
+++ b/kernel/kernel/Device/FramebufferDevice.cpp
@@ -3,11 +3,14 @@
 #include <kernel/Device/FramebufferDevice.h>
 #include <kernel/FS/DevFS/FileSystem.h>
 #include <kernel/Memory/Heap.h>
+#include <kernel/Terminal/TerminalDriver.h>
 
 #include <sys/framebuffer.h>
 #include <sys/mman.h>
 #include <sys/sysmacros.h>
 
+extern Kernel::TerminalDriver* g_terminal_driver;
+
 namespace Kernel
 {
 
@@ -250,16 +253,21 @@ namespace Kernel
 		{
 			if (flags != MS_SYNC)
 				return BAN::Error::from_errno(ENOTSUP);
+			if (vaddr % (BANAN_FB_BPP / 8))
+				return BAN::Error::from_errno(EINVAL);
 
-			if (vaddr < m_vaddr)
-				vaddr = m_vaddr;
-			if (vaddr + size > m_vaddr + m_size)
-				size = (vaddr - m_vaddr) + m_size;
+			if (auto rem = size % (BANAN_FB_BPP / 8))
+				size += (BANAN_FB_BPP / 8) - rem;
 
-			m_framebuffer->sync_pixels_linear(
-				(vaddr - m_vaddr) / (BANAN_FB_BPP / 8),
-				BAN::Math::div_round_up<uint32_t>((vaddr % (BANAN_FB_BPP / 8)) + size, (BANAN_FB_BPP / 8))
-			);
+			const vaddr_t start = BAN::Math::max(vaddr, m_vaddr);
+			const size_t  end   = BAN::Math::min(vaddr + size, m_vaddr + m_size);
+			if (start < end)
+			{
+				do_msync(
+					(start - m_vaddr) / (BANAN_FB_BPP / 8),
+					(end   - start)   / (BANAN_FB_BPP / 8)
+				);
+			}
 
 			return {};
 		}
@@ -300,6 +308,53 @@ namespace Kernel
 			, m_framebuffer(framebuffer)
 		{ }
 
+		void do_msync(uint32_t first_pixel, uint32_t pixel_count)
+		{
+			if (!Processor::get_should_print_cpu_load())
+				return m_framebuffer->sync_pixels_linear(first_pixel, pixel_count);
+
+			const uint32_t fb_width = m_framebuffer->width();
+
+			const auto& font = g_terminal_driver->font();
+
+			const uint32_t x = first_pixel % fb_width;
+			const uint32_t y = first_pixel / fb_width;
+
+			const uint32_t load_w = 16                 * font.width();
+			const uint32_t load_h = Processor::count() * font.height();
+
+			if (y >= load_h || x + pixel_count <= fb_width - load_w)
+				return m_framebuffer->sync_pixels_linear(first_pixel, pixel_count);
+
+			if (x >= fb_width - load_w && x + pixel_count <= fb_width)
+				return;
+
+			if (x < fb_width - load_w)
+				m_framebuffer->sync_pixels_linear(first_pixel, fb_width - load_w - x);
+
+			if (x + pixel_count > fb_width)
+			{
+				const uint32_t past_last_pixel = first_pixel + pixel_count;
+
+				first_pixel = (y + 1) * fb_width;
+				pixel_count = past_last_pixel - first_pixel;
+
+				const uint32_t cpu_load_end = load_h * fb_width;
+
+				while (pixel_count && first_pixel < cpu_load_end)
+				{
+					m_framebuffer->sync_pixels_linear(first_pixel, BAN::Math::min(pixel_count, fb_width - load_w));
+
+					const uint32_t advance = BAN::Math::min(pixel_count, fb_width);
+					pixel_count -= advance;
+					first_pixel += advance;
+				}
+
+				if (pixel_count)
+					m_framebuffer->sync_pixels_linear(first_pixel, pixel_count);
+			}
+		}
+
 	private:
 		BAN::RefPtr<FramebufferDevice> m_framebuffer;
 	};