From d0c71e87b8a5468ea048a73ba8ca0118b629d5b8 Mon Sep 17 00:00:00 2001
From: Federico Villa <61495946+federicovilla55@users.noreply.github.com>
Date: Tue, 3 Sep 2024 08:01:11 +0000
Subject: [PATCH] Add stack-heap collision check to BrkSyscall execution (#346)

Implemented additional checking in the `execute()` method of the brk syscall, defined in the `BrkSyscall` class, to ensure that memory allocation does not overflow stack space.
The System Call has as its argument the new "program break", which is now compared with the current "stack pointer" before being set to avoid overlaps.

Added unit tests to verify that when the `brk` syscall is called with illegal arguments (causing a stack-heap collision), the syscall returns a non-zero value, indicating that it was not executed successfully.
---
 src/syscall/control.h             | 15 +++++++++++--
 test/riscv-tests-64/brk_syscall.S | 35 +++++++++++++++++++++++++++++++
 test/riscv-tests/brk_syscall.s    | 27 ++++++++++++++++++++++++
 3 files changed, 75 insertions(+), 2 deletions(-)
 create mode 100644 test/riscv-tests-64/brk_syscall.S
 create mode 100644 test/riscv-tests/brk_syscall.s

diff --git a/src/syscall/control.h b/src/syscall/control.h
index 179a254a..3822e766 100644
--- a/src/syscall/control.h
+++ b/src/syscall/control.h
@@ -56,8 +56,19 @@ class BrkSyscall : public BaseSyscall {
 
   void execute() {
     // Nothing to do - at the moment we allow infinite growth of the heap.
-    // @todo: Add checks to ensure that the heap doesn't grow into the stack
-    // segment...!
+
+    // Retrieves the argument of the brk syscall, the new program break
+    uint64_t newBreak = BaseSyscall::getArg(BaseSyscall::REG_FILE, 0);
+    // Retrieve the current stack pointer from the processor handler
+    uint64_t stackPointer =
+        ProcessorHandler::getProcessor()->getRegister(BaseSyscall::REG_FILE, 2);
+    if (newBreak >= stackPointer) {
+      SystemIO::printString(
+          "Error: Attempted to allocate memory overlapping stack segment\n");
+      BaseSyscall::setRet(BaseSyscall::REG_FILE, 0, -1); // Syscall error code
+      return;
+    }
+
     BaseSyscall::setRet(BaseSyscall::REG_FILE, 0, 0);
   }
 };
diff --git a/test/riscv-tests-64/brk_syscall.S b/test/riscv-tests-64/brk_syscall.S
new file mode 100644
index 00000000..22560e75
--- /dev/null
+++ b/test/riscv-tests-64/brk_syscall.S
@@ -0,0 +1,35 @@
+.text
+ .globl _start
+ _start: nop
+
+#-------------------------------------------------------------
+# Stack-Heap Collision brk syscall Test
+#-------------------------------------------------------------
+
+test_0: # Test Ecall Return Value
+    mv a0, sp # Assign the ecall parameter to the stack pointer value 
+  	lw a7, BRK # Load the brk syscall immediate value
+    ecall
+
+    # If the return value of the ecall is zero, then the ecall
+    # was executed successfully and therefore the test failed
+    beqz a0, fail
+
+  bne x0, gp, pass
+ 
+pass: 
+ li a0, 42
+ li a7, 93
+ ecall
+
+
+fail: 
+ li a0, 0
+ li a7, 93
+ ecall
+
+.data
+.align 4
+
+# Syscall immediate value
+BRK: .word 214
\ No newline at end of file
diff --git a/test/riscv-tests/brk_syscall.s b/test/riscv-tests/brk_syscall.s
new file mode 100644
index 00000000..13df7d61
--- /dev/null
+++ b/test/riscv-tests/brk_syscall.s
@@ -0,0 +1,27 @@
+.text
+
+#-------------------------------------------------------------
+# Stack-Heap Collision brk syscall Test
+#-------------------------------------------------------------
+
+test_0: # Test Ecall Return Value
+    mv a0, sp # Assign the ecall parameter to the stack pointer value 
+  	li a7, 214 # Load the brk syscall immediate value
+    ecall
+
+    # If the return value of the ecall is zero, then the ecall
+    # was executed successfully and therefore the test failed
+    beqz a0, fail
+
+  bne x0, gp, pass
+ 
+pass: 
+ li a0, 42
+ li a7, 93
+ ecall
+
+
+fail: 
+ li a0, 0
+ li a7, 93
+ ecall
\ No newline at end of file