From 780263dedaa289dd707af387529674ef63804b60 Mon Sep 17 00:00:00 2001
From: Amlal <amlal@nekernel.org>
Date: Sun, 27 Apr 2025 13:35:09 +0200
Subject: StorageKit: made DMA.h stronger and more resilient.

Signed-off-by: Amlal <amlal@nekernel.org>
---
 dev/kernel/HALKit/AMD64/Storage/AHCI+Generic.cc |  4 ++++
 dev/kernel/StorageKit/DMA.h                     | 17 ++++++++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

(limited to 'dev/kernel')

diff --git a/dev/kernel/HALKit/AMD64/Storage/AHCI+Generic.cc b/dev/kernel/HALKit/AMD64/Storage/AHCI+Generic.cc
index a28fa0ab..722baf57 100644
--- a/dev/kernel/HALKit/AMD64/Storage/AHCI+Generic.cc
+++ b/dev/kernel/HALKit/AMD64/Storage/AHCI+Generic.cc
@@ -220,10 +220,14 @@ STATIC Void drv_std_input_output_ahci(UInt64 lba, UInt8* buffer, SizeT sector_sz
 
   if (kSATAHba->Is & kHBAErrTaskFile) {
     kout << "AHCI Task File Error during I/O.\r";
+    
+    rtl_dma_free(size_buffer);
     err_global_get() = kErrorDiskIsCorrupted;
     return;
   } else {
     rt_copy_memory(ptr, buffer, size_buffer);
+    rtl_dma_free(size_buffer);
+
     err_global_get() = kErrorSuccess;
   }
 }
diff --git a/dev/kernel/StorageKit/DMA.h b/dev/kernel/StorageKit/DMA.h
index 953683eb..17dda29a 100644
--- a/dev/kernel/StorageKit/DMA.h
+++ b/dev/kernel/StorageKit/DMA.h
@@ -33,6 +33,10 @@ namespace Kernel {
 inline UInt8* kDmaPoolPtr = (UInt8*) DMA_POOL_START;
 
 inline VoidPtr rtl_dma_alloc(SizeT size, SizeT align) {
+  if (!size) {
+    return nullptr;
+  }
+
   UIntPtr addr = (UIntPtr) kDmaPoolPtr;
 
   addr = (addr + (align - 1)) & ~(align - 1);  // Align up
@@ -47,12 +51,23 @@ inline VoidPtr rtl_dma_alloc(SizeT size, SizeT align) {
   return (VoidPtr) addr;
 }
 
+inline Void rtl_dma_free(SizeT size) {
+  if (!size)
+    return;
+
+  kDmaPoolPtr = (UInt8*) (kDmaPoolPtr - size);
+}
+
 inline Void rtl_dma_flush(Void* ptr, SizeT size_buffer) {
+  if (ptr > (Void*)(DMA_POOL_START + DMA_POOL_SIZE))
+  {
+    return;
+  }
+
   for (SizeT i = 0; i < size_buffer; ++i) {
     asm volatile("clflush (%0)" : : "r"((UInt8*) ptr + i) : "memory");
   }
   asm volatile("mfence" ::: "memory");
 }
 }  // namespace Kernel
-#else
 #endif
\ No newline at end of file
-- 
cgit v1.2.3