Add transfer_buffer_length to handle_urb

jiegec · Aug 23, 2023 · 2233140 · 2233140
1 parent 8506ede
commit 2233140
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 16 deletions.
diff --git a/src/cdc.rs b/src/cdc.rs
@@ -47,6 +47,7 @@ impl UsbInterfaceHandler for UsbCdcAcmHandler {
         &mut self,
         _interface: &UsbInterface,
         ep: UsbEndpoint,
+        _transfer_buffer_length: u32,
         _setup: SetupPacket,
         req: &[u8],
     ) -> Result<Vec<u8>> {

diff --git a/src/device.rs b/src/device.rs
@@ -189,6 +189,7 @@ impl UsbDevice {
         &self,
         ep: UsbEndpoint,
         intf: Option<&UsbInterface>,
+        transfer_buffer_length: u32,
         setup_packet: SetupPacket,
         out_data: &Vec<u8>,
     ) -> Result<Vec<u8>> {
@@ -374,15 +375,14 @@ impl UsbDevice {
                         // only low 8 bits are valid
                         let intf = &self.interfaces[setup_packet.index as usize & 0xFF];
                         let mut handler = intf.handler.lock().unwrap();
-                        let resp = handler.handle_urb(intf, ep, setup_packet, &out_data)?;
-                        return Ok(resp);
+                        handler.handle_urb(intf, ep, transfer_buffer_length, setup_packet, out_data)
                     }
                     _ if setup_packet.request_type & 0xF == 0 && self.device_handler.is_some() => {
                         // to device
                         // see https://www.beyondlogic.org/usbnutshell/usb6.shtml
                         let lock = self.device_handler.as_ref().unwrap();
                         let mut handler = lock.lock().unwrap();
-                        return Ok(handler.handle_urb(setup_packet, &out_data)?);
+                        handler.handle_urb(transfer_buffer_length, setup_packet, out_data)
                     }
                     _ => unimplemented!("control in"),
                 }
@@ -411,15 +411,14 @@ impl UsbDevice {
                         // only low 8 bits are valid
                         let intf = &self.interfaces[setup_packet.index as usize & 0xFF];
                         let mut handler = intf.handler.lock().unwrap();
-                        let resp = handler.handle_urb(intf, ep, setup_packet, &out_data)?;
-                        return Ok(resp);
+                        handler.handle_urb(intf, ep, transfer_buffer_length, setup_packet, out_data)
                     }
                     _ if setup_packet.request_type & 0xF == 0 && self.device_handler.is_some() => {
                         // to device
                         // see https://www.beyondlogic.org/usbnutshell/usb6.shtml
                         let lock = self.device_handler.as_ref().unwrap();
                         let mut handler = lock.lock().unwrap();
-                        return Ok(handler.handle_urb(setup_packet, &out_data)?);
+                        handler.handle_urb(transfer_buffer_length, setup_packet, out_data)
                     }
                     _ => unimplemented!("control out"),
                 }
@@ -428,8 +427,7 @@ impl UsbDevice {
                 // others
                 let intf = intf.unwrap();
                 let mut handler = intf.handler.lock().unwrap();
-                let resp = handler.handle_urb(intf, ep, setup_packet, &out_data)?;
-                return Ok(resp);
+                handler.handle_urb(intf, ep, transfer_buffer_length, setup_packet, out_data)
             }
             _ => unimplemented!("transfer to {:?}", ep),
         }
