From efc7cbf65312178d5a050fa5112d2a276673a81c Mon Sep 17 00:00:00 2001
From: plutoo <plutoo@universe>
Date: Fri, 12 Jan 2018 19:10:23 +0100
Subject: [PATCH] Implement heap override

---
 nx/include/switch/result.h      |  2 ++
 nx/include/switch/runtime/env.h |  9 +++++--
 nx/source/runtime/env.c         | 14 +++++++++-
 nx/source/runtime/init.c        | 48 ++++++++++++++++++++++++++-------
 4 files changed, 60 insertions(+), 13 deletions(-)

diff --git a/nx/include/switch/result.h b/nx/include/switch/result.h
index 8f4b28af..0b42e141 100644
--- a/nx/include/switch/result.h
+++ b/nx/include/switch/result.h
@@ -31,4 +31,6 @@
 #define LIBNX_BADINPUT       10
 #define LIBNX_BADREENT       11
 #define LIBNX_BUFFERPRODUCER_ERROR 12
+#define LIBNX_HANDLETOOEARLY 13
+#define LIBNX_HEAPALLOCFAILED 14
 #define LIBNX_PARCEL_ERRBASE 100
diff --git a/nx/include/switch/runtime/env.h b/nx/include/switch/runtime/env.h
index 2af06ebc..1cf51116 100644
--- a/nx/include/switch/runtime/env.h
+++ b/nx/include/switch/runtime/env.h
@@ -1,5 +1,10 @@
 // Copyright 2018 plutoo
 
-void   envParse(void* ctx, Handle main_thread);
+void envParse(void* ctx, Handle main_thread);
+
 Handle envGetMainThreadHandle(void);
-bool   envIsNso(void);
+bool envIsNso(void);
+
+bool  envHasHeapOverride(void);
+void* envGetHeapOverrideAddr(void);
+u64   envGetHeapOverrideSize(void);
diff --git a/nx/source/runtime/env.c b/nx/source/runtime/env.c
index 70a5afa7..239e554c 100644
--- a/nx/source/runtime/env.c
+++ b/nx/source/runtime/env.c
@@ -83,9 +83,21 @@ Handle envGetMainThreadHandle(void) {
         return g_mainThreadHandle;
     }
 
-    fatalSimple(1);
+    fatalSimple(MAKERESULT(MODULE_LIBNX, LIBNX_HANDLETOOEARLY));
 }
 
 bool envIsNso(void) {
     return g_isNso;
 }
+
+bool envHasHeapOverride(void) {
+    return g_overrideHeapAddr != NULL;
+}
+
+void* envGetHeapOverrideAddr(void) {
+    return g_overrideHeapAddr;
+}
+
+u64 envGetHeapOverrideSize(void) {
+    return g_overrideHeapSize;
+}
diff --git a/nx/source/runtime/init.c b/nx/source/runtime/init.c
index 1ff19801..585ded30 100644
--- a/nx/source/runtime/init.c
+++ b/nx/source/runtime/init.c
@@ -7,20 +7,48 @@ void newlibSetup(void);
 
 void __system_initArgv(void);
 
-#define INNER_HEAP_SIZE 0x200000
-__attribute__((weak)) size_t __nx_inner_heap_size = INNER_HEAP_SIZE;
-__attribute__((weak)) char   __nx_inner_heap[INNER_HEAP_SIZE];
-__attribute__((weak)) size_t __nx_outer_heap_size = 0x2000000*16;//Must be a multiple of 0x2000000.
+// Must be a multiple of 0x2000000.
+__attribute__((weak)) size_t __nx_heap_size = 0x2000000*16;
+
+/*
+  There are three ways of allocating heap:
+
+    - Normal syscall:
+
+    Allocates heap using |svcSetHeapSize|. The size is provided by a weak symbol
+    called |__nx_heap_size|.
+
+    - Heap override:
+
+    Uses existing heap segment as provided by the homebrew loader environment
+    block. This happens automatically if such a setting is provided by the
+    homebrew environment.
+
+    - Custom override:
+
+    A program can override the weak symbol |__libnx_initheap| to setup a
+    different heap. In this case, the global variables |fake_heap_start|
+    and |fake_heap_end| needs to be set appropriately.
+
+    A custom override be used to implement an "inner heap" located in the .bss
+    segment of a process, for example.
+ */
 
 void __attribute__((weak)) __libnx_initheap(void)
 {
-    u64 addr;
-    Result rc   = svcSetHeapSize((void**)&addr, __nx_outer_heap_size);
-    size_t size = __nx_outer_heap_size;
+    void*  addr;
+    size_t size;
 
-    if (R_FAILED(rc)) {
-        addr = (u64) &__nx_inner_heap[0];
-        size = __nx_inner_heap_size;
+    if (envHasHeapOverride()) {
+        addr = envGetHeapOverrideAddr();
+        size = envGetHeapOverrideSize();
+    }
+    else {
+        Result rc = svcSetHeapSize(&addr, __nx_heap_size);
+        size = __nx_heap_size;
+
+        if (R_FAILED(rc))
+            fatalSimple(MAKERESULT(MODULE_LIBNX, LIBNX_HEAPALLOCFAILED));
     }
 
     // Newlib