Support Fastboot over USB

Implement Fastboot over USB using EFI_ANDROID_BOOT_PROTOCOL. TCP & USB
are supported concurrently. GBL keeps polling TCP USB interfaces and
process Fastboot command based on the following rule:

1. If a USB packet arrives, GBL runs Fastboot over USB exclusively
   until the command in the paket is done. In the case of downloads,
   it will process until the download is complete. It then polls TCP
   connection.
2. If a TCP connection is established, GBL runs Fastboot TCP
   exclusively until connection is lost. Then returns to 1.

Bug: 333406178
Change-Id: I522733340557e86dbf436e7791b91813e23398fd

8 files changed, 339 insertions, 146 deletions

diff --git a/gbl/docs/EFI_ANDROID_BOOT_PROTOCOL.md b/gbl/docs/EFI_ANDROID_BOOT_PROTOCOL.md
index 6f4a08b..acdf90e 100644
--- a/gbl/docs/EFI_ANDROID_BOOT_PROTOCOL.md
+++ b/gbl/docs/EFI_ANDROID_BOOT_PROTOCOL.md
@@ -197,9 +197,9 @@ Receives a USB packet from the interface started by
 typedef
 EFI_STATUS
 (EFIAPI * EFI_ANDROID_BOOT_FASTBOOT_USB_RECEIVE)(
-  IN EFI_ANDROID_BOOT_PROTOCOL *This,
-  IN OUT UINTN        *BufferSize,
-  OUT VOID            *Buffer,
+  IN EFI_ANDROID_BOOT_PROTOCOL  *This,
+  IN OUT UINTN                  *BufferSize,
+  OUT VOID                      *Buffer,
   );
 ```
 
@@ -245,9 +245,9 @@ Sends a USB packet from the USB interface started by
 typedef
 EFI_STATUS
 (EFIAPI * EFI_ANDROID_BOOT_FASTBOOT_USB_SEND)(
-  IN EFI_ANDROID_BOOT_PROTOCOL *This,
-  IN OUT UINTN        *BufferSize,
-  IN VOID             *Buffer,
+  IN EFI_ANDROID_BOOT_PROTOCOL  *This,
+  IN OUT UINTN                  *BufferSize,
+  IN CONST VOID                 *Buffer,
   );
 ```
 
diff --git a/gbl/efi/src/fastboot.rs b/gbl/efi/src/fastboot.rs
index e2359af..0040e4b 100644
--- a/gbl/efi/src/fastboot.rs
+++ b/gbl/efi/src/fastboot.rs
@@ -17,96 +17,215 @@
 // `android_boot:android_boot_demo()` and `fuchsia_boot:fuchsia_boot_demo()` for
 // supported/unsupported features at the moment.
 
-use crate::error::Result;
+use crate::error::{EfiAppError, GblEfiError, Result as GblResult};
 use crate::net::{with_efi_network, EfiTcpSocket};
-use crate::utils::find_gpt_devices;
-use core::fmt::Write;
-use efi::{efi_print, efi_println, EfiEntry};
-use fastboot::{Fastboot, TcpStream, TransportError};
+use crate::utils::{find_gpt_devices, loop_with_timeout};
+use core::{fmt::Write, result::Result};
+use efi::{
+    defs::{EFI_STATUS_NOT_READY, EFI_STATUS_NOT_STARTED},
+    efi_print, efi_println,
+    protocol::{android_boot::AndroidBootProtocol, Protocol},
+    EfiEntry,
+};
+use fastboot::{Fastboot, TcpStream, Transport, TransportError};
 use libgbl::fastboot::GblFastboot;
 
 const DEFAULT_TIMEOUT_MS: u64 = 5_000;
 const FASTBOOT_TCP_PORT: u16 = 5554;
 
 struct EfiFastbootTcpTransport<'a, 'b, 'c> {
-    transport_error: &'c mut Result<()>,
+    last_err: GblResult<()>,
     socket: &'c mut EfiTcpSocket<'a, 'b>,
 }
 
