usb_util: implement usb_transfer

Wrap libusb_transfer and callback.

BUG=chromium:831850
TEST=local build

Change-Id: I1bc84e68cb36796e919f647e3a072f1599f80a4a
Reviewed-on: https://chromium-review.googlesource.com/1138643
Commit-Ready: ChromeOS CL Exonerator Bot <chromiumos-cl-exonerator@appspot.gserviceaccount.com>
Tested-by: Jingkui Wang <jkwang@google.com>
Reviewed-by: Zach Reizner <zachr@chromium.org>

3 files changed, 401 insertions, 0 deletions

diff --git a/usb_util/src/device_handle.rs b/usb_util/src/device_handle.rs
index e07f6f2..57432b2 100644
--- a/usb_util/src/device_handle.rs
+++ b/usb_util/src/device_handle.rs
@@ -8,6 +8,7 @@ use std::sync::Arc;
 use bindings;
 use error::{Error, Result};
 use libusb_context::LibUsbContextInner;
+use usb_transfer::{UsbTransfer, UsbTransferBuffer};
 
 /// DeviceHandle wraps libusb_device_handle.
 pub struct DeviceHandle {
@@ -126,4 +127,16 @@ impl DeviceHandle {
         try_libusb!(unsafe { bindings::libusb_clear_halt(self.handle, endpoint) });
         Ok(())
     }
+
+    /// Libusb asynchronous I/O interface has a 5 step process. It gives lots of
+    /// flexibility but makes it hard to manage object life cycle and easy to
+    /// write unsafe code. We wrap this interface to a simple "transfer" and "cancel"
+    /// interface. Resubmission is not supported and deallocation is handled safely
+    /// here.
+    pub fn submit_async_transfer<T: UsbTransferBuffer>(
+        &self,
+        transfer: UsbTransfer<T>,
+    ) -> Result<()> {
+        unsafe { transfer.submit(self.handle) }
+    }
 }
diff --git a/usb_util/src/lib.rs b/usb_util/src/lib.rs
index f92511c..e115f62 100644
--- a/usb_util/src/lib.rs
+++ b/usb_util/src/lib.rs
@@ -23,3 +23,4 @@ pub mod interface_descriptor;
 pub mod libusb_context;
 pub mod libusb_device;
 pub mod types;