@@ -440,8 +438,14 @@ impl UsbDevice {
 pub trait UsbDeviceHandler {
     /// Handle a URB(USB Request Block) targeting at this device
     ///
-    /// When the lower 4 bits of bmRequestType is zero and the URB is not handled by the library, this function is called
-    fn handle_urb(&mut self, setup: SetupPacket, req: &[u8]) -> Result<Vec<u8>>;
+    /// When the lower 4 bits of `bmRequestType` is zero and the URB is not handled by the library, this function is called.
+    /// The resulting data should not exceed `transfer_buffer_length`
+    fn handle_urb(
+        &mut self,
+        transfer_buffer_length: u32,
+        setup: SetupPacket,
+        req: &[u8],
+    ) -> Result<Vec<u8>>;
 
     /// Helper to downcast to actual struct
     ///

diff --git a/src/hid.rs b/src/hid.rs
@@ -89,6 +89,7 @@ impl UsbInterfaceHandler for UsbHidKeyboardHandler {
         &mut self,
         _interface: &UsbInterface,
         ep: UsbEndpoint,
+        _transfer_buffer_length: u32,
         setup: SetupPacket,
         _req: &[u8],
     ) -> Result<Vec<u8>> {

diff --git a/src/host.rs b/src/host.rs
@@ -18,14 +18,15 @@ impl UsbInterfaceHandler for UsbHostInterfaceHandler {
         &mut self,
         _interface: &UsbInterface,
         ep: UsbEndpoint,
+        transfer_buffer_length: u32,
         setup: SetupPacket,
         req: &[u8],
     ) -> Result<Vec<u8>> {
         debug!(
             "To host device: ep={:?} setup={:?} req={:?}",
             ep, setup, req
         );
-        let mut buffer = vec![0u8; ep.max_packet_size as usize];
+        let mut buffer = vec![0u8; transfer_buffer_length as usize];
         let timeout = std::time::Duration::new(1, 0);
         let handle = self.handle.lock().unwrap();
         if ep.attributes == EndpointAttributes::Control as u8 {
@@ -104,9 +105,14 @@ impl UsbHostDeviceHandler {
 }
 
 impl UsbDeviceHandler for UsbHostDeviceHandler {
-    fn handle_urb(&mut self, setup: SetupPacket, req: &[u8]) -> Result<Vec<u8>> {
+    fn handle_urb(
+        &mut self,
+        transfer_buffer_length: u32,
+        setup: SetupPacket,
+        req: &[u8],
+    ) -> Result<Vec<u8>> {
         debug!("To host device: setup={:?} req={:?}", setup, req);
-        let mut buffer = [0u8; 1024];
+        let mut buffer = vec![0u8; transfer_buffer_length as usize];
         let timeout = std::time::Duration::new(1, 0);
         let handle = self.handle.lock().unwrap();
         // control

diff --git a/src/interface.rs b/src/interface.rs
@@ -19,11 +19,13 @@ pub trait UsbInterfaceHandler {
 
     /// Handle a URB(USB Request Block) targeting at this interface
     ///
-    /// Can be one of: control transfer to ep0 or other types of transfer to its endpoint
+    /// Can be one of: control transfer to ep0 or other types of transfer to its endpoint.
+    /// The resulting data should not exceed `transfer_buffer_length`.
     fn handle_urb(
         &mut self,
         interface: &UsbInterface,
         ep: UsbEndpoint,
+        transfer_buffer_length: u32,
         setup: SetupPacket,
         req: &[u8],
     ) -> Result<Vec<u8>>;

diff --git a/src/lib.rs b/src/lib.rs
@@ -292,7 +292,13 @@ async fn handler<T: AsyncReadExt + AsyncWriteExt + Unpin>(
                 trace!("->Setup {:02x?}", setup);
                 trace!("->Request {:02x?}", out_data);
                 let resp = device
-                    .handle_urb(usb_ep, intf, SetupPacket::parse(&setup), &out_data)
+                    .handle_urb(
+                        usb_ep,
+                        intf,
+                        transfer_buffer_length,
+                        SetupPacket::parse(&setup),
+                        &out_data,
+                    )
                     .await?;
 
                 if out {
@@ -315,7 +321,7 @@ async fn handler<T: AsyncReadExt + AsyncWriteExt + Unpin>(
                     // In the out endpoint case, the actual_length field should be
                     // same as the data length received in the original URB transaction.
                     // No data bytes are sent
-                    transfer_buffer_length as u32
+                    transfer_buffer_length
                 } else {
                     resp.len() as u32
                 };