+impl<'a, 'b, 'c> EfiFastbootTcpTransport<'a, 'b, 'c> {
+    fn new(socket: &'c mut EfiTcpSocket<'a, 'b>) -> Self {
+        Self { last_err: Ok(()), socket: socket }
+    }
+}
+
 impl TcpStream for EfiFastbootTcpTransport<'_, '_, '_> {
     /// Reads to `out` for exactly `out.len()` number bytes from the TCP connection.
-    fn read_exact(&mut self, out: &mut [u8]) -> core::result::Result<(), TransportError> {
-        *self.transport_error = self.socket.receive_exact(out, DEFAULT_TIMEOUT_MS);
-        self.transport_error.as_ref().map_err(|_| TransportError::Others(""))?;
+    fn read_exact(&mut self, out: &mut [u8]) -> Result<(), TransportError> {
+        self.last_err = self.socket.receive_exact(out, DEFAULT_TIMEOUT_MS);
+        self.last_err.as_ref().map_err(|_| TransportError::Others("TCP read error"))?;
         Ok(())
     }
 
     /// Sends exactly `data.len()` number bytes from `data` to the TCP connection.
-    fn write_exact(&mut self, data: &[u8]) -> core::result::Result<(), TransportError> {
-        *self.transport_error = self.socket.send_exact(data, DEFAULT_TIMEOUT_MS);
-        self.transport_error.as_ref().map_err(|_| TransportError::Others(""))?;
+    fn write_exact(&mut self, data: &[u8]) -> Result<(), TransportError> {
+        self.last_err = self.socket.send_exact(data, DEFAULT_TIMEOUT_MS);
+        self.last_err.as_ref().map_err(|_| TransportError::Others("TCP write error"))?;
         Ok(())
     }
 }
 
-/// Internal helper for performing Fastboot over TCP.
-fn fastboot_tcp_usb(
-    socket: &mut EfiTcpSocket,
+/// `UsbTransport` implements the `fastboot::Transport` trait using USB interfaces from
+/// EFI_ANDROID_BOOT_PROTOCOL.
+pub struct UsbTransport<'a, 'b> {
+    last_err: GblResult<()>,
+    max_packet_size: usize,
+    protocol: &'b Protocol<'a, AndroidBootProtocol>,
+}
+
+impl<'a, 'b> UsbTransport<'a, 'b> {
+    fn new(max_packet_size: usize, protocol: &'b Protocol<'a, AndroidBootProtocol>) -> Self {
+        Self { last_err: Ok(()), max_packet_size: max_packet_size, protocol: protocol }
+    }
+
+    /// Waits for the previous send to complete up to `DEFAULT_TIMEOUT_MS` timeout.
+    fn wait_for_send(&self) -> GblResult<()> {
+        loop_with_timeout(self.protocol.efi_entry(), DEFAULT_TIMEOUT_MS, || {
+            match (|| -> GblResult<bool> {
+                Ok(self
+                    .protocol
+                    .efi_entry()
+                    .system_table()
+                    .boot_services()
+                    .check_event(&self.protocol.wait_for_send_completion()?)?)
+            })() {
+                Ok(true) => Ok(Ok(())),
+                Ok(false) => Err(false),
+                Err(e) => Ok(Err(e)),
+            }
+        })?
+        .ok_or(EfiAppError::Timeout)??;
+        Ok(())
+    }
+}
+
+impl Transport for UsbTransport<'_, '_> {
+    fn receive_packet(&mut self, out: &mut [u8]) -> Result<usize, TransportError> {
+        let mut out_size = 0;
+        self.last_err = Ok(());
+        match self.protocol.fastboot_usb_receive(out, &mut out_size) {
+            Ok(()) => Ok(out_size),
+            Err(e) if e.is_efi_err(EFI_STATUS_NOT_READY) => Ok(0),
+            Err(e) => {
+                self.last_err = Err(e.into());
+                Err(TransportError::Others("USB receive error"))
+            }
+        }
+    }
+
+    fn send_packet(&mut self, packet: &[u8]) -> Result<(), TransportError> {
+        let mut sent = 0;
+        self.last_err = (|| -> GblResult<()> {
+            while sent < packet.len() {
+                let to_send = core::cmp::min(packet.len() - sent, self.max_packet_size);
+                let mut out_size = 0;
+                self.protocol.fastboot_usb_send(&packet[sent..][..to_send], &mut out_size)?;
+                self.wait_for_send()?;
+                sent += to_send;
+            }
+            Ok(())
+        })();
+        Ok(*self.last_err.as_ref().map_err(|_| TransportError::Others("USB send error"))?)
+    }
+}
+
+/// Loops and polls both USB and TCP transport. Runs Fastboot if any is available.
+fn fastboot_loop(
     efi_entry: &EfiEntry,
-    download_buffer: &mut [u8],
-) -> Result<()> {
-    let mut gpt_devices = find_gpt_devices(efi_entry)?;
-    let mut gbl_fastboot = GblFastboot::new(&mut gpt_devices);
-    efi_println!(efi_entry, "Listening for Fastboot over TCP...");
-    efi_println!(efi_entry, "IP addresses:");
-    socket.interface().ip_addrs().iter().for_each(|v| {
-        efi_println!(efi_entry, "\t{}", v.address());
-    });
+    gbl_fb: &mut GblFastboot,
+    fastboot: &mut Fastboot,
+    mut socket: Option<&mut EfiTcpSocket>,
+    mut usb: Option<&mut UsbTransport>,
+) -> GblResult<()> {
+    if socket.is_none() && usb.is_none() {
+        return Err(EfiAppError::Unsupported.into());
+    }
+
+    efi_println!(efi_entry, "Fastboot USB: {}", usb.as_ref().map_or("No", |_| "Yes"));
+    if let Some(socket) = socket.as_ref() {
+        efi_println!(efi_entry, "Fastboot TCP: Yes");
+        efi_println!(efi_entry, "Device IP addresses:");
+        socket.interface().ip_addrs().iter().for_each(|v| {
+            efi_println!(efi_entry, "\t{}", v.address());
+        });
+    } else {
+        efi_println!(efi_entry, "Fastboot TCP: No");
+    }
+
+    let mut listen_start_timestamp = EfiTcpSocket::timestamp(0);
     loop {
-        socket.listen(FASTBOOT_TCP_PORT)?;
-        let listen_start = EfiTcpSocket::timestamp(0);
-        loop {
+        // Checks and processes commands over USB.
+        if let Some(usb) = usb.as_mut() {
+            if fastboot.process_next_command(*usb, gbl_fb).is_err() {
+                efi_println!(efi_entry, "Fastboot USB error: {:?}", usb.last_err);
+            }
+        }
+
+        // Checks and processes commands over TCP.
+        if let Some(socket) = socket.as_mut() {
             socket.poll();
+            let mut reset_socket = false;
             if socket.check_active() {
-                // Has connection.
-                efi_println!(
-                    efi_entry,
-                    "Connection from {}",
-                    socket.get_socket().remote_endpoint().unwrap()
-                );
-
-                let mut transport_error = Ok(());
-                let mut transport = EfiFastbootTcpTransport {
-                    transport_error: &mut transport_error,
-                    socket: socket,
-                };
-                let mut fastboot = Fastboot::new(&mut download_buffer[..]);
-                let _ = fastboot.run_tcp_session(&mut transport, &mut gbl_fastboot);
-                efi_println!(efi_entry, "Fastboot TCP session ends. {:?}", transport_error);
-                break;
-            } else if EfiTcpSocket::timestamp(listen_start) > DEFAULT_TIMEOUT_MS {
+                let remote = socket.get_socket().remote_endpoint().unwrap();
+                efi_println!(efi_entry, "TCP connection from {}", remote);
+                let mut transport = EfiFastbootTcpTransport::new(socket);
+                let _ = fastboot.run_tcp_session(&mut transport, gbl_fb);
+                match transport.last_err {
+                    Ok(()) | Err(GblEfiError::EfiAppError(EfiAppError::PeerClosed)) => {}
+                    Err(e) => {
+                        efi_println!(efi_entry, "Fastboot TCP error {:?}", e);
+                    }
+                }
+                reset_socket = true;
+            } else if EfiTcpSocket::timestamp(listen_start_timestamp) > DEFAULT_TIMEOUT_MS {
                 // Reset once in a while in case a remote client disconnects in the middle of
                 // TCP handshake and leaves the socket in a half open state.
-                break;
+                reset_socket = true;
             }
 
-            // Perform Fastboot over USB here.
+            if reset_socket {
+                listen_start_timestamp = EfiTcpSocket::timestamp(0);
+                if let Err(e) = socket.listen(FASTBOOT_TCP_PORT) {
+                    efi_println!(efi_entry, "TCP listen error: {:?}", e);
+                }
+            }
         }
     }
 }
 
+/// Initializes the Fastboot USB interface and returns a `UsbTransport`.
+fn init_usb<'a, 'b>(
+    android_boot_protocol: &Option<&'b Protocol<'a, AndroidBootProtocol>>,
+) -> GblResult<UsbTransport<'a, 'b>> {
+    let protocol = android_boot_protocol.ok_or(EfiAppError::Unsupported)?;
+    match protocol.fastboot_usb_interface_stop() {
+        Err(e) if !e.is_efi_err(EFI_STATUS_NOT_STARTED) => return Err(e.into()),
+        _ => {}
+    };
+    Ok(UsbTransport::new(protocol.fastboot_usb_interface_start()?, protocol))
+}
+
 /// Runs Fastboot.
-pub fn run_fastboot(efi_entry: &EfiEntry) -> Result<()> {
+pub fn run_fastboot(
+    efi_entry: &EfiEntry,
+    android_boot_protocol: Option<&Protocol<'_, AndroidBootProtocol>>,
+) -> GblResult<()> {
+    let mut gpt_devices = find_gpt_devices(efi_entry)?;
+    let mut gbl_fb = GblFastboot::new(&mut gpt_devices);
     // TODO(b/328786603): Figure out where to get download buffer size.
     let mut download_buffer = vec![0u8; 512 * 1024 * 1024];
-    match with_efi_network(efi_entry, |socket| -> Result<()> {
-        fastboot_tcp_usb(socket, efi_entry, &mut download_buffer[..])
-    })? {
+    let mut fastboot = Fastboot::new(&mut download_buffer[..]);
+
+    let mut usb = match init_usb(&android_boot_protocol) {
+        Ok(v) => Some(v),
         Err(e) => {
-            efi_println!(efi_entry, "Failed to initilaize network {:?}", e);
-            // Perform Fastboot over USB here.
+            efi_println!(efi_entry, "Failed to start Fastboot over USB. {:?}.", e);
+            None
         }
-        _ => {}
     };
-    Ok(())
+
+    match with_efi_network(efi_entry, |socket| -> GblResult<()> {
+        fastboot_loop(efi_entry, &mut gbl_fb, &mut fastboot, Some(socket), usb.as_mut())
+    }) {
+        Err(e) => {
+            efi_println!(efi_entry, "Failed to start EFI network. {:?}.", e);
+            fastboot_loop(efi_entry, &mut gbl_fb, &mut fastboot, None, usb.as_mut())
+        }
+        v => v?,
+    }
 }
diff --git a/gbl/efi/src/main.rs b/gbl/efi/src/main.rs
index b1c3247..e06a9b6 100644
--- a/gbl/efi/src/main.rs
+++ b/gbl/efi/src/main.rs
@@ -26,6 +26,7 @@ extern crate alloc;
 use core::fmt::Write;
 
 use efi::defs::EfiSystemTable;
+use efi::protocol::android_boot::AndroidBootProtocol;
 use efi::{efi_print, efi_println, initialize};
 
 #[macro_use]
@@ -56,7 +57,12 @@ fn main(image_handle: *mut core::ffi::c_void, systab_ptr: *mut EfiSystemTable) -
     match wait_key_stroke(&entry, 'f', 2000) {
         Ok(true) => {
             efi_println!(entry, "'f' pressed.");
-            fastboot::run_fastboot(&entry)?;
+            let android_boot_protocol = entry
+                .system_table()
+                .boot_services()
+                .find_first_and_open::<AndroidBootProtocol>()
+                .ok();
+            fastboot::run_fastboot(&entry, android_boot_protocol.as_ref())?;
         }
         _ => {}
     }
diff --git a/gbl/libefi/defs/protocols/android_boot_protocol.h b/gbl/libefi/defs/protocols/android_boot_protocol.h
index 3d4c115..14c727b 100644
--- a/gbl/libefi/defs/protocols/android_boot_protocol.h
+++ b/gbl/libefi/defs/protocols/android_boot_protocol.h
@@ -31,7 +31,7 @@ typedef struct EfiAndroidBootProtocol {
   EfiStatus (*fastboot_usb_receive)(struct EfiAndroidBootProtocol* self,
                                     size_t* buffer_size, void* buffer);
   EfiStatus (*fastboot_usb_send)(struct EfiAndroidBootProtocol* self,
-                                 size_t* buffer_size, void* buffer);
+                                 size_t* buffer_size, const void* buffer);
   EfiEvent wait_for_send_completion;
 } EfiAndroidBootProtocol;
 
diff --git a/gbl/libefi/src/lib.rs b/gbl/libefi/src/lib.rs
index c7889d2..f2d0e84 100644
--- a/gbl/libefi/src/lib.rs
+++ b/gbl/libefi/src/lib.rs
@@ -428,11 +428,11 @@ impl<'a> BootServices<'a> {
                 &mut efi_event
             )?;
         }
-        Ok(Event {
-            efi_entry: self.efi_entry,
-            efi_event: efi_event,
-            _cb: cb.map::<&'n mut dyn FnMut(EfiEvent), _>(|v| v.cb),
-        })
+        Ok(Event::new(
+            Some(self.efi_entry),
+            efi_event,
+            cb.map::<&'n mut dyn FnMut(EfiEvent), _>(|v| v.cb),
+        ))
     }
 
     /// Wrapper of `EFI_BOOT_SERVICE.CloseEvent()`.
@@ -508,14 +508,33 @@ impl<'e> EventNotify<'e> {
 /// `Event` wraps the raw `EfiEvent` handle and internally enforces a borrow of the registered
 /// callback for the given life time `e. The event is automatically closed when going out of scope.
 pub struct Event<'a, 'n> {
-    efi_entry: &'a EfiEntry,
+    // If `efi_entry` is None, it represents an unowned Event and won't get closed on drop.
+    efi_entry: Option<&'a EfiEntry>,
     efi_event: EfiEvent,
     _cb: Option<&'n mut dyn FnMut(EfiEvent)>,
 }
 
+impl<'a, 'n> Event<'a, 'n> {
+    /// Creates an instance of owned `Event`. The `Event` is closed when going out of scope.
+    fn new(
+        efi_entry: Option<&'a EfiEntry>,
+        efi_event: EfiEvent,
+        _cb: Option<&'n mut dyn FnMut(EfiEvent)>,
+    ) -> Self {
+        Self { efi_entry, efi_event, _cb }
+    }
+
+    /// Creates an  unowned `Event`. The `Event` is not closed when going out of scope.
+    fn new_unowned(efi_event: EfiEvent) -> Self {
+        Self { efi_entry: None, efi_event: efi_event, _cb: None }
+    }
+}
+
 impl Drop for Event<'_, '_> {
     fn drop(&mut self) {
-        self.efi_entry.system_table().boot_services().close_event(self).unwrap();
+        if let Some(efi_entry) = self.efi_entry {
+            efi_entry.system_table().boot_services().close_event(self).unwrap();
+        }
     }
 }
 
diff --git a/gbl/libefi/src/protocol.rs b/gbl/libefi/src/protocol.rs
index a6092b5..309d4a7 100644
--- a/gbl/libefi/src/protocol.rs
+++ b/gbl/libefi/src/protocol.rs
@@ -70,6 +70,11 @@ impl<'a, T: ProtocolInfo> Protocol<'a, T> {
         unsafe { self.interface.as_ref() }.ok_or_else(|| EFI_STATUS_INVALID_PARAMETER.into())
     }
 
+    /// Returns the reference to EFI entry.
+    pub fn efi_entry(&self) -> &'a EfiEntry {
+        self.efi_entry
+    }
+
     /// Returns the mutable pointer of the interface. Invisible from outside. Application should
     /// not have any need to alter the content of interface data.
     pub(crate) fn interface_ptr(&self) -> *mut T::InterfaceType {
diff --git a/gbl/libefi/src/protocol/android_boot.rs b/gbl/libefi/src/protocol/android_boot.rs
index f7d9ec8..931d6e6 100644
--- a/gbl/libefi/src/protocol/android_boot.rs
+++ b/gbl/libefi/src/protocol/android_boot.rs
@@ -14,7 +14,7 @@
 
 use crate::defs::{EfiAndroidBootProtocol, EfiGuid, EFI_STATUS_NOT_FOUND};
 use crate::protocol::{Protocol, ProtocolInfo};
-use crate::{efi_call, map_efi_err, EfiResult};
+use crate::{efi_call, map_efi_err, EfiResult, Event};
 
 /// EFI_ANDROID_BOOT_PROTOCOL
 pub struct AndroidBootProtocol;
@@ -74,7 +74,7 @@ impl Protocol<'_, AndroidBootProtocol> {
     }
 
     /// Wrapper of `EFI_ANDROID_BOOT_PROTOCOL.fastboot_usb_send()`
-    pub fn fastboot_usb_send(&self, data: &mut [u8], out_size: &mut usize) -> EfiResult<()> {
+    pub fn fastboot_usb_send(&self, data: &[u8], out_size: &mut usize) -> EfiResult<()> {
         *out_size = data.len();
         // SAFETY:
         // `self.interface()?` guarantees self.interface is non-null and points to a valid object
@@ -86,8 +86,13 @@ impl Protocol<'_, AndroidBootProtocol> {
                 self.interface()?.fastboot_usb_send,
                 self.interface,
                 out_size,
-                data.as_mut_ptr() as _,
+                data.as_ptr() as _,
             )
         }
     }
+
+    /// Returns the `EFI_ANDROID_BOOT_PROTOCOL.wait_for_send_completion` EFI event.
+    pub fn wait_for_send_completion(&self) -> EfiResult<Event> {
+        Ok(Event::new_unowned(self.interface()?.wait_for_send_completion))
+    }
 }
