From b590cd65c81da9637fc5a72eff967ed6bab49f80 Mon Sep 17 00:00:00 2001
From: Gregor Haas <gregorhaas1997@gmail.com>
Date: Tue, 7 Jan 2020 11:17:34 -0500
Subject: [PATCH] Same as prev, but on remote side

---
 c8_remote/src/command.c     |  8 ++++----
 c8_remote/src/usb_helpers.c | 31 ++++++++++---------------------
 2 files changed, 14 insertions(+), 25 deletions(-)

diff --git a/c8_remote/src/command.c b/c8_remote/src/command.c
index d2a94b0..30d5736 100644
--- a/c8_remote/src/command.c
+++ b/c8_remote/src/command.c
@@ -98,7 +98,7 @@ struct dev_cmd_resp *command(struct pwned_device *dev,
         return cmd_resp;
     }
 
-    ret = dfu_send_data(dev, args, arg_len, 1);
+    ret = dfu_send_data(dev, args, arg_len, 0);
     if(IS_CHECKM8_FAIL(ret))
     {
         cmd_resp->ret = ret;
@@ -110,7 +110,7 @@ struct dev_cmd_resp *command(struct pwned_device *dev,
         ret = ctrl_transfer(dev,
                             0xA1, 2, 0xFFFF, 0,
                             resp_buf, response_len + 1,
-                            100, 0);
+                            100, 1);
         if(ret >= 0) checkm8_debug_indent("\tfinal request transferred %i bytes\n", ret);
         else
         {
@@ -124,7 +124,7 @@ struct dev_cmd_resp *command(struct pwned_device *dev,
         ret = ctrl_transfer(dev,
                             0xA1, 2, 0xFFFF, 0,
                             resp_buf, response_len,
-                            100, 0);
+                            100, 1);
         if(ret >= 0) checkm8_debug_indent("\tfinal request transferred %i bytes\n", ret);
         else
         {
@@ -261,5 +261,5 @@ struct dev_cmd_resp *dev_write_memory(struct pwned_device *dev, long long addr,
     ((unsigned long *) cmd_args)[4] = len;
     memcpy(&cmd_args[40], data, len);
 
-    return command(dev, (unsigned char *) &cmd_args, 40 + len, 1 * sizeof(unsigned long long));
+    return command(dev, cmd_args, 40 + len, 1 * sizeof(unsigned long long));
 }
\ No newline at end of file
diff --git a/c8_remote/src/usb_helpers.c b/c8_remote/src/usb_helpers.c
index 3de78fb..ebe0bfa 100644
--- a/c8_remote/src/usb_helpers.c
+++ b/c8_remote/src/usb_helpers.c
@@ -435,23 +435,12 @@ int no_error_ctrl_transfer_data(struct pwned_device *dev,
         checkm8_debug_indent("\treceived argument ack\n");
         while(index < data_len)
         {
-            if(data_len - index > ARD_BUF_SIZE) amount = ARD_BUF_SIZE;
-            else amount = data_len - index;
-
+            amount = 0;
+            while(read(dev->ard_fd, &amount, 2) == 0);
             checkm8_debug_indent("\twriting data chunk of size %i\n", amount);
             write(dev->ard_fd, &data[index], amount);
 
-            while(read(dev->ard_fd, &buf, 1) == 0);
-            if(buf == PROT_ACK)
-            {
-                checkm8_debug_indent("\treceived data ack\n");
-                index += amount;
-            }
-            else
-            {
-                checkm8_debug_indent("\treceived unexpected response %x\n", buf);
-                return CHECKM8_FAIL_PROT;
-            }
+            index += amount;
         }
 
         while(read(dev->ard_fd, &buf, 1) == 0);
@@ -487,7 +476,7 @@ int ctrl_transfer(struct pwned_device *dev,
             dev, bmRequestType, bRequest, wValue, wIndex, data, data_len, timeout);
 
 #ifdef WITH_ARDUINO
-    int amount, index;
+    int amount, index, size;
     char buf;
     struct usb_xfer_args args;
     args.bmRequestType = bmRequestType;
@@ -511,16 +500,16 @@ int ctrl_transfer(struct pwned_device *dev,
             while(amount < data_len)
             {
                 // get the size of this chunk
-                while(read(dev->ard_fd, &buf, 1) == 0);
-                checkm8_debug_indent("\treceiving data chunk of size %i\n", buf);
+                while(read(dev->ard_fd, &size, 2) == 0);
+                checkm8_debug_indent("\treceiving data chunk of size %i\n", size);
 
                 index = 0;
-                while(index < buf)
+                while(index < size)
                 {
-                    index += read(dev->ard_fd, &data[amount + index], buf - index);
+                    index += read(dev->ard_fd, &data[amount + index], size - index);
                 }
 
-                amount += buf;
+                amount += size;
             }
         }
         else
@@ -529,7 +518,7 @@ int ctrl_transfer(struct pwned_device *dev,
             while(index < data_len)
             {
                 amount = 0;
-                while(read(dev->ard_fd, &amount, 1) == 0);
+                while(read(dev->ard_fd, &amount, 2) == 0);
                 checkm8_debug_indent("\twriting data chunk of size %i\n", amount);
                 write(dev->ard_fd, &data[index], amount);