Add timeout and overflow check in RX buffer handling to prevent hang …

…on disconnect
sekigon-gonnoc · Dec 16, 2024 · 447ea43 · 447ea43
1 parent bc14d3e
commit 447ea43
Showing 1 changed file with 4 additions and 13 deletions.
diff --git a/src/pio_usb.c b/src/pio_usb.c
@@ -178,6 +178,8 @@ int __no_inline_not_in_flash_func(pio_usb_bus_receive_packet_and_handshake)(
   bool crc_match = false;
   int16_t t = 240;
   uint16_t idx = 0;
+  uint16_t nak_timeout = 10000;
+  const uint16_t rx_buf_len = sizeof(pp->usb_rx_buffer) / sizeof(pp->usb_rx_buffer[0]);
 
   while (t--) {
     if (pio_sm_get_rx_fifo_level(pp->pio_usb_rx, pp->sm_rx)) {
@@ -192,25 +194,18 @@ int __no_inline_not_in_flash_func(pio_usb_bus_receive_packet_and_handshake)(
   // timing critical start
   if (t > 0) {
     if (handshake == USB_PID_ACK) {
-      uint32_t timeout = 240;
-      while ((pp->pio_usb_rx->irq & IRQ_RX_COMP_MASK) == 0 && timeout--) {
+      while ((pp->pio_usb_rx->irq & IRQ_RX_COMP_MASK) == 0 && idx < rx_buf_len - 1) {
         if (pio_sm_get_rx_fifo_level(pp->pio_usb_rx, pp->sm_rx)) {
           uint8_t data = pio_sm_get(pp->pio_usb_rx, pp->sm_rx) >> 24;
           crc_prev2 = crc_prev;
           crc_prev = crc;
           crc = update_usb_crc16(crc, data);
-          if (idx > (sizeof(pp->usb_rx_buffer) / sizeof(pp->usb_rx_buffer[0]))) {
-            return -1;
-          }
           pp->usb_rx_buffer[idx++] = data;
           crc_receive = (crc_receive >> 8) | (data << 8);
           crc_receive_inverse = crc_receive ^ 0xffff;
           crc_match = (crc_receive_inverse == crc_prev2);
         }
       }
-      if (timeout == 0) {
-        return -1;
-      }
 
       if (idx >= 4 && crc_match) {
         pio_usb_bus_send_handshake(pp, USB_PID_ACK);
@@ -219,14 +214,10 @@ int __no_inline_not_in_flash_func(pio_usb_bus_receive_packet_and_handshake)(
       }
     } else {
       // just discard received data since we NAK/STALL anyway
-      uint32_t timeout = 240;
-      while ((pp->pio_usb_rx->irq & IRQ_RX_COMP_MASK) == 0 && timeout--) {
+      while ((pp->pio_usb_rx->irq & IRQ_RX_COMP_MASK) == 0 && nak_timeout--) {
         continue;
       }
       pio_sm_clear_fifos(pp->pio_usb_rx, pp->sm_rx);
-      if (timeout == 0) {
-        return -1;
-      }
 
       pio_usb_bus_send_handshake(pp, handshake);
     }