diff --git a/gbl/libfastboot/src/lib.rs b/gbl/libfastboot/src/lib.rs
index 2254154..4b5f18b 100644
--- a/gbl/libfastboot/src/lib.rs
+++ b/gbl/libfastboot/src/lib.rs
@@ -77,6 +77,7 @@ const OKAY: &'static str = "OKAY";
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
 pub enum TransportError {
     InvalidHanshake,
+    InvalidState,
     PacketSizeOverflow,
     PacketSizeExceedMaximum,
     NotEnoughUpload,
@@ -123,29 +124,45 @@ pub trait TcpStream {
     fn write_exact(&mut self, data: &[u8]) -> Result<(), TransportError>;
 }
 
-impl Transport for &mut dyn TcpStream {
+/// `TcpTransport` implements `Transport` with a `TcpStream` object.
+pub struct TcpTransport<'a>(&'a mut dyn TcpStream);
+
+impl<'a> TcpTransport<'a> {
+    /// Creates an instance from a newly connected TcpStream and performs handshake.
+    pub fn new_and_handshake(tcp_stream: &'a mut dyn TcpStream) -> Result<Self, TransportError> {
+        let mut handshake = [0u8; 4];
+        tcp_stream.write_exact(TCP_HANDSHAKE_MESSAGE)?;
+        tcp_stream.read_exact(&mut handshake[..])?;
+        match handshake == *TCP_HANDSHAKE_MESSAGE {
+            true => Ok(Self(tcp_stream)),
+            _ => Err(TransportError::InvalidHanshake),
+        }
+    }
+}
+
+impl Transport for TcpTransport<'_> {
     fn receive_packet(&mut self, out: &mut [u8]) -> Result<usize, TransportError> {
         let mut length_prefix = [0u8; 8];
-        self.read_exact(&mut length_prefix[..])?;
+        self.0.read_exact(&mut length_prefix[..])?;
         let packet_size: usize = u64::from_be_bytes(length_prefix)
             .try_into()
             .map_err(|_| TransportError::PacketSizeOverflow)?;
         match out.len() < packet_size {
             true => Err(TransportError::PacketSizeExceedMaximum),
             _ => {
-                self.read_exact(&mut out[..packet_size])?;
+                self.0.read_exact(&mut out[..packet_size])?;
                 Ok(packet_size)
             }
         }
     }
 
     fn send_packet(&mut self, packet: &[u8]) -> Result<(), TransportError> {
-        self.write_exact(
+        self.0.write_exact(
             &mut u64::try_from(packet.len())
                 .map_err(|_| TransportError::PacketSizeOverflow)?
                 .to_be_bytes()[..],
         )?;
-        self.write_exact(packet)
+        self.0.write_exact(packet)
     }
 }
 
@@ -580,85 +597,113 @@ impl<'a> Fastboot<'a> {
         }
     }
 
