From 5e4e6d2db892811d269e7aef5055c0e5ed6095fe Mon Sep 17 00:00:00 2001
From: xamidev <121681048+xamidev@users.noreply.github.com>
Date: Tue, 10 Sep 2024 20:14:24 +0200
Subject: [PATCH] Fix: safety: malloc and free BMP image; memcpy sanitizing

---
 src/kernel/initrd.c | 38 +++++++++++++++++++++++++++++++++++---
 src/kernel/initrd.h |  1 +
 src/kernel/kmain.c  |  2 +-
 src/programs/bmp.c  |  9 ++++-----
 4 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/src/kernel/initrd.c b/src/kernel/initrd.c
index 70788d6..b1ab2fe 100644
--- a/src/kernel/initrd.c
+++ b/src/kernel/initrd.c
@@ -131,9 +131,16 @@ int tar_file_to_buffer(uint8_t* initrd, const char* filename, char* buffer)
 		if (strcmp(file_name, filename) == 0)
 		{
 			uint8_t* file_data = current_block + TAR_BLOCK_SIZE;
-			memcpy(buffer, file_data, file_size);
-			buffer[file_size] = '\0';
-			return 0;
+			if (sizeof(buffer) >= sizeof(file_data))
+			{
+				memcpy(buffer, file_data, file_size);
+				buffer[file_size] = '\0';
+				return 0;
+			} else {
+				printf("Invalid destination buffer size %d bytes < %d bytes\n", sizeof(buffer), sizeof(file_data));
+				return -1;
+			}
+			return -1;
 		}
 
 		uint32_t total_size = ((file_size + TAR_BLOCK_SIZE - 1) / TAR_BLOCK_SIZE) * TAR_BLOCK_SIZE;
@@ -142,3 +149,28 @@ int tar_file_to_buffer(uint8_t* initrd, const char* filename, char* buffer)
 	printf("[tar] file '%s' not found\n", filename);
 	return -1;
 }
+
+uint32_t tar_get_file_size(uint8_t* initrd, const char* filename)
+{
+	uint8_t* current_block = initrd;
+
+	while (1)
+	{
+		if (current_block[0] == '\0')
+		{
+			return -1;
+		}
+
+		const char* file_name = (const char*)current_block;
+		uint32_t file_size = tar_parse_size((const char*)(current_block+124));
+
+		if (strcmp(file_name, filename) == 0)
+		{
+			return file_size;
+		}
+
+		uint32_t total_size = ((file_size + TAR_BLOCK_SIZE - 1) / TAR_BLOCK_SIZE) * TAR_BLOCK_SIZE;
+		current_block += TAR_BLOCK_SIZE + total_size;
+	}
+	return -1;
+}
diff --git a/src/kernel/initrd.h b/src/kernel/initrd.h
index 43ffd29..714bf84 100644
--- a/src/kernel/initrd.h
+++ b/src/kernel/initrd.h
@@ -34,5 +34,6 @@ void tar_find_file(uint8_t *tar_start, const char* filename);
 void ls_initrd(uint8_t* initrd, int verbose);
 void cat_initrd(uint8_t* initrd, const char* filename);
 int tar_file_to_buffer(uint8_t* initrd, const char* filename, char* buffer);
+uint32_t tar_get_file_size(uint8_t* initrd, const char* filename);
 
 #endif
diff --git a/src/kernel/kmain.c b/src/kernel/kmain.c
index ac7f6c3..03a9886 100644
--- a/src/kernel/kmain.c
+++ b/src/kernel/kmain.c
@@ -79,7 +79,7 @@ void kmain(multiboot2_info *mb_info)
 		if (mmap->addr != 0)
 		{
 			
-			serial_printf(3, "base addr=0x%x%x, length=0x%x%x, type=%u\n",
+			serial_printf(3, "base addr=0x%x%x, length=0x%x%x, type=%u",
 				(uint32_t) (mmap->addr >> 32),
 				(uint32_t) (mmap->addr & 0xFFFFFFFF),
 				(uint32_t) (mmap->len >> 32),
diff --git a/src/programs/bmp.c b/src/programs/bmp.c
index 16e8654..95c47b0 100644
--- a/src/programs/bmp.c
+++ b/src/programs/bmp.c
@@ -9,6 +9,7 @@
 #include "../drivers/framebuffer.h"
 #include "../libc/stdio.h"
 #include "../drivers/serial.h"
+#include "../kernel/kheap.h"
 
 #pragma pack(push, 1)
 typedef struct
@@ -38,11 +39,8 @@ typedef struct
 
 void display_bmp(uint32_t* fb, int pitch, int bpp, uint8_t* initrd, const char* filename)
 {
-	// Should use dynamic allocation when heap works
-	// Cannot go more than ~500k size for buffer
-	// Fail zone 450k->470k
-	// So right now the max should be 400kb img size
-	char buffer[400*1000];
+	uint32_t buf_size = tar_get_file_size(initrd, filename);
+	char* buffer = (char*)malloc(buf_size);
 	int file_status = tar_file_to_buffer(initrd, filename, buffer);
 
 	if (file_status != 0)
@@ -87,6 +85,7 @@ void display_bmp(uint32_t* fb, int pitch, int bpp, uint8_t* initrd, const char*
 
 	// Update cursor pos after image drawing
 	move_cursor(get_cursor_x(), get_cursor_y()+(height/16)+2);
+	free(buffer);
 }
 
 void program_bmp(int argc, char* argv[])