fix: Error when reusing same connection

This commit is contained in:
LIAUD Corentin
2024-07-13 18:53:37 +02:00
committed by cocool97
parent 2a551182ec
commit 6fe4905bb7
16 changed files with 113 additions and 75 deletions

View File

@@ -29,3 +29,4 @@ thiserror = { version = "1.0.61" }
## Marked as optional so that lib users do not depend on them
[dev-dependencies]
clap = { version = "4.5.4", features = ["derive"] }
rand = { version = "0.8.5" }

View File

@@ -47,9 +47,8 @@ use std::fs::File;
use std::path::Path;
let mut connection = AdbTcpConnection::new(Ipv4Addr::from([127,0,0,1]), 5037).unwrap();
let serial: Option<&str> = None;
let mut input = File::open(Path::new("/tmp")).unwrap();
connection.send(serial, &mut input, "/data/local/tmp");
connection.send::<&str,&str>(None, &mut input, "/data/local/tmp");
```
## Rust binary

View File

@@ -111,16 +111,16 @@ fn main() -> Result<(), RustADBError> {
}
Command::Pull { path, filename } => {
let mut output = File::create(Path::new(&filename)).unwrap(); // TODO: Better error handling
connection.recv(opt.serial, &path, &mut output)?;
connection.recv(opt.serial.as_ref(), &path, &mut output)?;
println!("Downloaded {path} as {filename}");
}
Command::Push { filename, path } => {
let mut input = File::open(Path::new(&filename)).unwrap(); // TODO: Better error handling
connection.send(opt.serial, &mut input, &path)?;
connection.send(opt.serial.as_ref(), &mut input, &path)?;
println!("Uploaded {filename} to {path}");
}
Command::List { path } => {
connection.list(opt.serial, path)?;
connection.list(opt.serial.as_ref(), path)?;
}
Command::Stat { path } => {
let stat_response = connection.stat(opt.serial, path)?;
@@ -128,20 +128,20 @@ fn main() -> Result<(), RustADBError> {
}
Command::Shell { command } => {
if command.is_empty() {
connection.shell(&opt.serial)?;
connection.shell(opt.serial.as_ref())?;
} else {
connection.shell_command(&opt.serial, command)?;
connection.shell_command(opt.serial.as_ref(), command)?;
}
}
Command::HostFeatures => {
println!("Available host features");
for feature in connection.host_features(&opt.serial)? {
for feature in connection.host_features(opt.serial.as_ref())? {
println!("- {}", feature);
}
}
Command::Reboot { sub_command } => {
println!("Reboots device");
connection.reboot(&opt.serial, sub_command.into())?
connection.reboot(opt.serial.as_ref(), sub_command.into())?
}
}

View File

@@ -42,8 +42,9 @@ impl AdbTcpConnection {
&mut self,
adb_command: AdbCommand,
with_response: bool,
fresh_connection: bool
) -> Result<Vec<u8>> {
self.send_adb_request(adb_command)?;
self.send_adb_request(adb_command, fresh_connection)?;
if with_response {
let length = self.get_hex_body_length()?;
@@ -65,7 +66,16 @@ impl AdbTcpConnection {
/// Sends the given [AdbCommand] to ADB server, and checks that the request has been taken in consideration.
/// If an error occurred, a [RustADBError] is returned with the response error string.
pub(crate) fn send_adb_request(&mut self, command: AdbCommand) -> Result<()> {
pub(crate) fn send_adb_request(
&mut self,
command: AdbCommand,
fresh_connection: bool,
) -> Result<()> {
if fresh_connection {
// Recreate a new connection (likely because command does not need "state" server side)
self.new_connection()?;
}
let adb_command_string = command.to_string();
let adb_request = format!("{:04x}{}", adb_command_string.len(), adb_command_string);

View File

@@ -5,7 +5,7 @@ use crate::{models::AdbCommand, AdbTcpConnection, Device, DeviceLong, Result, Ru
impl AdbTcpConnection {
/// Gets a list of connected devices.
pub fn devices(&mut self) -> Result<Vec<Device>> {
let devices = self.proxy_connection(AdbCommand::Devices, true)?;
let devices = self.proxy_connection(AdbCommand::Devices, true, true)?;
let mut vec_devices: Vec<Device> = vec![];
for device in devices.split(|x| x.eq(&b'\n')) {
@@ -21,7 +21,7 @@ impl AdbTcpConnection {
/// Gets an extended list of connected devices including the device paths in the state.
pub fn devices_long(&mut self) -> Result<Vec<DeviceLong>> {
let devices_long = self.proxy_connection(AdbCommand::DevicesLong, true)?;
let devices_long = self.proxy_connection(AdbCommand::DevicesLong, true, true)?;
let mut vec_devices: Vec<DeviceLong> = vec![];
for device in devices_long.split(|x| x.eq(&b'\n')) {
@@ -38,7 +38,7 @@ impl AdbTcpConnection {
/// Tracks new devices showing up.
// TODO: Change with Generator when feature stabilizes
pub fn track_devices(&mut self, callback: impl Fn(Device) -> Result<()>) -> Result<()> {
self.send_adb_request(AdbCommand::TrackDevices)?;
self.send_adb_request(AdbCommand::TrackDevices, true)?;
loop {
let length = self.get_hex_body_length()?;

View File

@@ -5,15 +5,15 @@ use crate::{
impl AdbTcpConnection {
/// Lists available ADB server features.
pub fn host_features<S: ToString>(&mut self, serial: &Option<S>) -> Result<Vec<HostFeatures>> {
pub fn host_features<S: ToString>(&mut self, serial: Option<&S>) -> Result<Vec<HostFeatures>> {
match serial {
None => self.send_adb_request(AdbCommand::TransportAny)?,
None => self.send_adb_request(AdbCommand::TransportAny, true)?,
Some(serial) => {
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()))?
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()), true)?
}
}
let features = self.proxy_connection(AdbCommand::HostFeatures, true)?;
let features = self.proxy_connection(AdbCommand::HostFeatures, true, false)?;
Ok(features
.split(|x| x.eq(&b','))

View File

@@ -3,6 +3,6 @@ use crate::{models::AdbCommand, AdbTcpConnection, Result};
impl AdbTcpConnection {
/// Asks the ADB server to quit immediately.
pub fn kill(&mut self) -> Result<()> {
self.proxy_connection(AdbCommand::Kill, false).map(|_| ())
self.proxy_connection(AdbCommand::Kill, false, true).map(|_| ())
}
}

View File

@@ -10,18 +10,16 @@ use std::{
impl AdbTcpConnection {
/// Lists files in [path] on the device.
pub fn list<S: ToString, A: AsRef<str>>(&mut self, serial: Option<S>, path: A) -> Result<()> {
self.new_connection()?;
pub fn list<S: ToString, A: AsRef<str>>(&mut self, serial: Option<&S>, path: A) -> Result<()> {
match serial {
None => self.send_adb_request(AdbCommand::TransportAny)?,
None => self.send_adb_request(AdbCommand::TransportAny, false)?,
Some(serial) => {
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()))?
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()), false)?
}
}
// Set device in SYNC mode
self.send_adb_request(AdbCommand::Sync)?;
self.send_adb_request(AdbCommand::Sync, false)?;
// Send a list command
self.send_sync_request(SyncCommand::List)?;

View File

@@ -7,17 +7,17 @@ impl AdbTcpConnection {
/// Reboots the device
pub fn reboot<S: ToString>(
&mut self,
serial: &Option<S>,
serial: Option<&S>,
reboot_type: RebootType,
) -> Result<()> {
match serial {
None => self.send_adb_request(AdbCommand::TransportAny)?,
None => self.send_adb_request(AdbCommand::TransportAny, true)?,
Some(serial) => {
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()))?
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()), true)?
}
}
self.proxy_connection(AdbCommand::Reboot(reboot_type), false)
self.proxy_connection(AdbCommand::Reboot(reboot_type), false, false)
.map(|_| ())
}
}

View File

@@ -9,21 +9,19 @@ impl AdbTcpConnection {
/// Receives [path] to [stream] from the device.
pub fn recv<S: ToString, A: AsRef<str>>(
&mut self,
serial: Option<S>,
serial: Option<&S>,
path: A,
stream: &mut dyn Write,
) -> Result<()> {
self.new_connection()?;
match serial {
None => self.send_adb_request(AdbCommand::TransportAny)?,
None => self.send_adb_request(AdbCommand::TransportAny, true)?,
Some(serial) => {
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()))?
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()), true)?
}
}
// Set device in SYNC mode
self.send_adb_request(AdbCommand::Sync)?;
self.send_adb_request(AdbCommand::Sync, false)?;
// Send a recv command
self.send_sync_request(SyncCommand::Recv)?;
@@ -54,7 +52,7 @@ impl AdbTcpConnection {
// Handle received data
let length: usize = self.get_body_length()?.try_into().unwrap();
self.tcp_stream.read_exact(&mut buffer[..length])?;
output.write_all(&buffer)?;
output.write_all(&buffer[..length])?;
}
b"DONE" => break, // We're done here
b"FAIL" => {

View File

@@ -14,21 +14,19 @@ impl AdbTcpConnection {
/// Sends [stream] to [path] on the device.
pub fn send<S: ToString, A: AsRef<str>>(
&mut self,
serial: Option<S>,
serial: Option<&S>,
stream: &mut dyn Read,
path: A,
) -> Result<()> {
self.new_connection()?;
match serial {
None => self.send_adb_request(AdbCommand::TransportAny)?,
None => self.send_adb_request(AdbCommand::TransportAny, true)?,
Some(serial) => {
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()))?
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()), true)?
}
}
// Set device in SYNC mode
self.send_adb_request(AdbCommand::Sync)?;
self.send_adb_request(AdbCommand::Sync, false)?;
// Send a send command
self.send_sync_request(SyncCommand::Send)?;

View File

@@ -29,7 +29,7 @@ impl AdbTcpConnection {
/// Runs 'command' in a shell on the device, and return its output and error streams.
pub fn shell_command<S: ToString>(
&mut self,
serial: &Option<S>,
serial: Option<&S>,
command: impl IntoIterator<Item = S>,
) -> Result<Vec<u8>> {
let supported_features = self.host_features(serial)?;
@@ -39,24 +39,25 @@ impl AdbTcpConnection {
return Err(RustADBError::ADBShellNotSupported);
}
self.new_connection()?;
match serial {
None => self.send_adb_request(AdbCommand::TransportAny)?,
None => self.send_adb_request(AdbCommand::TransportAny, true)?,
Some(serial) => {
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()))?
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()), true)?
}
}
self.send_adb_request(AdbCommand::ShellCommand(
command
.into_iter()
.map(|v| v.to_string())
.collect::<Vec<_>>()
.join(" "),
))?;
self.send_adb_request(
AdbCommand::ShellCommand(
command
.into_iter()
.map(|v| v.to_string())
.collect::<Vec<_>>()
.join(" "),
),
false,
)?;
const BUFFER_SIZE: usize = 512;
let result = (|| {
(|| {
let mut result = Vec::new();
loop {
let mut buffer = [0; BUFFER_SIZE];
@@ -73,20 +74,17 @@ impl AdbTcpConnection {
}
}
}
})();
self.new_connection()?;
result
})()
}
/// Starts an interactive shell session on the device. Redirects stdin/stdout/stderr as appropriate.
pub fn shell<S: ToString>(&mut self, serial: &Option<S>) -> Result<()> {
pub fn shell<S: ToString>(&mut self, serial: Option<&S>) -> Result<()> {
let mut adb_termios = ADBTermios::new(std::io::stdin())?;
adb_termios.set_adb_termios()?;
self.tcp_stream.set_nodelay(true)?;
// FORWARD CTRL+C !!
// TODO: FORWARD CTRL+C !!
let supported_features = self.host_features(serial)?;
if !supported_features.contains(&HostFeatures::ShellV2)
@@ -95,15 +93,13 @@ impl AdbTcpConnection {
return Err(RustADBError::ADBShellNotSupported);
}
self.new_connection()?;
match serial {
None => self.send_adb_request(AdbCommand::TransportAny)?,
None => self.send_adb_request(AdbCommand::TransportAny, true)?,
Some(serial) => {
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()))?
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()), true)?
}
}
self.send_adb_request(AdbCommand::Shell)?;
self.send_adb_request(AdbCommand::Shell, false)?;
// let read_stream = Arc::new(self.tcp_stream);
let mut read_stream = self.tcp_stream.try_clone()?;

View File

@@ -79,17 +79,15 @@ impl AdbTcpConnection {
serial: Option<S>,
path: A,
) -> Result<AdbStatResponse> {
self.new_connection()?;
match serial {
None => self.send_adb_request(AdbCommand::TransportAny)?,
None => self.send_adb_request(AdbCommand::TransportAny, false)?,
Some(serial) => {
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()))?
self.send_adb_request(AdbCommand::TransportSerial(serial.to_string()), false)?
}
}
// Set device in SYNC mode
self.send_adb_request(AdbCommand::Sync)?;
self.send_adb_request(AdbCommand::Sync, false)?;
// Send a "Stat" command
self.send_sync_request(SyncCommand::Stat)?;

View File

@@ -3,7 +3,7 @@ use crate::{models::AdbCommand, AdbTcpConnection, Result};
impl AdbTcpConnection {
/// Asks ADB server to switch the connection to either the device or emulator connect to/running on the host. Will fail if there is more than one such device/emulator available.
pub fn transport_any(&mut self) -> Result<()> {
self.proxy_connection(AdbCommand::TransportAny, false)
self.proxy_connection(AdbCommand::TransportAny, false, true)
.map(|_| ())
}
}

View File

@@ -3,7 +3,7 @@ use crate::{models::AdbCommand, AdbTcpConnection, AdbVersion, Result};
impl AdbTcpConnection {
/// Gets server's internal version number.
pub fn version(&mut self) -> Result<AdbVersion> {
let version = self.proxy_connection(AdbCommand::Version, true)?;
let version = self.proxy_connection(AdbCommand::Version, true, true)?;
AdbVersion::try_from(version)
}

View File

@@ -1,9 +1,11 @@
#[cfg(test)]
mod tests {
use std::io::Cursor;
use std::net::Ipv4Addr;
use std::str::FromStr;
use adb_client::AdbTcpConnection;
use rand::Rng;
fn new_client() -> AdbTcpConnection {
let address = Ipv4Addr::from_str("127.0.0.1").unwrap();
@@ -19,8 +21,8 @@ mod tests {
#[test]
fn test_shell() {
let mut adb = new_client();
adb.shell_command(&None, vec!["ls"]).unwrap();
adb.shell_command(&None, vec!["pwd"]).unwrap();
adb.shell_command(None, vec!["ls"]).unwrap();
adb.shell_command(None, vec!["pwd"]).unwrap();
}
#[test]
@@ -41,4 +43,42 @@ mod tests {
let address = Ipv4Addr::from_str("127.0.0.300").unwrap();
let _ = AdbTcpConnection::new(address, 5037).expect("Could not create ADB connection...");
}
#[test]
fn test_send_recv() {
// Create random "Reader" in memory
let mut key = [0u8; 1000];
rand::thread_rng().fill(&mut key[..]);
let mut c: Cursor<Vec<u8>> = Cursor::new(key.to_vec());
let mut connection = new_client();
const TEST_FILENAME: &'static str = "/data/local/tmp/test_file";
// Send it
connection
.send::<&str, &str>(None, &mut c, TEST_FILENAME)
.expect("cannot send file");
// Pull it to memory
let mut res = vec![];
connection
.recv::<&str, &str>(None, TEST_FILENAME, &mut res)
.expect("cannot recv file");
// diff
assert_eq!(c.get_ref(), &res);
connection
.shell_command::<&str>(None, [format!("rm {TEST_FILENAME}").as_str()])
.expect("cannot remove test file");
}
#[test]
fn multiple_connexions() {
let mut connection = new_client();
for _ in 0..2 {
let _ = connection.devices().expect("cannot get version");
}
}
}