+pub mod usb_transfer;
diff --git a/usb_util/src/usb_transfer.rs b/usb_util/src/usb_transfer.rs
new file mode 100644
index 0000000..93f3979
--- /dev/null
+++ b/usb_util/src/usb_transfer.rs
@@ -0,0 +1,387 @@
+// Copyright 2018 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use std::mem::size_of;
+use std::os::raw::c_void;
+use std::sync::{Arc, Weak};
+
+use bindings::{
+    libusb_alloc_transfer, libusb_cancel_transfer, libusb_device_handle, libusb_free_transfer,
+    libusb_submit_transfer, libusb_transfer, libusb_transfer_status, LIBUSB_TRANSFER_CANCELLED,
+    LIBUSB_TRANSFER_COMPLETED, LIBUSB_TRANSFER_ERROR, LIBUSB_TRANSFER_NO_DEVICE,
+    LIBUSB_TRANSFER_OVERFLOW, LIBUSB_TRANSFER_STALL, LIBUSB_TRANSFER_TIMED_OUT,
+    LIBUSB_TRANSFER_TYPE_BULK, LIBUSB_TRANSFER_TYPE_CONTROL, LIBUSB_TRANSFER_TYPE_INTERRUPT,
+};
+use error::{Error, Result};
+use types::UsbRequestSetup;
+
+/// Status of transfer.
+#[derive(PartialEq)]
+pub enum TransferStatus {
+    Completed,
+    Error,
+    TimedOut,
+    Cancelled,
+    Stall,
+    NoDevice,
+    Overflow,
+}
+
+impl From<libusb_transfer_status> for TransferStatus {
+    fn from(s: libusb_transfer_status) -> Self {
+        match s {
+            LIBUSB_TRANSFER_COMPLETED => TransferStatus::Completed,
+            LIBUSB_TRANSFER_ERROR => TransferStatus::Error,
+            LIBUSB_TRANSFER_TIMED_OUT => TransferStatus::TimedOut,
+            LIBUSB_TRANSFER_CANCELLED => TransferStatus::Cancelled,
+            LIBUSB_TRANSFER_STALL => TransferStatus::Stall,
+            LIBUSB_TRANSFER_NO_DEVICE => TransferStatus::NoDevice,
+            LIBUSB_TRANSFER_OVERFLOW => TransferStatus::Overflow,
+            _ => TransferStatus::Error,
+        }
+    }
+}
+
+/// Trait for usb transfer buffer.
+pub trait UsbTransferBuffer: Send {
+    fn as_ptr(&mut self) -> *mut u8;
+    fn len(&self) -> i32;
+}
+
+/// Default buffer size for control data transfer.
+const CONTROL_DATA_BUFFER_SIZE: usize = 1024;
+
+/// Buffer type for control transfer. The first 8-bytes is a UsbRequestSetup struct.
+#[repr(C, packed)]
+pub struct ControlTransferBuffer {
+    pub setup_buffer: UsbRequestSetup,
+    pub data_buffer: [u8; CONTROL_DATA_BUFFER_SIZE],
+}
+
+impl ControlTransferBuffer {
+    fn new() -> ControlTransferBuffer {
+        ControlTransferBuffer {
+            setup_buffer: UsbRequestSetup {
+                request_type: 0,
+                request: 0,
+                value: 0,
+                index: 0,
+                length: 0,
+            },
+            data_buffer: [0; CONTROL_DATA_BUFFER_SIZE],
+        }
+    }
+
+    /// Set request setup for this control buffer.
+    pub fn set_request_setup(&mut self, request_setup: &UsbRequestSetup) {
+        self.setup_buffer = request_setup.clone();
+    }
+}
+
+impl UsbTransferBuffer for ControlTransferBuffer {
+    fn as_ptr(&mut self) -> *mut u8 {
+        self as *mut ControlTransferBuffer as *mut u8
+    }
+
+    fn len(&self) -> i32 {
+        if self.setup_buffer.length as usize > CONTROL_DATA_BUFFER_SIZE {
+            panic!("Setup packet has an oversize length");
+        }
+        self.setup_buffer.length as i32 + size_of::<UsbRequestSetup>() as i32
+    }
+}
+
+/// Buffer type for Bulk transfer.
+pub struct BulkTransferBuffer {
+    buffer: Vec<u8>,
+}
+
+impl BulkTransferBuffer {
+    fn with_size(buffer_size: usize) -> Self {
+        BulkTransferBuffer {
+            buffer: vec![0; buffer_size],
+        }
+    }
+
+    /// Get mutable interal slice of this buffer.
+    pub fn as_mut_slice(&mut self) -> &mut [u8] {
+        &mut self.buffer
+    }
+
+    /// Get interal slice of this buffer.
+    pub fn as_slice(&self) -> &[u8] {
+        &self.buffer
+    }
+}
+
+impl UsbTransferBuffer for BulkTransferBuffer {
+    fn as_ptr(&mut self) -> *mut u8 {
+        if self.buffer.len() == 0 {
+            // Vec::as_mut_ptr() won't give 0x0 even if len() is 0.
+            std::ptr::null_mut()
+        } else {
+            self.buffer.as_mut_ptr()
+        }
+    }
+
+    fn len(&self) -> i32 {
+        self.buffer.len() as i32
+    }
+}
+
+type UsbTransferCompletionCallback<T> = Fn(UsbTransfer<T>) + Send + 'static;
+
+// This wraps libusb_transfer pointer.
+struct LibUsbTransfer {
+    ptr: *mut libusb_transfer,
+}
+
+impl Drop for LibUsbTransfer {
+    fn drop(&mut self) {
+        // Safe because 'self.ptr' is allocated by libusb_alloc_transfer.
+        unsafe {
+            libusb_free_transfer(self.ptr);
+        }
+    }
+}
+
+// It is safe to invoke libusb functions from multiple threads.
+// We cannot modify libusb_transfer safely from multiple threads. All the modifications happens
+// in construct (UsbTransfer::new) or consume (UsbTransfer::into_raw), we can consider this thread
+// safe.
+unsafe impl Send for LibUsbTransfer {}
+unsafe impl Sync for LibUsbTransfer {}
+
+/// TransferCanceller can cancel the transfer.
+pub struct TransferCanceller {
+    transfer: Weak<LibUsbTransfer>,
+}
+
+impl TransferCanceller {
+    /// Return false if fail to cancel.
+    pub fn try_cancel(&self) -> bool {
+        match self.transfer.upgrade() {
+            Some(t) => {
+                // Safe because self.transfer has ownership of the raw pointer.
+                let r = unsafe { libusb_cancel_transfer(t.ptr) };
+                if r == 0 {
+                    true
+                } else {
+                    false
+                }
+            }
+            None => false,
+        }
+    }
+}
+
+struct UsbTransferInner<T: UsbTransferBuffer> {
+    transfer: Arc<LibUsbTransfer>,
+    callback: Option<Box<UsbTransferCompletionCallback<T>>>,
+    buffer: T,
+}
+
+/// UsbTransfer owns a LibUsbTransfer, it's buffer and callback.
+pub struct UsbTransfer<T: UsbTransferBuffer> {
+    inner: Box<UsbTransferInner<T>>,
+}
+
+/// Build a control transfer.
+pub fn control_transfer(timeout: u32) -> UsbTransfer<ControlTransferBuffer> {
+    UsbTransfer::<ControlTransferBuffer>::new(
+        0,
+        LIBUSB_TRANSFER_TYPE_CONTROL as u8,
+        timeout,
+        ControlTransferBuffer::new(),
+    )
+}
+
+/// Build a data transfer.
+pub fn bulk_transfer(endpoint: u8, timeout: u32, size: usize) -> UsbTransfer<BulkTransferBuffer> {
+    UsbTransfer::<BulkTransferBuffer>::new(
+        endpoint,
+        LIBUSB_TRANSFER_TYPE_BULK as u8,
+        timeout,
+        BulkTransferBuffer::with_size(size),
+    )
+}
+
+/// Build a data transfer.
+pub fn interrupt_transfer(
+    endpoint: u8,
+    timeout: u32,
+    size: usize,
+) -> UsbTransfer<BulkTransferBuffer> {
+    UsbTransfer::<BulkTransferBuffer>::new(
+        endpoint,
+        LIBUSB_TRANSFER_TYPE_INTERRUPT as u8,
+        timeout,
+        BulkTransferBuffer::with_size(size),
+    )
+}
+
+impl<T: UsbTransferBuffer> UsbTransfer<T> {
+    fn new(endpoint: u8, type_: u8, timeout: u32, buffer: T) -> Self {
+        // Safe because alloc is safe.
+        let transfer: *mut libusb_transfer = unsafe { libusb_alloc_transfer(0) };
+        // Just panic on OOM.
+        assert!(!transfer.is_null());
+        let inner = Box::new(UsbTransferInner {
+            transfer: Arc::new(LibUsbTransfer { ptr: transfer }),
+            callback: None,
+            buffer,
+        });
+        // Safe because we inited transfer.
+        let raw_transfer: &mut libusb_transfer = unsafe { &mut *(inner.transfer.ptr) };
+        raw_transfer.endpoint = endpoint;
+        raw_transfer.type_ = type_;
+        raw_transfer.timeout = timeout;
+        raw_transfer.callback = Some(UsbTransfer::<T>::on_transfer_completed);
+        UsbTransfer { inner }
+    }
+
+    /// Get canceller of this transfer.
+    pub fn get_canceller(&self) -> TransferCanceller {
+        let weak_transfer = Arc::downgrade(&self.inner.transfer);
+        TransferCanceller {
+            transfer: weak_transfer,
+        }
+    }
+
+    /// Set callback function for transfer completion.
+    pub fn set_callback<C: 'static + Fn(UsbTransfer<T>) + Send>(&mut self, cb: C) {
+        self.inner.callback = Some(Box::new(cb));
+    }
+
+    /// Get a reference to the buffer.
+    pub fn buffer(&self) -> &T {
+        &self.inner.buffer
+    }
+
+    /// Get a mutable reference to the buffer.
+    pub fn buffer_mut(&mut self) -> &mut T {
+        &mut self.inner.buffer
+    }
+
+    /// Get actual length of data that was transferred.
+    pub fn actual_length(&self) -> i32 {
+        let transfer = self.inner.transfer.ptr;
+        // Safe because inner.ptr is always allocated by libusb_alloc_transfer.
+        unsafe { (*transfer).actual_length }
+    }
+
+    /// Get the transfer status of this transfer.
+    pub fn status(&self) -> TransferStatus {
+        let transfer = self.inner.transfer.ptr;
+        // Safe because inner.ptr is always allocated by libusb_alloc_transfer.
+        unsafe { TransferStatus::from((*transfer).status) }
+    }
+
+    /// Submit this transfer to device handle. 'self' is consumed. On success, the memory will be
+    /// 'leaked' (and stored in user_data) and sent to libusb, when the async operation is done,
+    /// on_transfer_completed will recreate 'self' and deliver it to callback/free 'self'. On
+    /// faliure, 'self' is returned with an error.
+    ///
+    /// # Safety
+    ///
+    /// Assumes libusb_device_handle is an handled opened by libusb, self.inner.transfer.ptr is
+    /// initialized with correct buffer and length.
+    pub unsafe fn submit(self, handle: *mut libusb_device_handle) -> Result<()> {
+        let transfer = self.into_raw();
+        (*transfer).dev_handle = handle;
+        match Error::from(libusb_submit_transfer(transfer)) {
+            Error::Success(_e) => Ok(()),
+            err => {
+                UsbTransfer::<T>::from_raw(transfer);
+                Err(err)
+            }
+        }
+    }
+
+    /// Invoke callback when transfer is completed.
+    ///
+    /// # Safety
+    ///
+    /// Assumes libusb_tranfser is finished. This function is called by libusb, don't call it
+    /// manually.
+    unsafe extern "C" fn on_transfer_completed(transfer: *mut libusb_transfer) {
+        let mut transfer = UsbTransfer::<T>::from_raw(transfer);
+        // Callback is reset to None.
+        if let Some(cb) = transfer.inner.callback.take() {
+            cb(transfer);
+        }
+    }
+
+    fn into_raw(mut self) -> *mut libusb_transfer {
+        let transfer: *mut libusb_transfer = self.inner.transfer.ptr;
+        // Safe because transfer is allocated by libusb_alloc_transfer.
+        unsafe {
+            (*transfer).buffer = self.buffer_mut().as_ptr();
+            (*transfer).length = self.buffer_mut().len();
+            (*transfer).user_data = Box::into_raw(self.inner) as *mut c_void;
+        }
+        transfer
+    }
+
+    unsafe fn from_raw(transfer: *mut libusb_transfer) -> Self {
+        UsbTransfer {
+            inner: Box::<UsbTransferInner<T>>::from_raw(
+                (*transfer).user_data as *mut UsbTransferInner<T>,
+            ),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::Mutex;
+
+    pub fn fake_submit_transfer<T: UsbTransferBuffer>(transfer: UsbTransfer<T>) {
+        let transfer = transfer.into_raw();
+        unsafe {
+            match (*transfer).callback {
+                Some(cb) => cb(transfer),
+                // Although no callback is invoked, we still need on_transfer_completed to
+                // free memory.
+                None => panic!("Memory leak!"),
+            };
+        }
+    }
+
+    #[test]
+    fn check_control_buffer_size() {
+        assert_eq!(
+            size_of::<ControlTransferBuffer>(),
+            size_of::<UsbRequestSetup>() + CONTROL_DATA_BUFFER_SIZE
+        );
+    }
+
+    #[test]
+    fn submit_transfer_no_callback_test() {
+        let t = control_transfer(0);
+        fake_submit_transfer(t);
+        let t = bulk_transfer(0, 0, 1);
+        fake_submit_transfer(t);
+    }
+
+    struct FakeTransferController {
+        data: Mutex<u8>,
+    }
+
+    #[test]
+    fn submit_transfer_with_callback() {
+        let c = Arc::new(FakeTransferController {
+            data: Mutex::new(0),
+        });
+        let c1 = Arc::downgrade(&c);
+        let mut t = control_transfer(0);
+        t.set_callback(move |_t| {
+            let c = c1.upgrade().unwrap();
+            *c.data.lock().unwrap() = 3;
+        });
+        fake_submit_transfer(t);
+        assert_eq!(*c.data.lock().unwrap(), 3);
+    }
+}