feat: add framebuffer over USB/TCP direct devices

adb_cli/src/commands/mod.rs

@@ -7,5 +7,5 @@ mod usb;
 pub use emu::EmuCommand;
 pub use host::{HostCommand, MdnsCommand};
 pub use local::LocalCommand;
-pub use tcp::{TcpCommand, TcpCommands};
-pub use usb::{UsbCommand, UsbCommands};
+pub use tcp::TcpCommand;
+pub use usb::{DeviceCommands, UsbCommand};

adb_cli/src/commands/tcp.rs

@@ -1,44 +1,11 @@
-use std::net::SocketAddr;
-use std::path::PathBuf;
-
 use clap::Parser;
+use std::net::SocketAddr;
 
-use crate::models::RebootTypeCommand;
+use super::DeviceCommands;
 
 #[derive(Parser, Debug)]
 pub struct TcpCommand {
     pub address: SocketAddr,
     #[clap(subcommand)]
-    pub commands: TcpCommands,
-}
-
-#[derive(Parser, Debug)]
-pub enum TcpCommands {
-    /// Spawn an interactive shell or run a list of commands on the device
-    Shell { commands: Vec<String> },
-    /// Pull a file from device
-    Pull { source: String, destination: String },
-    /// Push a file on device
-    Push { filename: String, path: String },
-    /// Stat a file on device
-    Stat { path: String },
-    /// Run an activity on device specified by the intent
-    Run {
-        /// The package whose activity is to be invoked
-        #[clap(short = 'p', long = "package")]
-        package: String,
-        /// The activity to be invoked itself, Usually it is MainActivity
-        #[clap(short = 'a', long = "activity")]
-        activity: String,
-    },
-    /// Reboot the device
-    Reboot {
-        #[clap(subcommand)]
-        reboot_type: RebootTypeCommand,
-    },
-    /// Install an APK on device
-    Install {
-        /// Path to APK file. Extension must be ".apk"
-        path: PathBuf,
-    },
+    pub commands: DeviceCommands,
 }

adb_cli/src/commands/usb.rs

@@ -21,11 +21,11 @@ pub struct UsbCommand {
     #[clap(short = 'k', long = "private-key")]
     pub path_to_private_key: Option<PathBuf>,
     #[clap(subcommand)]
-    pub commands: UsbCommands,
+    pub commands: DeviceCommands,
 }
 
 #[derive(Parser, Debug)]
