From 6e998c8cd96e7e5c35baa4f64ae2197e44725dc5 Mon Sep 17 00:00:00 2001
From: Yury Popov <github@phoenix.dj>
Date: Sun, 19 May 2019 18:31:47 +0300
Subject: [PATCH] Ranged random support

---
 src/kernlib/include/kernlib/rand.hpp | 4 ++++
 src/platform/pagetable.cpp           | 4 ++--
 src/process/process.cpp              | 4 ++--
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/kernlib/include/kernlib/rand.hpp b/src/kernlib/include/kernlib/rand.hpp
index ce512d3..53ffce2 100644
--- a/src/kernlib/include/kernlib/rand.hpp
+++ b/src/kernlib/include/kernlib/rand.hpp
@@ -10,4 +10,8 @@ class RAND {
   static uint64_t _get64();
  public:
   template<typename T> static inline T get() { return static_cast<T>(_get64()); }
+  template<typename T> static inline T get(T max) { return get(0, max); }
+  template<typename T> static inline T get(T min, T max) {
+    return min + (get<T>() % (max - min));
+  }
 };
diff --git a/src/platform/pagetable.cpp b/src/platform/pagetable.cpp
index 3806ac4..c80b6c4 100644
--- a/src/platform/pagetable.cpp
+++ b/src/platform/pagetable.cpp
@@ -102,7 +102,7 @@ void Pagetable::init() {
         EFI::getImageHandle(), &EFI::GUID_LoadedImageProtocol,
         reinterpret_cast<void**>(&loaded_image));
 
-    uintptr_t ptbase = 0x600000 + (RAND::get<uintptr_t>() & 0x3FFF000);
+    uintptr_t ptbase = RAND::get<uintptr_t>(0x800, 0x10000) << 12;
 
     pagetable = static_cast<PTE*>(efiAllocatePage(ptbase, ST));
     efiMapPage(pagetable, nullptr, ST, 0);
@@ -135,7 +135,7 @@ void Pagetable::init() {
     static const size_t pdpe_num = 64;
     static const size_t ptsz = (3 + pdpe_num) * 0x1000;
 
-    uintptr_t ptbase = 0x600000 - ptsz + ((RAND::get<uintptr_t>() & 0x3FFF) << 12);
+    uintptr_t ptbase = RAND::get<uintptr_t>(0x800, 0x8000 - (3 + pdpe_num)) << 12;
 
     PTE *newpt = reinterpret_cast<PTE*>(ptbase);
     Memory::fill(newpt, 0, ptsz);
diff --git a/src/process/process.cpp b/src/process/process.cpp
index f48b1ff..7c1a543 100644
--- a/src/process/process.cpp
+++ b/src/process/process.cpp
@@ -12,8 +12,8 @@ Process::Process() {
   id = size_t(-1);
   pagetable = nullptr;
   entry = 0;
-  _aslrCode = ((RAND::get<uintptr_t>() << 12) & 0x7FFFFFFF000) | 0x40000000000;
-  _aslrStack = ((RAND::get<uintptr_t>() << 12) & 0x7FFFFFFF000) | 0x80000000000;
+  _aslrCode = RAND::get<uintptr_t>(0x80000000llu, 0x100000000llu) << 12;
+  _aslrStack = RAND::get<uintptr_t>(0x40000000llu, 0x80000000llu) << 12;
 }
 Process::~Process() {
   if (pagetable != nullptr) {