better data housekeeping since we're using the on-device heap now... don't want to leave a mess

c8_remote/include/checkm8.h

@@ -36,7 +36,8 @@ struct pwned_device
     unsigned int idVendor;
     unsigned int idProduct;
 
-    struct payload *installed;
+    struct payload *inst_pl;
+    struct data *inst_data;
 
 #ifdef WITH_ARDUINO
     int ard_fd;

c8_remote/include/payload.h

@@ -23,6 +23,7 @@ typedef unsigned long long DEV_PTR_T;
 
 int install_payload(struct pwned_device *dev, PAYLOAD_T p, LOCATION_T loc);
 int uninstall_payload(struct pwned_device *dev, PAYLOAD_T p);
+int uninstall_all_payloads(struct pwned_device *dev);
 DEV_PTR_T get_payload_address(struct pwned_device *dev, PAYLOAD_T p);
 
 struct dev_cmd_resp *execute_payload(struct pwned_device *dev, PAYLOAD_T p, int response_len, int nargs, ...);
@@ -32,6 +33,7 @@ int kill_payload_async(struct pwned_device *dev, PAYLOAD_T p, DEV_PTR_T buf_addr
 
 DEV_PTR_T install_data(struct pwned_device *dev, LOCATION_T loc, unsigned char *data, int len);
 int uninstall_data(struct pwned_device *dev, DEV_PTR_T ptr);
+int uninstall_all_data(struct pwned_device *dev);
 
 struct dev_cmd_resp *read_gadget(struct pwned_device *dev, DEV_PTR_T addr, int len);
 struct dev_cmd_resp *write_gadget(struct pwned_device *dev, DEV_PTR_T addr, unsigned char *data, int len);

c8_remote/main.c

@@ -44,62 +44,55 @@ void checkm8_debug_block(const char *format, ...)
 #endif
 }
 
-int floppysleep(struct pwned_device *dev)
+void floppysleep(struct pwned_device *dev)
 {
     struct dev_cmd_resp *resp;
 
     if(IS_CHECKM8_FAIL(open_device_session(dev)))
     {
         printf("failed to open device session\n");
-        return -1;
+        return;
     }
 
     if(IS_CHECKM8_FAIL(install_payload(dev, PAYLOAD_SYNC, SRAM)))
     {
         printf("failed to install sync payload\n");
-        return -1;
+        return;
     }
 
     if(IS_CHECKM8_FAIL(install_payload(dev, PAYLOAD_FLOPPYSLEEP, SRAM)))
     {
         printf("failed to install task sleep payload\n");
-        return -1;
+        return;
     }
 
     float init_a = -7.504355E-39f;
-    unsigned long long init_a_ptr = install_data(dev, SRAM, (unsigned char *) &init_a, sizeof(float));
+    DEV_PTR_T init_a_ptr = install_data(dev, SRAM, (unsigned char *) &init_a, sizeof(float));
     if(init_a_ptr == DEV_PTR_NULL)
     {
         printf("failed to write initial data\n");
-        return -1;
+        return;
     }
 
     resp = execute_payload(dev, PAYLOAD_SYNC, 0, 0);
     if(IS_CHECKM8_FAIL(resp->ret))
     {
         printf("failed to execute bootstrap\n");
-        return -1;
+        return;
     }
 
     free_dev_cmd_resp(resp);
 
-    while(1)
+    resp = execute_payload(dev, PAYLOAD_FLOPPYSLEEP, 0, 1, init_a_ptr);
+    if(IS_CHECKM8_FAIL(resp->ret))
     {
-        resp = execute_payload(dev, PAYLOAD_FLOPPYSLEEP, 0, 1, init_a_ptr);
+        printf("failed to execute flopsleep payload\n");
-        if(IS_CHECKM8_FAIL(resp->ret))
+        return;
-        {
-            printf("failed to execute flopsleep payload\n");
-            return -1;
-        }
-
-        printf("retval is %08lli\n", resp->retval);
-        free_dev_cmd_resp(resp);
-
-        usleep(2000000);
     }
 
+    printf("retval is %08lli\n", resp->retval);
+    free_dev_cmd_resp(resp);
     close_device_session(dev);
-    free_device(dev);
 }
 
 void aes_sw(struct pwned_device *dev)
