chore: minor code improvements

This allows multiple operations to be performed simultaneously

Co-authored-by: gaykitty

adb_client/src/device/adb_message_device.rs

@@ -1,5 +1,6 @@
 use super::{ADBRsaKey, ADBTransportMessage, MessageCommand, models::MessageSubcommand};
 use crate::device::adb_transport_message::{AUTH_RSAPUBLICKEY, AUTH_SIGNATURE, AUTH_TOKEN};
+use crate::device::models::ADBSession;
 use crate::{ADBMessageTransport, AdbStatResponse, Result, RustADBError, constants::BUFFER_SIZE};
 use bincode::config::{Configuration, Fixint, LittleEndian, NoLimit};
 use byteorder::ReadBytesExt;
@@ -27,18 +28,12 @@ pub(crate) fn bincode_deserialize_from_slice<D: DeserializeOwned>(data: &[u8]) -
 #[derive(Debug)]
 pub struct ADBMessageDevice<T: ADBMessageTransport> {
     transport: T,
-    local_id: Option<u32>,
-    remote_id: Option<u32>,
 }
 
 impl<T: ADBMessageTransport> ADBMessageDevice<T> {
     /// Instantiate a new [`ADBMessageTransport`]
     pub fn new(transport: T) -> Self {
-        Self {
-            transport,
-            local_id: None,
-            remote_id: None,
-        }
+        Self { transport }
     }
 
     pub(crate) fn get_transport(&mut self) -> &T {
@@ -110,12 +105,15 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
     }
 
     /// Receive a message and acknowledge it by replying with an `OKAY` command
-    pub(crate) fn recv_and_reply_okay(&mut self) -> Result<ADBTransportMessage> {
+    pub(crate) fn recv_and_reply_okay(
+        &mut self,
+        session: ADBSession,
+    ) -> Result<ADBTransportMessage> {
         let message = self.transport.read_message()?;
         self.transport.write_message(ADBTransportMessage::new(
             MessageCommand::Okay,
-            self.get_local_id()?,
-            self.get_remote_id()?,
+            session.local_id(),
+            session.remote_id(),
             &[],
         ))?;
         Ok(message)
@@ -136,11 +134,12 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
 
     pub(crate) fn recv_file<W: std::io::Write>(
         &mut self,
+        session: ADBSession,
         mut output: W,
     ) -> std::result::Result<(), RustADBError> {
         let mut len: Option<u64> = None;
         loop {
-            let payload = self.recv_and_reply_okay()?.into_payload();
+            let payload = self.recv_and_reply_okay(session)?.into_payload();
             let mut rdr = Cursor::new(&payload);
             while rdr.position() != payload.len() as u64 {
                 match len.take() {
@@ -173,8 +172,7 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
 
     pub(crate) fn push_file<R: std::io::Read>(
         &mut self,
-        local_id: u32,
-        remote_id: u32,
+        session: ADBSession,
         mut reader: R,
     ) -> std::result::Result<(), RustADBError> {
         let mut buffer = vec![0; BUFFER_SIZE].into_boxed_slice();
@@ -186,8 +184,8 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
 
         let message = ADBTransportMessage::new(
             MessageCommand::Write,
-            local_id,
-            remote_id,
+            session.local_id(),
+            session.remote_id(),
             &serialized_message,
         );
 
@@ -204,8 +202,8 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
                     let serialized_message = bincode_serialize_to_vec(&subcommand_data)?;
                     let message = ADBTransportMessage::new(
                         MessageCommand::Write,
-                        local_id,
-                        remote_id,
+                        session.local_id(),
+                        session.remote_id(),
                         &serialized_message,
                     );
 
@@ -230,8 +228,8 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
 
                     let message = ADBTransportMessage::new(
                         MessageCommand::Write,
-                        local_id,
-                        remote_id,
+                        session.local_id(),
+                        session.remote_id(),
                         &serialized_message,
                     );
 
@@ -244,24 +242,27 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
         }
     }
 
-    pub(crate) fn begin_synchronization(&mut self) -> Result<()> {
-        self.open_session(b"sync:\0")?;
-        Ok(())
+    pub(crate) fn begin_synchronization(&mut self) -> Result<ADBSession> {
+        self.open_session(b"sync:\0")
     }
 
-    pub(crate) fn stat_with_explicit_ids(&mut self, remote_path: &str) -> Result<AdbStatResponse> {
+    pub(crate) fn stat_with_explicit_ids(
+        &mut self,
+        session: ADBSession,
+        remote_path: &str,
+    ) -> Result<AdbStatResponse> {
         let stat_buffer = MessageSubcommand::Stat.with_arg(u32::try_from(remote_path.len())?);
         let message = ADBTransportMessage::new(
             MessageCommand::Write,
-            self.get_local_id()?,
-            self.get_remote_id()?,
+            session.local_id(),
+            session.remote_id(),
             &bincode_serialize_to_vec(&stat_buffer)?,
         );
         self.send_and_expect_okay(message)?;
         self.send_and_expect_okay(ADBTransportMessage::new(
             MessageCommand::Write,
-            self.get_local_id()?,
-            self.get_remote_id()?,
+            session.local_id(),
+            session.remote_id(),
             remote_path.as_bytes(),
         ))?;
         let response = self.transport.read_message()?;
@@ -271,24 +272,25 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
         bincode_deserialize_from_slice(&response.into_payload()[4..])
     }
 
-    pub(crate) fn end_transaction(&mut self) -> Result<()> {
+    pub(crate) fn end_transaction(&mut self, session: ADBSession) -> Result<()> {
         let quit_buffer = MessageSubcommand::Quit.with_arg(0u32);
         self.send_and_expect_okay(ADBTransportMessage::new(
             MessageCommand::Write,
-            self.get_local_id()?,
-            self.get_remote_id()?,
+            session.local_id(),
+            session.remote_id(),
             &bincode_serialize_to_vec(&quit_buffer)?,
         ))?;
         let _discard_close = self.transport.read_message()?;
         Ok(())
     }
 
-    pub(crate) fn open_session(&mut self, data: &[u8]) -> Result<ADBTransportMessage> {
+    pub(crate) fn open_session(&mut self, data: &[u8]) -> Result<ADBSession> {
         let mut rng = rand::rng();
+        let local_id: u32 = rng.random();
 
         let message = ADBTransportMessage::new(
             MessageCommand::Open,
-            rng.random(), // Our 'local-id'
+            local_id, // Our 'local-id'
             0,
             data,
         );
@@ -296,21 +298,20 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
 
         let response = self.get_transport_mut().read_message()?;
 
-        self.local_id = Some(response.header().arg1());
-        self.remote_id = Some(response.header().arg0());
+        if response.header().command() != MessageCommand::Okay {
+            return Err(RustADBError::ADBRequestFailed(format!(
+                "Open session failed: got {} in respone instead of OKAY",
+                response.header().command()
+            )));
+        }
 
-        Ok(response)
-    }
+        if response.header().arg1() != local_id {
+            return Err(RustADBError::ADBRequestFailed(format!(
+                "Open session failed: respones used {} for our local_id instead of {local_id}",
+                response.header().arg1()
+            )));
+        }
 
-    pub(crate) fn get_local_id(&self) -> Result<u32> {
-        self.local_id.ok_or(RustADBError::ADBRequestFailed(
-            "connection not opened, no local_id".into(),
-        ))
-    }
-
-    pub(crate) fn get_remote_id(&self) -> Result<u32> {
-        self.remote_id.ok_or(RustADBError::ADBRequestFailed(
-            "connection not opened, no remote_id".into(),
-        ))
+        Ok(ADBSession::new(local_id, response.header().arg0()))
     }
 }

adb_client/src/device/commands/framebuffer.rs

@@ -11,9 +11,9 @@ use crate::{
 
 impl<T: ADBMessageTransport> ADBMessageDevice<T> {
     pub(crate) fn framebuffer_inner(&mut self) -> Result<ImageBuffer<Rgba<u8>, Vec<u8>>> {
-        self.open_session(b"framebuffer:\0")?;
+        let session = self.open_session(b"framebuffer:\0")?;
 
-        let response = self.recv_and_reply_okay()?;
+        let response = self.recv_and_reply_okay(session)?;
 
         let mut payload_cursor = Cursor::new(response.payload());
 
@@ -36,7 +36,7 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
                         break;
                     }
 
-                    let response = self.recv_and_reply_okay()?;
+                    let response = self.recv_and_reply_okay(session)?;
 
                     framebuffer_data.extend_from_slice(&response.into_payload());
 
@@ -69,7 +69,7 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
                         break;
                     }
 
-                    let response = self.recv_and_reply_okay()?;
+                    let response = self.recv_and_reply_okay(session)?;
 
                     framebuffer_data.extend_from_slice(&response.into_payload());
 

adb_client/src/device/commands/install.rs

@@ -14,11 +14,12 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
 
         let file_size = apk_file.metadata()?.len();
 
-        self.open_session(format!("exec:cmd package 'install' -S {file_size}\0").as_bytes())?;
+        let session =
+            self.open_session(format!("exec:cmd package 'install' -S {file_size}\0").as_bytes())?;
 
         let transport = self.get_transport().clone();
 
-        let mut writer = MessageWriter::new(transport, self.get_local_id()?, self.get_remote_id()?);
+        let mut writer = MessageWriter::new(transport, session.local_id(), session.remote_id());
 
         std::io::copy(&mut apk_file, &mut writer)?;
 

adb_client/src/device/commands/list.rs

@@ -13,11 +13,11 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
     /// List the entries in the given directory on the device.
     /// note: path uses internal file paths, so Documents is at /storage/emulated/0/Documents
     pub(crate) fn list<A: AsRef<str>>(&mut self, path: A) -> Result<Vec<ADBListItem>> {
-        self.begin_synchronization()?;
+        let session = self.begin_synchronization()?;
 
-        let output = self.handle_list(path);
+        let output = self.handle_list(path, session.local_id(), session.remote_id());
 
-        self.end_transaction()?;
+        self.end_transaction(session)?;
         output
     }
 
@@ -70,11 +70,14 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
         }
     }
 
-    fn handle_list<A: AsRef<str>>(&mut self, path: A) -> Result<Vec<ADBListItem>> {
+    fn handle_list<A: AsRef<str>>(
+        &mut self,
+        path: A,
+        local_id: u32,
+        remote_id: u32,
+    ) -> Result<Vec<ADBListItem>> {
         // TODO: use LIS2 to support files over 2.14 GB in size.
         // SEE: https://github.com/cstyan/adbDocumentation?tab=readme-ov-file#adb-list
-        let local_id = self.get_local_id()?;
-        let remote_id = self.get_remote_id()?;
         {
             let mut len_buf = Vec::from([0_u8; 4]);
             LittleEndian::write_u32(&mut len_buf, path.as_ref().len() as u32);

adb_client/src/device/commands/pull.rs

@@ -11,10 +11,10 @@ use crate::{
 
 impl<T: ADBMessageTransport> ADBMessageDevice<T> {
     pub(crate) fn pull<A: AsRef<str>, W: Write>(&mut self, source: A, output: W) -> Result<()> {
-        self.begin_synchronization()?;
+        let session = self.begin_synchronization()?;
         let source = source.as_ref();
 
-        let adb_stat_response = self.stat_with_explicit_ids(source)?;
+        let adb_stat_response = self.stat_with_explicit_ids(session, source)?;
 
         if adb_stat_response.file_perm == 0 {
             return Err(RustADBError::UnknownResponseType(
@@ -22,11 +22,13 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
             ));
         }
 
-        let local_id = self.get_local_id()?;
-        let remote_id = self.get_remote_id()?;
-
         self.get_transport_mut().write_message_with_timeout(
-            ADBTransportMessage::new(MessageCommand::Okay, local_id, remote_id, &[]),
+            ADBTransportMessage::new(
+                MessageCommand::Okay,
+                session.local_id(),
+                session.remote_id(),
+                &[],
+            ),
             std::time::Duration::from_secs(4),
         )?;
 
@@ -34,19 +36,19 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
         let recv_buffer = adb_message_device::bincode_serialize_to_vec(&recv_buffer)?;
         self.send_and_expect_okay(ADBTransportMessage::new(
             MessageCommand::Write,
-            self.get_local_id()?,
-            self.get_remote_id()?,
+            session.local_id(),
+            session.remote_id(),
             &recv_buffer,
         ))?;
         self.send_and_expect_okay(ADBTransportMessage::new(
             MessageCommand::Write,
-            self.get_local_id()?,
-            self.get_remote_id()?,
+            session.local_id(),
+            session.remote_id(),
             source.as_bytes(),
         ))?;
 
-        self.recv_file(output)?;
-        self.end_transaction()?;
+        self.recv_file(session, output)?;
+        self.end_transaction(session)?;
         Ok(())
     }
 }

adb_client/src/device/commands/push.rs

@@ -10,7 +10,7 @@ use crate::{
 
 impl<T: ADBMessageTransport> ADBMessageDevice<T> {
     pub(crate) fn push<R: Read, A: AsRef<str>>(&mut self, stream: R, path: A) -> Result<()> {
-        self.begin_synchronization()?;
+        let session = self.begin_synchronization()?;
 
         let path_header = format!("{},0777", path.as_ref());
 
@@ -20,14 +20,13 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
 
         self.send_and_expect_okay(ADBTransportMessage::new(
             MessageCommand::Write,
-            self.get_local_id()?,
-            self.get_remote_id()?,
+            session.local_id(),
+            session.remote_id(),
             &send_buffer,
         ))?;
 
-        self.push_file(self.get_local_id()?, self.get_remote_id()?, stream)?;
-
-        self.end_transaction()?;
+        self.push_file(session, stream)?;
+        self.end_transaction(session)?;
 
         Ok(())
     }

adb_client/src/device/commands/shell.rs

@@ -11,20 +11,10 @@ use crate::{
 impl<T: ADBMessageTransport> ADBMessageDevice<T> {
     /// Runs 'command' in a shell on the device, and write its output and error streams into output.
     pub(crate) fn shell_command(&mut self, command: &[&str], output: &mut dyn Write) -> Result<()> {
-        let response = self.open_session(format!("shell:{}\0", command.join(" "),).as_bytes())?;
+        let session = self.open_session(format!("shell:{}\0", command.join(" "),).as_bytes())?;
 
         let mut transport = self.get_transport().clone();
 
-        let local_id = self.get_local_id()?;
-        let remote_id = self.get_remote_id()?;
-
-        if response.header().command() != MessageCommand::Okay {
-            return Err(RustADBError::ADBRequestFailed(format!(
-                "wrong command {}",
-                response.header().command()
-            )));
-        }
-
         loop {
             let message = transport.read_message()?;
             let command = message.header().command();
@@ -36,8 +26,8 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
             self.get_transport_mut()
                 .write_message(ADBTransportMessage::new(
                     MessageCommand::Okay,
-                    local_id,
-                    remote_id,
+                    session.local_id(),
+                    session.remote_id(),
                     &[],
                 ))?;
 
@@ -48,9 +38,7 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
         // the client has acknowledged it. Read them quickly if present.
         while let Ok(_discard_close_message) =
             transport.read_message_with_timeout(Duration::from_millis(20))
-        {
-            continue;
-        }
+        {}
 
         Ok(())
     }
@@ -62,21 +50,22 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
         mut reader: &mut dyn Read,
         mut writer: Box<dyn Write + Send>,
     ) -> Result<()> {
-        self.open_session(b"shell:\0")?;
+        let session = self.open_session(b"shell:\0")?;
 
         let mut transport = self.get_transport().clone();
 
-        let local_id = self.get_local_id()?;
-        let remote_id = self.get_remote_id()?;
-
         // Reading thread, reads response from adbd
         std::thread::spawn(move || -> Result<()> {
             loop {
                 let message = transport.read_message()?;
 
                 // Acknowledge for more data
-                let response =
-                    ADBTransportMessage::new(MessageCommand::Okay, local_id, remote_id, &[]);
+                let response = ADBTransportMessage::new(
+                    MessageCommand::Okay,
+                    session.local_id(),
+                    session.remote_id(),
+                    &[],
+                );
                 transport.write_message(response)?;
 
                 match message.header().command() {
@@ -91,7 +80,8 @@ impl<T: ADBMessageTransport> ADBMessageDevice<T> {
         });
 
         let transport = self.get_transport().clone();
-        let mut shell_writer = ShellMessageWriter::new(transport, local_id, remote_id);
+        let mut shell_writer =
+            ShellMessageWriter::new(transport, session.local_id(), session.remote_id());
 
         // Read from given reader (that could be stdin e.g), and write content to device adbd
         if let Err(e) = std::io::copy(&mut reader, &mut shell_writer) {

adb_client/src/device/commands/stat.rs

@@ -4,9 +4,9 @@ use crate::{
 
 impl<T: ADBMessageTransport> ADBMessageDevice<T> {
     pub(crate) fn stat(&mut self, remote_path: &str) -> Result<AdbStatResponse> {
-        self.begin_synchronization()?;
-        let adb_stat_response = self.stat_with_explicit_ids(remote_path)?;
-        self.end_transaction()?;
+        let session = self.begin_synchronization()?;
+        let adb_stat_response = self.stat_with_explicit_ids(session, remote_path)?;
+        self.end_transaction(session)?;
         Ok(adb_stat_response)
     }
 }

adb_client/src/device/models/adb_session.rs

@@ -0,0 +1,23 @@
+/// Represent a session between an ADBDevice and remote `adbd`.
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct ADBSession {
+    local_id: u32,
+    remote_id: u32,
+}
+
+impl ADBSession {
+    pub fn new(local_id: u32, remote_id: u32) -> Self {
+        Self {
+            local_id,
+            remote_id,
+        }
+    }
+
+    pub fn local_id(self) -> u32 {
+        self.local_id
+    }
+
+    pub fn remote_id(self) -> u32 {
+        self.remote_id
+    }
+}

adb_client/src/device/models/mod.rs

@@ -1,5 +1,7 @@
 mod adb_rsa_key;
+mod adb_session;
 mod message_commands;
 
 pub use adb_rsa_key::ADBRsaKey;
+pub(crate) use adb_session::ADBSession;
 pub use message_commands::{MessageCommand, MessageSubcommand};