-    /// Process fastboot command/data from a given transport.
+    /// Processes the next fastboot packet from a transport.
     ///
     /// # Args
     ///
-    ///   * `transport`: An implementation of `Transport`
-    ///   * `fb_impl`: An implementation of `FastbootImplementation`.
+    /// * `transport`: An implementation of `Transport`
+    /// * `fb_impl`: An implementation of `FastbootImplementation`.
     ///
     /// # Returns
     ///
-    ///   The method returns on any errors from calls to `transport` methods.
-    pub fn run(
+    /// Returns error if any calls to `transport` methods return error.
+    /// Returns Ok(()) if transport doesn't have any next packet.
+    pub fn process_next_packet(
         &mut self,
         transport: &mut impl Transport,
         fb_impl: &mut impl FastbootImplementation,
     ) -> Result<(), TransportError> {
-        loop {
-            match self.state {
-                ProtocolState::Command => {
-                    let mut packet = [0u8; MAX_COMMAND_SIZE];
-                    let cmd_size = transport.receive_packet(&mut packet[..])?;
-                    if cmd_size == 0 {
-                        continue;
-                    }
+        match self.state {
+            ProtocolState::Command => {
+                let mut packet = [0u8; MAX_COMMAND_SIZE];
+                let cmd_size = transport.receive_packet(&mut packet[..])?;
+                if cmd_size == 0 {
+                    return Ok(());
+                }
 
-                    let mut res = [0u8; MAX_RESPONSE_SIZE];
-                    let cmd_str = match from_utf8(&packet[..cmd_size]) {
-                        Ok(s) => s,
-                        _ => {
-                            transport.send_packet(fastboot_fail!(res, "Invalid Command"))?;
-                            continue;
-                        }
-                    };
-                    let mut args = cmd_str.split(':');
-                    let Some(cmd) = args.next() else {
-                        transport.send_packet(fastboot_fail!(res, "No command"))?;
-                        continue;
-                    };
-                    match cmd {
-                        "getvar" => self.get_var(args, transport, fb_impl)?,
-                        "download" => self.download(args, transport, fb_impl)?,
-                        "flash" => self.flash(cmd_str, transport, fb_impl)?,
-                        "upload" => self.upload(transport, fb_impl)?,
-                        "fetch" => self.fetch(&cmd_str, args, transport, fb_impl)?,
-                        _ if cmd_str.starts_with("oem ") => {
-                            self.oem(&cmd_str[4..], transport, fb_impl)?
-                        }
-                        _ => {
-                            transport.send_packet(fastboot_fail!(res, "Command not found"))?;
-                        }
+                let mut res = [0u8; MAX_RESPONSE_SIZE];
+                let cmd_str = match from_utf8(&packet[..cmd_size]) {
+                    Ok(s) => s,
+                    _ => {
+                        return transport.send_packet(fastboot_fail!(res, "Invalid Command"));
+                    }
+                };
+                let mut args = cmd_str.split(':');
+                let Some(cmd) = args.next() else {
+                    return transport.send_packet(fastboot_fail!(res, "No command"));
+                };
+                match cmd {
+                    "getvar" => self.get_var(args, transport, fb_impl)?,
+                    "download" => self.download(args, transport, fb_impl)?,
+                    "flash" => self.flash(cmd_str, transport, fb_impl)?,
+                    "upload" => self.upload(transport, fb_impl)?,
+                    "fetch" => self.fetch(&cmd_str, args, transport, fb_impl)?,
+                    _ if cmd_str.starts_with("oem ") => {
+                        self.oem(&cmd_str[4..], transport, fb_impl)?;
+                    }
+                    _ => {
+                        return transport.send_packet(fastboot_fail!(res, "Command not found"));
                     }
                 }
-                ProtocolState::Download => {
-                    let (_, remains) = &mut self.download_buffer[..self.total_download_size]
-                        .split_at_mut(self.downloaded_size);
-                    match transport.receive_packet(remains) {
-                        Ok(size) if size > remains.len() => {
-                            let mut res = [0u8; MAX_RESPONSE_SIZE];
-                            transport.send_packet(
-                                snprintf!(res, "FAILMore data received then expected").as_bytes(),
-                            )?;
-                            self.total_download_size = 0;
-                            self.downloaded_size = 0;
+            }
+            ProtocolState::Download => {
+                let (_, remains) = &mut self.download_buffer[..self.total_download_size]
+                    .split_at_mut(self.downloaded_size);
+                match transport.receive_packet(remains) {
+                    Ok(size) if size > remains.len() => {
+                        let mut res = [0u8; MAX_RESPONSE_SIZE];
+                        transport.send_packet(
+                            snprintf!(res, "FAILMore data received then expected").as_bytes(),
+                        )?;
+                        self.total_download_size = 0;
+                        self.downloaded_size = 0;
+                        self.state = ProtocolState::Command;
+                    }
+                    Ok(size) => {
+                        self.downloaded_size = self.downloaded_size.checked_add(size).unwrap();
+                        if self.downloaded_size == self.total_download_size {
                             self.state = ProtocolState::Command;
+                            transport.send_packet(OKAY.as_bytes())?;
                         }
-                        Ok(size) => {
-                            self.downloaded_size = self.downloaded_size.checked_add(size).unwrap();
-                            if self.downloaded_size == self.total_download_size {
-                                self.state = ProtocolState::Command;
-                                transport.send_packet(OKAY.as_bytes())?;
-                            }
-                        }
-                        Err(e) => {
-                            self.total_download_size = 0;
-                            self.downloaded_size = 0;
-                            return Err(e);
-                        }
+                    }
+                    Err(e) => {
+                        self.total_download_size = 0;
+                        self.downloaded_size = 0;
+                        return Err(e);
                     }
                 }
             }
+        };
+        Ok(())
+    }
+
+    /// Fetches and processes the next fastboot command from the transport.
+    ///
+    /// Returns Ok(()) if transport doesn't have any next packet.
+    pub fn process_next_command(
+        &mut self,
+        transport: &mut impl Transport,
+        fb_impl: &mut impl FastbootImplementation,
+    ) -> Result<(), TransportError> {
+        if !matches!(self.state, ProtocolState::Command) {
+            return Err(TransportError::InvalidState);
+        }
+        self.process_next_packet(transport, fb_impl)?;
+        // Keep processing until it is back to the command state.
+        while !matches!(self.state, ProtocolState::Command) {
+            self.process_next_packet(transport, fb_impl)?;
+        }
+        Ok(())
+    }
+
+    /// Keeps polling and processing fastboot commands from the transport.
+    pub fn run(
+        &mut self,
+        transport: &mut impl Transport,
+        fb_impl: &mut impl FastbootImplementation,
+    ) -> Result<(), TransportError> {
+        loop {
+            self.process_next_command(transport, fb_impl)?;
         }
     }
 
@@ -667,16 +712,10 @@ impl<'a> Fastboot<'a> {
     /// The method performs fastboot over TCP handshake and then call `Self::run(...)`.
     pub fn run_tcp_session(
         &mut self,
-        mut tcp_stream: &mut dyn TcpStream,
+        tcp_stream: &mut dyn TcpStream,
         fb_impl: &mut impl FastbootImplementation,
     ) -> Result<(), TransportError> {
-        let mut handshake = [0u8; 4];
-        tcp_stream.write_exact(TCP_HANDSHAKE_MESSAGE)?;
-        tcp_stream.read_exact(&mut handshake[..])?;
-        match handshake == *TCP_HANDSHAKE_MESSAGE {
-            true => self.run(&mut tcp_stream, fb_impl),
-            _ => Err(TransportError::InvalidHanshake),
-        }
+        self.run(&mut TcpTransport::new_and_handshake(tcp_stream)?, fb_impl)
     }
 
     /// Method for handling "fastboot getvar ..."