@@ -347,8 +340,10 @@ int main()
     }
 
     demote_device(dev);
+    floppysleep(dev);
 
-    aes_sw(dev);
+    uninstall_all_payloads(dev);
+    uninstall_all_data(dev);
     free_device(dev);
 }
 

c8_remote/src/payload.c

@@ -23,6 +23,15 @@ struct payload
     struct payload *prev;
 };
 
+struct data
+{
+    DEV_PTR_T addr;
+    int len;
+
+    struct data *next;
+    struct data *prev;
+};
+
 struct payload *get_payload(PAYLOAD_T p)
 {
     struct payload *res;
@@ -75,15 +84,10 @@ struct payload *get_payload(PAYLOAD_T p)
     return res;
 }
 
-void free_payload(struct payload *p)
-{
-    free(p);
-}
-
 struct payload *dev_retrieve_payload(struct pwned_device *dev, PAYLOAD_T p)
 {
     struct payload *curr;
-    for(curr = dev->installed; curr != NULL; curr = curr->next)
+    for(curr = dev->inst_pl; curr != NULL; curr = curr->next)
     {
         if(curr->type == p) return curr;
     }
@@ -94,14 +98,14 @@ struct payload *dev_retrieve_payload(struct pwned_device *dev, PAYLOAD_T p)
 int dev_link_payload(struct pwned_device *dev, struct payload *pl)
 {
     struct payload *curr;
-    if(dev->installed == NULL)
+    if(dev->inst_pl == NULL)
     {
-        dev->installed = pl;
+        dev->inst_pl = pl;
         return CHECKM8_SUCCESS;
     }
     else
     {
-        for(curr = dev->installed; curr->next != NULL; curr = curr->next);
+        for(curr = dev->inst_pl; curr->next != NULL; curr = curr->next);
 
         curr->next = pl;
         pl->prev = curr;
@@ -109,11 +113,11 @@ int dev_link_payload(struct pwned_device *dev, struct payload *pl)
     }
 }
 
-int *dev_unlink_payload(struct pwned_device *dev, struct payload *pl)
+int dev_unlink_payload(struct pwned_device *dev, struct payload *pl)
 {
-    if(dev->installed == pl)
+    if(dev->inst_pl == pl)
     {
-        dev->installed = pl->next;
+        dev->inst_pl = pl->next;
         return CHECKM8_SUCCESS;
     }
     else
@@ -126,11 +130,58 @@ int *dev_unlink_payload(struct pwned_device *dev, struct payload *pl)
     }
 }
 
+struct data *dev_retrieve_data(struct pwned_device *dev, DEV_PTR_T addr)
+{
+    struct data *curr;
+    for(curr = dev->inst_data; curr != NULL; curr = curr->next)
+    {
+        if(curr->addr == addr) return curr;
+    }
+
+    return NULL;
+}
+
+int dev_link_data(struct pwned_device *dev, struct data *data)
+{
+    struct data *curr;
+    if(dev->inst_data == NULL)
+    {
+        dev->inst_data = data;
+        return CHECKM8_SUCCESS;
+    }
+    else
+    {
+        for(curr = dev->inst_data; curr->next != NULL; curr = curr->next);
+
+        curr->next = data;
+        data->prev = curr;
+        return CHECKM8_SUCCESS;
+    }
+}
+
+int dev_unlink_data(struct pwned_device *dev, struct data *data)
+{
+    if(dev->inst_data == data)
+    {
+        dev->inst_data = data->next;
+        return CHECKM8_SUCCESS;
+    }
+    else
+    {
+        data->prev->next = data->next;
+        if(data->next != NULL)
+            data->next->prev = data->prev;
+
+        return CHECKM8_SUCCESS;
+    }
+}
+
 DEV_PTR_T get_address(struct pwned_device *dev, LOCATION_T l, int len)
 {
     checkm8_debug_indent("get_address(dev = %p, loc = %i, len = %i)\n", dev, l, len);
     DEV_PTR_T retval;
     unsigned long long malloc_args[2] = {ADDR_DEV_MALLOC, (unsigned long long) len};
+    struct data *new_entry;
 
     struct dev_cmd_resp *resp = dev_exec(dev, 0, 2, malloc_args);
     if(IS_CHECKM8_FAIL(resp->ret))
@@ -143,6 +194,11 @@ DEV_PTR_T get_address(struct pwned_device *dev, LOCATION_T l, int len)
     retval = resp->retval;
     free_dev_cmd_resp(resp);
 
+    new_entry = malloc(sizeof(struct data));
+    new_entry->addr = retval;
+    new_entry->len = len;
+    dev_link_data(dev, new_entry);
+
     checkm8_debug_indent("\tgot address %llX\n", retval);
     return retval;
 }
