From 7d485e7c7e61eb392f021c2bd21b801e04691caf Mon Sep 17 00:00:00 2001
From: Will Toohey <will@mon.im>
Date: Sun, 4 Apr 2021 17:16:19 +1000
Subject: [PATCH] Improve robustness of p4iodrv/usb

---
 src/main/p4iodrv/device.c | 24 ++++++++++++++++--------
 src/main/p4iodrv/usb.c    | 28 +++++++++++-----------------
 2 files changed, 27 insertions(+), 25 deletions(-)

diff --git a/src/main/p4iodrv/device.c b/src/main/p4iodrv/device.c
index cdf7737..fc511c2 100644
--- a/src/main/p4iodrv/device.c
+++ b/src/main/p4iodrv/device.c
@@ -27,13 +27,24 @@ static bool p4io_transfer(
     void *resp_payload,
     size_t resp_payload_len
 ) {
+    size_t transferred_response_payload = resp_payload_len;
     bool ret = p4io_usb_transfer(ctx->bulk_handle, cmd, ctx->seq_no,
         req_payload, req_payload_len,
-        resp_payload, &resp_payload_len);
+        resp_payload, &transferred_response_payload);
 
     ctx->seq_no++;
 
-    return ret;
+    if(!ret) {
+        return false;
+    }
+
+    if(resp_payload_len && transferred_response_payload != resp_payload_len) {
+        log_warning("Asked for %zu bytes got %zu", resp_payload_len,
+            transferred_response_payload);
+        return false;
+    }
+
+    return true;
 }
 
 struct p4iodrv_ctx *p4iodrv_open(void) {
@@ -73,8 +84,7 @@ bool p4iodrv_read_jamma(struct p4iodrv_ctx *ctx, uint32_t jamma[4]) {
 
 // send something you don't expect a response for
 static bool p4io_send(struct p4iodrv_ctx *ctx, uint8_t cmd) {
-    uint8_t dummy[P4IO_MAX_PAYLOAD];
-    return p4io_transfer(ctx, cmd, NULL, 0, dummy, sizeof(dummy));
+    return p4io_transfer(ctx, cmd, NULL, 0, NULL, 0);
 }
 
 // real IO does not check the return value, so neither do we
@@ -112,11 +122,9 @@ bool p4iodrv_cmd_device_info(struct p4iodrv_ctx *ctx, struct p4io_resp_device_in
 }
 
 bool p4iodrv_cmd_portout(struct p4iodrv_ctx *ctx, const uint8_t buffer[16]) {
-    uint8_t dummy[P4IO_MAX_PAYLOAD];
-    return p4io_transfer(ctx, P4IO_CMD_SET_PORTOUT, buffer, 16, dummy, sizeof(dummy));
+    return p4io_transfer(ctx, P4IO_CMD_SET_PORTOUT, buffer, 16, NULL, 0);
 }
 
 bool p4iodrv_cmd_coinstock(struct p4iodrv_ctx *ctx, const uint8_t buffer[4]) {
-    uint8_t dummy[P4IO_MAX_PAYLOAD];
-    return p4io_transfer(ctx, P4IO_CMD_COINSTOCK, buffer, 4, dummy, sizeof(dummy));
+    return p4io_transfer(ctx, P4IO_CMD_COINSTOCK, buffer, 4, NULL, 0);
 }
diff --git a/src/main/p4iodrv/usb.c b/src/main/p4iodrv/usb.c
index f7d2af2..858a212 100644
--- a/src/main/p4iodrv/usb.c
+++ b/src/main/p4iodrv/usb.c
@@ -135,27 +135,21 @@ bool p4io_usb_transfer(
         return false;
     }
 
-    bytes_requested = P4IO_CMD_HEADER_LEN + *resp_payload_len;
-    // must do this or requests can stall
-    if(bytes_requested == 64) {
-        bytes_requested = 65;
-    }
-    if(!ReadFile(bulk_handle, &cmd_buf, bytes_requested, &bytes_xferred, 0)) {
-        log_warning("ReadFile failed");
+    // must be 65 bytes or requests can stall - only 64 will ever be returned
+    if(!ReadFile(bulk_handle, &cmd_buf, 65, &bytes_xferred, 0)) {
+        log_warning("ReadFile (%u) failed", bytes_requested);
         return false;
     }
 
-    if(bytes_xferred > 0) {
-        if(resp_payload_len) {
-            if(*resp_payload_len < cmd_buf.header.payload_len) {
-                log_warning("Response buffer too short");
-                return false;
-            }
-
-            memcpy(resp_payload, cmd_buf.payload, *resp_payload_len);
-            *resp_payload_len = cmd_buf.header.payload_len;
+    if(bytes_xferred >= sizeof(struct p4io_cmd_header)) {
+        if(*resp_payload_len < cmd_buf.header.payload_len) {
+            log_warning("Response buffer too short");
+            return false;
         }
 
+        memcpy(resp_payload, cmd_buf.payload, *resp_payload_len);
+        *resp_payload_len = cmd_buf.header.payload_len;
+
         if(cmd_buf.header.AA != 0xAA) {
             log_warning("Response bad header");
             return false;
@@ -165,7 +159,7 @@ bool p4io_usb_transfer(
             log_warning("seq_num mismatch (ours %d !=  theirs %d)", seq_no, cmd_buf.header.seq_num);
             return false;
         }
-    } else if(resp_payload_len) {
+    } else {
         *resp_payload_len = 0;
     }