Merge pull request #42 from h7x4/fix-string-pool-panic

Fix panic on request of invalid string pool index
jiegec · Jan 14, 2024 · 0b7c03a · 0b7c03a
2 parents 28f4f7a + 8c0d86c
commit 0b7c03a
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 11 deletions.
diff --git a/src/device.rs b/src/device.rs
@@ -354,11 +354,10 @@ impl UsbDevice {
                                         desc.resize(setup_packet.length as usize, 0);
                                     }
                                     Ok(desc)
-                                } else {
-                                    let s = &self.string_pool[&index];
+                                } else if let Some(s) = &self.string_pool.get(&index) {
                                     let bytes: Vec<u16> = s.encode_utf16().collect();
                                     let mut desc = vec![
-                                        (2 + bytes.len() * 2) as u8,  // bLength
+                                        2 + bytes.len() as u8 * 2,    // bLength
                                         DescriptorType::String as u8, // bDescriptorType
                                     ];
                                     for byte in bytes {
@@ -371,6 +370,11 @@ impl UsbDevice {
                                         desc.resize(setup_packet.length as usize, 0);
                                     }
                                     Ok(desc)
+                                } else {
+                                    Err(std::io::Error::new(
+                                        std::io::ErrorKind::InvalidInput,
+                                        format!("Invalid string index: {}", index),
+                                    ))
                                 }
                             }
                             Some(DeviceQualifier) => {
@@ -514,4 +518,33 @@ mod test {
         assert_eq!(device.string_pool[&3], "test");
         assert_eq!(device.string_pool[&4], "test");
     }
+
+    #[tokio::test]
+    async fn test_invalid_string_index() {
+        setup_test_logger();
+        let device = UsbDevice::new(0);
+        let res = device
+            .handle_urb(
+                UsbEndpoint {
+                    address: 0x80, // IN
+                    attributes: EndpointAttributes::Control as u8,
+                    max_packet_size: EP0_MAX_PACKET_SIZE,
+                    interval: 0,
+                },
+                None,
+                0,
+                SetupPacket {
+                    request_type: 0b10000000,
+                    request: StandardRequest::GetDescriptor as u8,
+                    // string pool only contains 4 strings, 5 should be invalid
+                    value: (DescriptorType::String as u16) << 8 | 5,
+                    index: 0,
+                    length: 0,
+                },
+                &[],
+            )
+            .await;
+
+        assert!(res.is_err());
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -372,15 +372,22 @@ pub async fn handler<T: AsyncReadExt + AsyncWriteExt + Unpin>(
                                 SetupPacket::parse(&setup),
                                 &data,
                             )
-                            .await?;
-
-                        if out {
-                            trace!("<-Wrote {}", data.len());
-                        } else {
-                            trace!("<-Resp {:02x?}", resp);
+                            .await;
+
+                        match resp {
+                            Ok(resp) => {
+                                if out {
+                                    trace!("<-Wrote {}", data.len());
+                                } else {
+                                    trace!("<-Resp {:02x?}", resp);
+                                }
+                                UsbIpResponse::usbip_ret_submit_success(&header, 0, 0, resp, vec![])
+                            }
+                            Err(err) => {
+                                warn!("Error handling URB: {}", err);
+                                UsbIpResponse::usbip_ret_submit_fail(&header)
+                            }
                         }
-
-                        UsbIpResponse::usbip_ret_submit_success(&header, 0, 0, resp, vec![])
                     }
                 };
                 res.write_to_socket(socket).await?;