@@ -150,8 +206,16 @@ DEV_PTR_T get_address(struct pwned_device *dev, LOCATION_T l, int len)
 int free_address(struct pwned_device *dev, LOCATION_T l, DEV_PTR_T ptr)
 {
     struct dev_cmd_resp *resp;
+    struct data *entry;
     unsigned long long free_args[2] = {ADDR_DEV_FREE, ptr};
 
+    entry = dev_retrieve_data(dev, ptr);
+    if(entry == NULL)
+    {
+        checkm8_debug_indent("\tthis pointer was not allocated through the payload interface, not freeing\n");
+        return CHECKM8_FAIL_NOINST;
+    }
+
     resp = dev_exec(dev, 0, 2, free_args);
     if(IS_CHECKM8_FAIL(resp->ret))
     {
@@ -161,6 +225,9 @@ int free_address(struct pwned_device *dev, LOCATION_T l, DEV_PTR_T ptr)
     }
 
     free_dev_cmd_resp(resp);
+    dev_unlink_data(dev, entry);
+    free(entry);
+
     return CHECKM8_SUCCESS;
 }
 
@@ -211,7 +278,24 @@ int uninstall_payload(struct pwned_device *dev, PAYLOAD_T p)
     }
 
     dev_unlink_payload(dev, pl);
-    free_payload(pl);
+    free(pl);
+    return CHECKM8_SUCCESS;
+}
+
+int uninstall_all_payloads(struct pwned_device *dev)
+{
+    checkm8_debug_indent("uninstall_all_payloads(dev = %p)\n");
+    int ret;
+    while(dev->inst_pl != NULL)
+    {
+        ret = uninstall_payload(dev, dev->inst_pl->type);
+        if(IS_CHECKM8_FAIL(ret))
+        {
+            checkm8_debug_indent("\terror while uninstalling\n");
+            return ret;
+        }
+    }
+
     return CHECKM8_SUCCESS;
 }
 
@@ -438,6 +522,24 @@ int uninstall_data(struct pwned_device *dev, DEV_PTR_T addr)
     return free_address(dev, SRAM, addr);
 }
 
+int uninstall_all_data(struct pwned_device *dev)
+{
+    checkm8_debug_indent("uninstall_all_data(dev = %p)\n", dev);
+    int retval;
+
+    while(dev->inst_data != NULL)
+    {
+        retval = uninstall_data(dev, dev->inst_data->addr);
+        if(IS_CHECKM8_FAIL(retval))
+        {
+            checkm8_debug_indent("\terror while uninstalling data\n");
+            return retval;
+        }
+    }
+
+    return CHECKM8_SUCCESS;
+}
+
 struct dev_cmd_resp *read_gadget(struct pwned_device *dev, DEV_PTR_T addr, int len)
 {
     checkm8_debug_indent("read_gadget(dev = %p, addr = %lx, len = %i)\n", dev, addr, len);