-pub enum UsbCommands {
+pub enum DeviceCommands {
     /// Spawn an interactive shell or run a list of commands on the device
     Shell { commands: Vec<String> },
     /// Pull a file from device
@@ -53,4 +53,9 @@ pub enum UsbCommands {
         /// Path to APK file. Extension must be ".apk"
         path: PathBuf,
     },
+    /// Dump framebuffer of device
+    Framebuffer { 
+        /// Framebuffer image destination path
+        path: String
+    },
 }

adb_cli/src/main.rs

@@ -12,7 +12,7 @@ use adb_client::{
 };
 use anyhow::{anyhow, Result};
 use clap::Parser;
-use commands::{EmuCommand, HostCommand, LocalCommand, MdnsCommand, TcpCommands, UsbCommands};
+use commands::{DeviceCommands, EmuCommand, HostCommand, LocalCommand, MdnsCommand};
 use models::{Command, Opts};
 use std::fs::File;
 use std::io::Write;
@@ -226,7 +226,7 @@ fn main() -> Result<()> {
             };
 
             match usb.commands {
-                UsbCommands::Shell { commands } => {
+                DeviceCommands::Shell { commands } => {
                     if commands.is_empty() {
                         // Need to duplicate some code here as ADBTermios [Drop] implementation resets terminal state.
                         // Using a scope here would call drop() too early..
@@ -245,42 +245,43 @@ fn main() -> Result<()> {
                         device.shell_command(commands, std::io::stdout())?;
                     }
                 }
-                UsbCommands::Pull {
+                DeviceCommands::Pull {
                     source,
                     destination,
                 } => {
                     let mut output = File::create(Path::new(&destination))?;
                     device.pull(&source, &mut output)?;
                     log::info!("Downloaded {source} as {destination}");
                 }
-                UsbCommands::Stat { path } => {
+                DeviceCommands::Stat { path } => {
                     let stat_response = device.stat(&path)?;
                     println!("{}", stat_response);
                 }
-                UsbCommands::Reboot { reboot_type } => {
+                DeviceCommands::Reboot { reboot_type } => {
                     log::info!("Reboots device in mode {:?}", reboot_type);
                     device.reboot(reboot_type.into())?
                 }
-                UsbCommands::Push { filename, path } => {
+                DeviceCommands::Push { filename, path } => {
                     let mut input = File::open(Path::new(&filename))?;
                     device.push(&mut input, &path)?;
                     log::info!("Uploaded {filename} to {path}");
                 }
-                UsbCommands::Run { package, activity } => {
+                DeviceCommands::Run { package, activity } => {
                     let output = device.run_activity(&package, &activity)?;
                     std::io::stdout().write_all(&output)?;
                 }
-                UsbCommands::Install { path } => {
+                DeviceCommands::Install { path } => {
                     log::info!("Starting installation of APK {}...", path.display());
                     device.install(path)?;
                 }
+                DeviceCommands::Framebuffer { path } => device.framebuffer(path)?,
             }
         }
         Command::Tcp(tcp) => {
             let mut device = ADBTcpDevice::new(tcp.address)?;
 
             match tcp.commands {
-                TcpCommands::Shell { commands } => {
+                DeviceCommands::Shell { commands } => {
                     if commands.is_empty() {
                         // Need to duplicate some code here as ADBTermios [Drop] implementation resets terminal state.
                         // Using a scope here would call drop() too early..
@@ -299,35 +300,36 @@ fn main() -> Result<()> {
                         device.shell_command(commands, std::io::stdout())?;
                     }
                 }
-                TcpCommands::Pull {
+                DeviceCommands::Pull {
                     source,
                     destination,
                 } => {
                     let mut output = File::create(Path::new(&destination))?;
                     device.pull(&source, &mut output)?;
                     log::info!("Downloaded {source} as {destination}");
                 }
-                TcpCommands::Stat { path } => {
+                DeviceCommands::Stat { path } => {
                     let stat_response = device.stat(&path)?;
                     println!("{}", stat_response);
                 }
-                TcpCommands::Reboot { reboot_type } => {
+                DeviceCommands::Reboot { reboot_type } => {
                     log::info!("Reboots device in mode {:?}", reboot_type);
                     device.reboot(reboot_type.into())?
                 }
-                TcpCommands::Push { filename, path } => {
+                DeviceCommands::Push { filename, path } => {
                     let mut input = File::open(Path::new(&filename))?;
                     device.push(&mut input, &path)?;
                     log::info!("Uploaded {filename} to {path}");
                 }
-                TcpCommands::Run { package, activity } => {
+                DeviceCommands::Run { package, activity } => {
                     let output = device.run_activity(&package, &activity)?;
                     std::io::stdout().write_all(&output)?;
                 }
-                TcpCommands::Install { path } => {
+                DeviceCommands::Install { path } => {
                     log::info!("Starting installation of APK {}...", path.display());
                     device.install(path)?;
                 }
+                DeviceCommands::Framebuffer { path } => device.framebuffer(path)?,
             }
         }
         Command::MdnsDiscovery => {

adb_client/src/adb_device_ext.rs

@@ -1,4 +1,4 @@
-use std::io::{Read, Write};
+use std::io::{Read, Seek, Write};
 use std::path::Path;
 
 use crate::models::AdbStatResponse;
@@ -43,4 +43,12 @@ pub trait ADBDeviceExt {
 
     /// Install an APK pointed to by `apk_path` on device.
     fn install<P: AsRef<Path>>(&mut self, apk_path: P) -> Result<()>;
+
+    /// Dump framebuffer of this device into given path
+    fn framebuffer<P: AsRef<Path>>(&mut self, path: P) -> Result<()>;
+
+    /// Dump framebuffer of this device and return corresponding bytes.
+    ///
+    /// Output data format is currently only `PNG`.
+    fn framebuffer_bytes<W: Write + Seek>(&mut self, writer: W) -> Result<()>;
 }

adb_client/src/device/adb_message_device_commands.rs

@@ -35,4 +35,12 @@ impl<T: ADBMessageTransport> ADBDeviceExt for ADBMessageDevice<T> {
     fn install<P: AsRef<std::path::Path>>(&mut self, apk_path: P) -> Result<()> {
         self.install(apk_path)
     }
+
+    fn framebuffer<P: AsRef<std::path::Path>>(&mut self, path: P) -> Result<()> {
+        self.framebuffer(path)
+    }
+
+    fn framebuffer_bytes<W: Write + std::io::Seek>(&mut self, writer: W) -> Result<()> {
+        self.framebuffer_bytes(writer)
+    }
 }

adb_client/src/device/adb_tcp_device.rs

@@ -110,6 +110,16 @@ impl ADBDeviceExt for ADBTcpDevice {
     fn install<P: AsRef<Path>>(&mut self, apk_path: P) -> Result<()> {
         self.inner.install(apk_path)
     }
+
+    #[inline]
+    fn framebuffer<P: AsRef<Path>>(&mut self, path: P) -> Result<()> {
+        self.inner.framebuffer(path)
+    }
+
+    #[inline]
+    fn framebuffer_bytes<W: std::io::Write + std::io::Seek>(&mut self, writer: W) -> Result<()> {
+        self.inner.framebuffer_bytes(writer)
+    }
 }
 
 impl Drop for ADBTcpDevice {

adb_client/src/device/adb_usb_device.rs

@@ -292,6 +292,16 @@ impl ADBDeviceExt for ADBUSBDevice {
     fn install<P: AsRef<Path>>(&mut self, apk_path: P) -> Result<()> {
         self.inner.install(apk_path)
     }
+
+    #[inline]
+    fn framebuffer<P: AsRef<Path>>(&mut self, path: P) -> Result<()> {
+        self.inner.framebuffer(path)
+    }
+
+    #[inline]
+    fn framebuffer_bytes<W: std::io::Write + std::io::Seek>(&mut self, writer: W) -> Result<()> {
+        self.inner.framebuffer_bytes(writer)
+    }
 }
 
 impl Drop for ADBUSBDevice {

adb_client/src/device/commands/framebuffer.rs

@@ -0,0 +1,108 @@
+use std::io::{Cursor, Read, Write};
+
+use byteorder::{LittleEndian, ReadBytesExt};
+use image::{ImageBuffer, ImageFormat, Rgba};
+
+use crate::{
+    device::{adb_message_device::ADBMessageDevice, ADBTransportMessage, MessageCommand},
+    models::{FrameBufferInfoV1, FrameBufferInfoV2},
+    ADBMessageTransport, Result, RustADBError,
+};
+
+impl<T: ADBMessageTransport> ADBMessageDevice<T> {
+    pub fn framebuffer<P: AsRef<std::path::Path>>(&mut self, path: P) -> Result<()> {
+        let img = self.framebuffer_inner()?;
+        Ok(img.save(path.as_ref())?)
+    }
+
+    pub fn framebuffer_bytes<W: Write + std::io::Seek>(&mut self, mut writer: W) -> Result<()> {
+        let img = self.framebuffer_inner()?;
+        Ok(img.write_to(&mut writer, ImageFormat::Png)?)
+    }
+
+    fn framebuffer_inner(&mut self) -> Result<ImageBuffer<Rgba<u8>, Vec<u8>>> {
+        let message =
+            ADBTransportMessage::new(MessageCommand::Open, 1, 0, b"framebuffer:\0".to_vec());
+
+        let response = self.send_and_expect_okay(message)?;
+
+        let local_id = response.header().arg1();
+        let remote_id = response.header().arg0();
+
+        let response = self.recv_and_reply_okay(local_id, remote_id)?;
+
+        let mut payload_cursor = Cursor::new(response.payload());
+
+        let version = payload_cursor.read_u32::<LittleEndian>()?;
+
+        let img = match version {
+            // RGBA_8888
+            1 => {
+                let mut buf = [0u8; std::mem::size_of::<FrameBufferInfoV1>()];
+
+                payload_cursor.read_exact(&mut buf)?;
+
+                let framebuffer_info: FrameBufferInfoV1 = buf.try_into()?;
+
+                let mut data = vec![
+                    0_u8;
+                    framebuffer_info
+                        .size
+                        .try_into()
+                        .map_err(|_| RustADBError::ConversionError)?
+                ];
+                payload_cursor.read_exact(&mut data)?;
+
+                ImageBuffer::<Rgba<u8>, Vec<u8>>::from_vec(
+                    framebuffer_info.width,
+                    framebuffer_info.height,
+                    data,
+                )
+                .ok_or_else(|| RustADBError::FramebufferConversionError)?
+            }
+            // RGBX_8888
+            2 => {
+                let mut buf = [0u8; std::mem::size_of::<FrameBufferInfoV2>()];
+
+                payload_cursor.read_exact(&mut buf)?;
+
+                let framebuffer_info: FrameBufferInfoV2 = buf.try_into()?;
+
+                let mut framebuffer_data = Vec::new();
+                payload_cursor.read_to_end(&mut framebuffer_data)?;
+
+                loop {
+                    if framebuffer_data.len() as u32 == framebuffer_info.size {
+                        break;
+                    }
+
+                    let response = self.recv_and_reply_okay(local_id, remote_id)?;
+
+                    framebuffer_data.extend_from_slice(&response.into_payload());
+
+                    log::debug!(
+                        "received framebuffer data. new size {}",
+                        framebuffer_data.len()
+                    );
+                }
+
+                ImageBuffer::<Rgba<u8>, Vec<u8>>::from_vec(
+                    framebuffer_info.width,
+                    framebuffer_info.height,
+                    framebuffer_data,
+                )
+                .ok_or_else(|| RustADBError::FramebufferConversionError)?
+            }
+            v => return Err(RustADBError::UnimplementedFramebufferImageVersion(v)),
+        };
+
+        let message = self.get_transport_mut().read_message()?;
+        match message.header().command() {
+            MessageCommand::Clse => Ok(img),
+            c => Err(RustADBError::ADBRequestFailed(format!(
+                "Wrong command received {}",
+                c
+            ))),
+        }
+    }
+}

adb_client/src/device/commands/mod.rs

@@ -1,3 +1,4 @@
+mod framebuffer;
 mod install;
 mod pull;
 mod push;