From 3bcc9282ed2a945a553cc89a122146f148d0b6e1 Mon Sep 17 00:00:00 2001 From: Bjorn Andersson Date: Fri, 3 Jul 2026 02:56:50 +0000 Subject: [PATCH 1/5] ssh: fall back to private keys without ssh-agent The SSH helper assumes ssh-agent is always available and panics when the agent cannot be contacted. That makes library users fail before they can try another usable key source. Try ssh-agent first, but treat agent connection and identity lookup failures as normal authentication misses. Fall back to SK8BRD_SSH_KEY, id_ed25519, and id_rsa, and only report failure after every candidate has been rejected. Signed-off-by: Bjorn Andersson --- proto/src/ssh.rs | 115 ++++++++++++++++++++++++++++++++++++----------- 1 file changed, 90 insertions(+), 25 deletions(-) diff --git a/proto/src/ssh.rs b/proto/src/ssh.rs index 94f8d3b..e63a33b 100644 --- a/proto/src/ssh.rs +++ b/proto/src/ssh.rs @@ -2,7 +2,9 @@ use anyhow::{Context as _, bail}; use asynchronous_codec::BytesMut; use russh::Channel; use russh::client::{self, Msg}; -use russh::keys::{HashAlg, ssh_key}; +use russh::keys::load_secret_key; +use russh::keys::{HashAlg, PrivateKeyWithHashAlg, ssh_key}; +use std::path::{Path, PathBuf}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -29,40 +31,78 @@ pub async fn ssh_connect(farm: &str, username: String) -> anyhow::Result key_pair, + Err(_) => continue, + }; + + let hash_alg = if key_pair.algorithm().is_rsa() { + sess.best_supported_rsa_hash() + .await + .with_context(|| format!("Could not check RSA signatures for {path:?}"))? + .flatten() + } else { + None + }; + let private_key = PrivateKeyWithHashAlg::new(Arc::new(key_pair), hash_alg); + + if sess + .authenticate_publickey(&username, private_key) + .await + .map(|result| result.success()) + .unwrap_or(false) + { + authenticated = true; + break; + } + } + } + + if !authenticated || sess.is_closed() { bail!("No key was accepted by the server"); } @@ -74,6 +114,31 @@ pub async fn ssh_connect(farm: &str, username: String) -> anyhow::Result Vec { + let mut keys = Vec::new(); + + if let Ok(key_path) = std::env::var("SK8BRD_SSH_KEY") { + keys.push(expand_tilde(key_path)); + } + if let Ok(home) = std::env::var("HOME") { + keys.push(Path::new(&home).join(".ssh").join("id_ed25519")); + keys.push(Path::new(&home).join(".ssh").join("id_rsa")); + } + + keys +} + +fn expand_tilde(path: String) -> PathBuf { + let path = path.trim().to_string(); + if let Some(tail) = path.strip_prefix("~/") { + return std::env::var("HOME") + .ok() + .map_or_else(PathBuf::new, |home| Path::new(&home).join(tail)); + } + + Path::new(&path).to_path_buf() +} + pub struct Wrap(Receiver>, BytesMut); impl Wrap { From 2d8c79c29e999bdf43f3bf885421f4cb24583071 Mon Sep 17 00:00:00 2001 From: Bjorn Andersson Date: Fri, 3 Jul 2026 02:58:03 +0000 Subject: [PATCH 2/5] client: read stdout and stderr concurrently The interactive and non-interactive clients drain stderr before they read the framed protocol from stdout. If the remote server has no pending stderr data, the clients can block forever and never process the stdout message that would let the boot flow progress. Read both SSH streams with tokio::select!, buffer partial stdout frames, and parse every complete protocol message as it arrives. Track stderr EOF so the select loop does not spin after the status stream closes, wake the interactive loop so Ctrl-A q can exit when the console is idle, and bound the shutdown power-off request so exit does not hang if the remote writer stalls. Signed-off-by: Bjorn Andersson --- cli/src/main.rs | 152 ++++++++++++++++++++++++++++----------------- client/src/main.rs | 138 ++++++++++++++++++++++++---------------- proto/src/ssh.rs | 38 +++++------- 3 files changed, 196 insertions(+), 132 deletions(-) diff --git a/cli/src/main.rs b/cli/src/main.rs index f938911..a867c86 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -10,9 +10,10 @@ use sk8brd::{ use std::fs; use std::io::{stdout, Write}; use std::sync::Arc; -use std::time::{Duration, SystemTime}; +use std::time::{Duration, Instant}; use tokio::io::AsyncReadExt; use tokio::sync::Mutex; +use tokio::time::{sleep_until, timeout}; #[derive(Parser, Debug)] #[command(version, about, long_about = None)] @@ -41,11 +42,15 @@ struct Args { #[tokio::main] async fn main() -> anyhow::Result<()> { + const MAX_MSG_LEN: usize = 1024 * 1024; let quit = Arc::new(Mutex::new(false)); - let mut buf = [0u8; SSH_BUFFER_SIZE]; - let mut time: SystemTime = SystemTime::now(); - let mut hdr_buf = [0u8; MSG_HDR_SIZE]; + let mut stderr_buf = [0u8; SSH_BUFFER_SIZE]; + let mut stdout_chunk = [0u8; SSH_BUFFER_SIZE]; + let mut stdout_buf: Vec = Vec::new(); + let mut stderr_open = true; let args = Args::parse(); + let mut deadline = Instant::now() + Duration::from_secs(args.timeout); + let mut should_exit = false; let fastboot_image = fs::read(args.image_path).expect("boot image not found"); @@ -60,9 +65,7 @@ async fn main() -> anyhow::Result<()> { .with_context(|| format!("Couldn't execute {CDBA_SERVER_BIN_NAME} on remote server"))?; let mut server_stdin = Arc::new(Mutex::new((*chan.lock().await).make_writer())); - let (server_stdout, server_stderr) = sk8brd::ssh::into_streams::(chan).await; - let server_stdout = Arc::new(Mutex::new(server_stdout)); - let server_stderr = Arc::new(Mutex::new(server_stderr)); + let (mut server_stdout, mut server_stderr) = sk8brd::ssh::into_streams::(chan).await; if args.board.is_empty() { send_ack(&mut server_stdin, Sk8brdMsgs::MsgListDevices).await?; @@ -71,67 +74,102 @@ async fn main() -> anyhow::Result<()> { } // Msg handler - // Read the message header first - while time.elapsed()? < Duration::from_secs(args.timeout) { - // Stream of "blue text" - status updates from the server - if let Ok(bytes_read) = (*server_stderr.lock().await).read(&mut buf).await { - let s = String::from_utf8_lossy(&buf[..bytes_read]); - print!( - "{}\r", - s.split('\n').collect::>().join("\r\n").blue() - ); - stdout().flush()?; - } + while Instant::now() < deadline { + tokio::select! { + _ = sleep_until(tokio::time::Instant::from_std(deadline)) => break, + + // Stream of "blue text" - status updates from the server + stderr_read = server_stderr.read(&mut stderr_buf), if stderr_open => { + if let Ok(bytes_read) = stderr_read { + if bytes_read == 0 { + stderr_open = false; + continue; + } - if (*server_stdout.lock().await) - .read_exact(&mut hdr_buf) - .await - .is_ok() - { - let msg = parse_recv_msg(&hdr_buf); - let mut msgbuf = vec![0u8; msg.len as usize]; - - // Now read the actual data... - (*server_stderr.lock().await) - .read_exact(&mut msgbuf) - .await?; - - // ..and process it - match msg.r#type.try_into() { - Ok(Sk8brdMsgs::MsgSelectBoard) => { - send_ack(&mut server_stdin, Sk8brdMsgs::MsgPowerOn).await? + let s = String::from_utf8_lossy(&stderr_buf[..bytes_read]); + print!( + "{}\r", + s.split('\n').collect::>().join("\r\n").blue() + ); + stdout().flush()?; } - Ok(Sk8brdMsgs::MsgConsole) => { - if args.verbose { - console_print(&msgbuf).await + } + + // Binary protocol stream on stdout + stdout_read = server_stdout.read(&mut stdout_chunk) => { + if let Ok(bytes_read) = stdout_read { + if bytes_read == 0 { + break; } - } - Ok(Sk8brdMsgs::MsgPowerOn) => { - // Refresh the timer so that the timeout actually makes sense - time = SystemTime::now(); - } - Ok(Sk8brdMsgs::MsgFastbootPresent) => { - if !msgbuf.is_empty() && msgbuf[0] != 0 { - send_image(&mut server_stdin, &fastboot_image, &quit).await? + + stdout_buf.extend_from_slice(&stdout_chunk[..bytes_read]); + + // Parse as many complete framed messages as available. + loop { + if stdout_buf.len() < MSG_HDR_SIZE { + break; + } + + let msg = parse_recv_msg(&stdout_buf[..MSG_HDR_SIZE]); + if Sk8brdMsgs::try_from(msg.r#type).is_err() || msg.len as usize > MAX_MSG_LEN { + // Resync in case stdout had unexpected text/noise. + stdout_buf.drain(..1); + continue; + } + + let total_len = MSG_HDR_SIZE + msg.len as usize; + if stdout_buf.len() < total_len { + break; + } + + let msgbuf = stdout_buf[MSG_HDR_SIZE..total_len].to_vec(); + stdout_buf.drain(..total_len); + match msg.r#type.try_into() { + Ok(Sk8brdMsgs::MsgSelectBoard) => { + send_ack(&mut server_stdin, Sk8brdMsgs::MsgPowerOn).await? + } + Ok(Sk8brdMsgs::MsgConsole) => { + if args.verbose { + console_print(&msgbuf).await + } + } + Ok(Sk8brdMsgs::MsgPowerOn) => { + // Refresh timeout window after power-on ack. + deadline = Instant::now() + Duration::from_secs(args.timeout); + } + Ok(Sk8brdMsgs::MsgFastbootPresent) => { + if !msgbuf.is_empty() && msgbuf[0] != 0 { + send_image(&mut server_stdin, &fastboot_image, &quit).await? + } + } + Ok(Sk8brdMsgs::MsgFastbootDownload) => (), + Ok(Sk8brdMsgs::MsgListDevices) => { + print_string_msg(&msgbuf); + if msgbuf.is_empty() { + should_exit = true; + break; + } + } + + // Ignore all other valid messages + Ok(_) => (), + Err(e) => todo!("Received unknown/invalid message: `{e}`"), + }; } - } - Ok(Sk8brdMsgs::MsgFastbootDownload) => (), - Ok(Sk8brdMsgs::MsgListDevices) => { - print_string_msg(&msgbuf); - if msgbuf.is_empty() { + if should_exit { break; } } - - // Ignore all other valid messages - Ok(_) => (), - Err(e) => todo!("Received unknown/invalid message: `{e}`"), - }; + } } } // Power off the board on goodbye - send_ack(&mut server_stdin, Sk8brdMsgs::MsgPowerOff).await?; + let _ = timeout( + Duration::from_secs(1), + send_ack(&mut server_stdin, Sk8brdMsgs::MsgPowerOff), + ) + .await; // ssh_disconnect(&mut sess).await?; diff --git a/client/src/main.rs b/client/src/main.rs index b0bb1f7..138bcef 100644 --- a/client/src/main.rs +++ b/client/src/main.rs @@ -12,6 +12,7 @@ use std::io::{stdout, Read, Write}; use std::sync::Arc; use tokio::io::{AsyncReadExt, AsyncWrite}; use tokio::sync::Mutex; +use tokio::time::timeout; macro_rules! get_arc { ($a: expr) => {{ @@ -79,8 +80,11 @@ async fn handle_keypress( #[allow(clippy::explicit_write)] #[tokio::main] async fn main() -> anyhow::Result<()> { - let mut hdr_buf = [0u8; MSG_HDR_SIZE]; - let mut buf = [0u8; SSH_BUFFER_SIZE]; + const MAX_MSG_LEN: usize = 1024 * 1024; + let mut stderr_buf = [0u8; SSH_BUFFER_SIZE]; + let mut stdout_chunk = [0u8; SSH_BUFFER_SIZE]; + let mut stdout_buf: Vec = Vec::new(); + let mut stderr_open = true; let mut key_buf = [0u8; 1]; let quit = Arc::new(Mutex::new(false)); let args = Args::parse(); @@ -99,9 +103,7 @@ async fn main() -> anyhow::Result<()> { let mut server_stdin = Arc::new(Mutex::new(get_arc!(chan).make_writer())); - let (server_stdout, server_stderr) = sk8brd::ssh::into_streams::(chan).await; - let server_stdout = Arc::new(Mutex::new(server_stdout)); - let server_stderr = Arc::new(Mutex::new(server_stderr)); + let (mut server_stdout, mut server_stderr) = sk8brd::ssh::into_streams::(chan).await; send_ack(&mut server_stdin, Sk8brdMsgs::MsgListDevices).await?; select_brd(&mut server_stdin, &args.board).await?; @@ -134,57 +136,83 @@ async fn main() -> anyhow::Result<()> { }); while !*get_arc!(quit) { - // Stream of "blue text" - status updates from the server - if let Ok(bytes_read) = (*get_arc!(server_stderr)).read(&mut buf).await { - let s = String::from_utf8_lossy(&buf[..bytes_read]); - print!( - "{}\r", - s.split('\n').collect::>().join("\r\n").blue() - ); - stdout().flush()?; - } + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => (), + + // Stream of "blue text" - status updates from the server + stderr_read = server_stderr.read(&mut stderr_buf), if stderr_open => { + match stderr_read { + Ok(0) => stderr_open = false, + Ok(bytes_read) => { + let s = String::from_utf8_lossy(&stderr_buf[..bytes_read]); + print!( + "{}\r", + s.split('\n').collect::>().join("\r\n").blue() + ); + stdout().flush()?; + } + Err(_) => stderr_open = false, + } + } - // Msg handler - // Read the message header first - if (*get_arc!(server_stdout)) - .read_exact(&mut hdr_buf) - .await - .is_ok() - { - let msg = parse_recv_msg(&hdr_buf); - let mut msgbuf = vec![0u8; msg.len as usize]; - - // Now read the actual data... - (*get_arc!(server_stdout)).read_exact(&mut msgbuf).await?; - - // ..and process it - match msg.r#type.try_into() { - Ok(Sk8brdMsgs::MsgSelectBoard) => { - send_msg(&mut server_stdin, Sk8brdMsgs::MsgPowerOn, &[]).await? + // Binary protocol stream on stdout + stdout_read = server_stdout.read(&mut stdout_chunk) => { + let bytes_read = stdout_read?; + if bytes_read == 0 { + break; } - Ok(Sk8brdMsgs::MsgConsole) => console_print(&msgbuf).await, - Ok(Sk8brdMsgs::MsgHardReset) => todo!("MsgHardReset is unused"), - Ok(Sk8brdMsgs::MsgPowerOn) => (), - Ok(Sk8brdMsgs::MsgPowerOff) => (), - Ok(Sk8brdMsgs::MsgFastbootPresent) => { - if !msgbuf.is_empty() && msgbuf[0] != 0 { - send_image(&mut server_stdin, &fastboot_image, &quit).await? + + stdout_buf.extend_from_slice(&stdout_chunk[..bytes_read]); + + // Parse as many complete framed messages as available. + loop { + if stdout_buf.len() < MSG_HDR_SIZE { + break; } + + let msg = parse_recv_msg(&stdout_buf[..MSG_HDR_SIZE]); + if Sk8brdMsgs::try_from(msg.r#type).is_err() || msg.len as usize > MAX_MSG_LEN { + // Resync in case stdout had unexpected text/noise. + stdout_buf.drain(..1); + continue; + } + + let total_len = MSG_HDR_SIZE + msg.len as usize; + if stdout_buf.len() < total_len { + break; + } + + let msgbuf = stdout_buf[MSG_HDR_SIZE..total_len].to_vec(); + stdout_buf.drain(..total_len); + match msg.r#type.try_into() { + Ok(Sk8brdMsgs::MsgSelectBoard) => { + send_msg(&mut server_stdin, Sk8brdMsgs::MsgPowerOn, &[]).await? + } + Ok(Sk8brdMsgs::MsgConsole) => console_print(&msgbuf).await, + Ok(Sk8brdMsgs::MsgHardReset) => todo!("MsgHardReset is unused"), + Ok(Sk8brdMsgs::MsgPowerOn) => (), + Ok(Sk8brdMsgs::MsgPowerOff) => (), + Ok(Sk8brdMsgs::MsgFastbootPresent) => { + if !msgbuf.is_empty() && msgbuf[0] != 0 { + send_image(&mut server_stdin, &fastboot_image, &quit).await? + } + } + Ok(Sk8brdMsgs::MsgFastbootDownload) => (), + Ok(Sk8brdMsgs::MsgFastbootBoot) => todo!("MsgFastbootBoot is unused"), + Ok(Sk8brdMsgs::MsgStatusUpdate) => todo!("MsgStatusUpdate: implement me!"), + Ok(Sk8brdMsgs::MsgVbusOn) => todo!("Unexpected MsgVbusOn"), + Ok(Sk8brdMsgs::MsgVbusOff) => todo!("Unexpected MsgVbusOff"), + Ok(Sk8brdMsgs::MsgFastbootReboot) => todo!("MsgFastbootReboot is unused"), + Ok(Sk8brdMsgs::MsgSendBreak) => todo!("MsgSendBreak: implement me!"), + Ok(Sk8brdMsgs::MsgListDevices) => print_string_msg(&msgbuf), + Ok(Sk8brdMsgs::MsgBoardInfo) => print_string_msg(&msgbuf), + Ok(Sk8brdMsgs::MsgFastbootContinue) => (), + + Ok(m) => todo!("{m:?} is unimplemented, skipping.."), + Err(e) => todo!("Received unknown/invalid message: `{e}`"), + }; } - Ok(Sk8brdMsgs::MsgFastbootDownload) => (), - Ok(Sk8brdMsgs::MsgFastbootBoot) => todo!("MsgFastbootBoot is unused"), - Ok(Sk8brdMsgs::MsgStatusUpdate) => todo!("MsgStatusUpdate: implement me!"), - Ok(Sk8brdMsgs::MsgVbusOn) => todo!("Unexpected MsgVbusOn"), - Ok(Sk8brdMsgs::MsgVbusOff) => todo!("Unexpected MsgVbusOff"), - Ok(Sk8brdMsgs::MsgFastbootReboot) => todo!("MsgFastbootReboot is unused"), - Ok(Sk8brdMsgs::MsgSendBreak) => todo!("MsgSendBreak: implement me!"), - Ok(Sk8brdMsgs::MsgListDevices) => print_string_msg(&msgbuf), - Ok(Sk8brdMsgs::MsgBoardInfo) => print_string_msg(&msgbuf), - Ok(Sk8brdMsgs::MsgFastbootContinue) => (), - - Ok(m) => todo!("{m:?} is unimplemented, skipping.."), - Err(e) => todo!("Received unknown/invalid message: `{e}`"), - }; + } } } @@ -195,7 +223,11 @@ async fn main() -> anyhow::Result<()> { crossterm::terminal::disable_raw_mode()?; // Power off the board on goodbye - send_ack(&mut server_stdin, Sk8brdMsgs::MsgPowerOff).await?; + let _ = timeout( + std::time::Duration::from_secs(1), + send_ack(&mut server_stdin, Sk8brdMsgs::MsgPowerOff), + ) + .await; // ssh_disconnect(&mut sess).await?; diff --git a/proto/src/ssh.rs b/proto/src/ssh.rs index e63a33b..23b3d75 100644 --- a/proto/src/ssh.rs +++ b/proto/src/ssh.rs @@ -171,13 +171,7 @@ where Some(russh::ChannelMsg::ExtendedData { data: _, ext }) => { println!("Received surprise data on stream {ext}"); } - Some(russh::ChannelMsg::Eof) => { - // Send a 0-length chunk to indicate EOF. - txo.send(vec![]) - .await - .map_err(|_| russh::Error::SendError)?; - break; - } + Some(russh::ChannelMsg::Eof) => break, None => break, _ => (), } @@ -196,23 +190,23 @@ impl AsyncRead for Wrap { cx: &mut Context<'_>, buf: &mut tokio::io::ReadBuf<'_>, ) -> Poll> { - let cache_size = self.1.len(); - buf.put_slice(&self.1.split_to(usize::min(buf.remaining(), cache_size))); - - if buf.remaining() > 0 { - match self.0.poll_recv(cx) { - Poll::Ready(Some(msg)) => { - self.1 = BytesMut::from(&msg[..]); - let len = self.1.len(); - buf.put_slice(&self.1.split_to(usize::min(buf.remaining(), len))); - Poll::Ready(Ok(())) - } + if !self.1.is_empty() { + let n = usize::min(buf.remaining(), self.1.len()); + buf.put_slice(&self.1.split_to(n)); + return Poll::Ready(Ok(())); + } - Poll::Ready(None) => Poll::Ready(Ok(())), - Poll::Pending => Poll::Pending, + match self.0.poll_recv(cx) { + Poll::Ready(Some(msg)) => { + self.1 = BytesMut::from(&msg[..]); + if !self.1.is_empty() { + let n = usize::min(buf.remaining(), self.1.len()); + buf.put_slice(&self.1.split_to(n)); + } + Poll::Ready(Ok(())) } - } else { - Poll::Ready(Ok(())) + Poll::Ready(None) => Poll::Ready(Ok(())), + Poll::Pending => Poll::Pending, } } } From 3977c901da71697f7198d2f0f9edf861ca75015e Mon Sep 17 00:00:00 2001 From: Bjorn Andersson Date: Fri, 3 Jul 2026 03:16:10 +0000 Subject: [PATCH 3/5] proto: reduce per-frame overhead during image upload Image upload sends every fastboot payload as a separate 2 KiB protocol message. With the russh channel writer this creates a large number of small async writes and mutex acquisitions for large Android boot images. Use 8 KiB image frames, which remain safely below the cdba-server 16 KiB receive ring, and hold the writer lock for the duration of an image transfer. Signed-off-by: Bjorn Andersson --- proto/src/lib.rs | 33 ++++++++++++++++++++++++++++----- 1 file changed, 28 insertions(+), 5 deletions(-) diff --git a/proto/src/lib.rs b/proto/src/lib.rs index 0e448da..262480a 100644 --- a/proto/src/lib.rs +++ b/proto/src/lib.rs @@ -67,6 +67,7 @@ pub struct Sk8brdMsg { pub len: u16, } pub const MSG_HDR_SIZE: usize = size_of::(); +const IMAGE_CHUNK_SIZE: usize = 8 * 1024; pub async fn send_msg( write_sink: &mut Arc>, @@ -75,9 +76,16 @@ pub async fn send_msg( ) -> anyhow::Result<()> { // Make sure we're not trying to send two messages at once let mut write_sink = write_sink.lock().await; + write_msg(&mut *write_sink, r#type, buf).await +} - let len = buf.len(); - let hdr = [r#type as u8, (len & 0xff) as u8, ((len >> 8) & 0xff) as u8]; +async fn write_msg( + write_sink: &mut (impl AsyncWrite + std::marker::Unpin), + r#type: Sk8brdMsgs, + buf: &[u8], +) -> anyhow::Result<()> { + let len = u16::try_from(buf.len())?; + let hdr = [r#type as u8, (len & 0xff) as u8, (len >> 8) as u8]; write_sink.write_all(&hdr).await?; write_sink.write_all(buf).await?; @@ -116,7 +124,9 @@ pub async fn send_image( let mut last_percent_done: usize = 0; let mut bytes_sent = 0; - for chunk in buf.chunks(2048) { + let mut write_sink = write_sink.lock().await; + + for chunk in buf.chunks(IMAGE_CHUNK_SIZE) { let percent_done = 100 * bytes_sent / buf.len(); if *quit.lock().await { @@ -129,7 +139,7 @@ pub async fn send_image( stdout().flush()?; } - send_msg(write_sink, Sk8brdMsgs::MsgFastbootDownload, chunk).await?; + write_msg(&mut *write_sink, Sk8brdMsgs::MsgFastbootDownload, chunk).await?; bytes_sent += chunk.len(); last_percent_done = percent_done; @@ -141,7 +151,20 @@ pub async fn send_image( } } - send_ack(write_sink, Sk8brdMsgs::MsgFastbootDownload).await + write_msg(&mut *write_sink, Sk8brdMsgs::MsgFastbootDownload, &[]).await +} + +pub async fn send_image_quiet( + write_sink: &mut Arc>, + buf: &[u8], +) -> anyhow::Result<()> { + let mut write_sink = write_sink.lock().await; + + for chunk in buf.chunks(IMAGE_CHUNK_SIZE) { + write_msg(&mut *write_sink, Sk8brdMsgs::MsgFastbootDownload, chunk).await?; + } + + write_msg(&mut *write_sink, Sk8brdMsgs::MsgFastbootDownload, &[]).await } pub async fn select_brd( From a9584dc7798a704eeff8346d059c98c44dcaa52a Mon Sep 17 00:00:00 2001 From: Bjorn Andersson Date: Fri, 3 Jul 2026 01:51:14 +0000 Subject: [PATCH 4/5] python: expose sk8brd as a Python module sk8brd allowed us to integrate cdba into Rust workflows, without having to shell out and deal with CLI parsing and process management. In the same way there are use cases for being able to integrate cdba workflows directly into Python tools. Add a PyO3 extension crate that builds the sk8brd_cdba module with maturin. Expose blocking Python methods for listing devices, booting an image, and toggling board power while Rust keeps the existing async protocol handling. Signed-off-by: Bjorn Andersson --- Cargo.lock | 80 ++++++ Cargo.toml | 2 +- README.md | 10 +- proto/src/lib.rs | 4 +- proto/src/ssh.rs | 2 +- pyproject.toml | 14 + python/Cargo.toml | 20 ++ python/README.md | 20 ++ python/src/lib.rs | 705 ++++++++++++++++++++++++++++++++++++++++++++++ 9 files changed, 852 insertions(+), 5 deletions(-) create mode 100644 pyproject.toml create mode 100644 python/Cargo.toml create mode 100644 python/README.md create mode 100644 python/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index 78e9051..4695480 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1209,6 +1209,12 @@ dependencies = [ "universal-hash", ] +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -1236,6 +1242,63 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "pyo3" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" +dependencies = [ + "libc", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "quote" version = "1.0.45" @@ -1630,6 +1693,17 @@ dependencies = [ "use", ] +[[package]] +name = "sk8brd-python" +version = "0.1.0" +dependencies = [ + "anyhow", + "pyo3", + "russh", + "sk8brd-proto", + "tokio", +] + [[package]] name = "slab" version = "0.4.12" @@ -1720,6 +1794,12 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + [[package]] name = "thiserror" version = "1.0.69" diff --git a/Cargo.toml b/Cargo.toml index a47a0c1..e6e01eb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,3 @@ [workspace] -members = ["client", "cli", "proto"] +members = ["client", "cli", "proto", "python"] resolver = "2" diff --git a/README.md b/README.md index 602b6be..5ae6b2e 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,14 @@ Keybinds: Make sure `ssh-agent` is running and has your keys imported. +### Python module: +The `sk8brd-cdba` package exposes the sk8brd cdba client as `sk8brd_cdba`: + +```sh +python -m pip install maturin +maturin develop +``` + ## License `BSD-3-Clause` @@ -30,4 +38,4 @@ Make sure `ssh-agent` is running and has your keys imported. ``` Author: Konrad Dybcio cdba contributors (for the original cdba implementation) -``` \ No newline at end of file +``` diff --git a/proto/src/lib.rs b/proto/src/lib.rs index 262480a..cf62127 100644 --- a/proto/src/lib.rs +++ b/proto/src/lib.rs @@ -67,7 +67,7 @@ pub struct Sk8brdMsg { pub len: u16, } pub const MSG_HDR_SIZE: usize = size_of::(); -const IMAGE_CHUNK_SIZE: usize = 8 * 1024; +pub const IMAGE_CHUNK_SIZE: usize = 8 * 1024; pub async fn send_msg( write_sink: &mut Arc>, @@ -79,7 +79,7 @@ pub async fn send_msg( write_msg(&mut *write_sink, r#type, buf).await } -async fn write_msg( +pub async fn write_msg( write_sink: &mut (impl AsyncWrite + std::marker::Unpin), r#type: Sk8brdMsgs, buf: &[u8], diff --git a/proto/src/ssh.rs b/proto/src/ssh.rs index 23b3d75..fab0fc1 100644 --- a/proto/src/ssh.rs +++ b/proto/src/ssh.rs @@ -109,7 +109,7 @@ pub async fn ssh_connect(farm: &str, username: String) -> anyhow::Result=1,<2"] +build-backend = "maturin" + +[project] +name = "sk8brd-cdba" +version = "0.1.0" +description = "Python bindings for the sk8brd cdba client" +requires-python = ">=3.9" + +[tool.maturin] +manifest-path = "python/Cargo.toml" +module-name = "sk8brd_cdba" +features = ["pyo3/extension-module"] diff --git a/python/Cargo.toml b/python/Cargo.toml new file mode 100644 index 0000000..313114d --- /dev/null +++ b/python/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "sk8brd-python" +version = "0.1.0" +edition = "2021" +authors = ["Konrad Dybcio "] +license = "BSD-3-Clause" +description = "Python bindings for the sk8brd cdba client" +repository = "https://github.com/linux-msm/sk8brd" +publish = false + +[lib] +name = "sk8brd_cdba" +crate-type = ["cdylib"] + +[dependencies] +anyhow = "1.0" +pyo3 = "0.29.0" +russh = "0.50.4" +sk8brd-proto = { path = "../proto", features = ["ssh"] } +tokio = { version = "1.43.0", features = ["full"] } diff --git a/python/README.md b/python/README.md new file mode 100644 index 0000000..91a9987 --- /dev/null +++ b/python/README.md @@ -0,0 +1,20 @@ +# sk8brd-cdba Python module + +This crate exposes the sk8brd cdba client as a Python module named +`sk8brd_cdba`. + +Build it in a virtual environment with: + +```sh +python -m pip install maturin +maturin develop +``` + +Then run the example application: + +```sh +SK8BRD_FARM= \ +SK8BRD_BOARD= \ +SK8BRD_BOOT_IMAGE= \ +python examples/python_cdba/boot_android.py +``` diff --git a/python/src/lib.rs b/python/src/lib.rs new file mode 100644 index 0000000..8ca4949 --- /dev/null +++ b/python/src/lib.rs @@ -0,0 +1,705 @@ +use anyhow::{bail, Context as _}; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use russh::client::Msg; +use sk8brd::{ + parse_recv_msg, select_brd, send_ack, send_console, send_image_quiet, write_msg, Sk8brdMsgs, + CDBA_SERVER_BIN_NAME, IMAGE_CHUNK_SIZE, MSG_HDR_SIZE, +}; +use std::fs; +use std::future::Future; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::mpsc::{self, Receiver, RecvTimeoutError, Sender, TryRecvError}; +use std::sync::{Arc, Mutex as StdMutex}; +use std::thread::{self, JoinHandle}; +use std::time::Duration; +use tokio::io::{AsyncReadExt, AsyncWrite}; +use tokio::sync::Mutex as TokioMutex; +use tokio::time::{sleep_until, timeout, Instant}; + +const MAX_MSG_LEN: usize = 1024 * 1024; +const STREAM_CHUNK_SIZE: usize = 2048; +const SESSION_POLL_INTERVAL: Duration = Duration::from_millis(10); + +#[pyclass(skip_from_py_object)] +#[derive(Clone)] +struct CdbaClient { + host: String, + port: u16, + user: String, + timeout_secs: u64, +} + +#[pymethods] +impl CdbaClient { + #[new] + #[pyo3(signature = (host, port = 22, user = "cdba".to_string(), timeout_secs = 60))] + fn new(host: String, port: u16, user: String, timeout_secs: u64) -> Self { + Self { + host, + port, + user, + timeout_secs, + } + } + + #[getter] + fn host(&self) -> &str { + &self.host + } + + #[getter] + fn port(&self) -> u16 { + self.port + } + + #[getter] + fn user(&self) -> &str { + &self.user + } + + fn list_devices(&self, py: Python<'_>) -> PyResult { + let client = self.clone(); + py.detach(move || run_blocking(client.list_devices_async())) + } + + #[pyo3(signature = (board, image_path, timeout_secs = None, collect_console = false))] + fn boot_image( + &self, + py: Python<'_>, + board: String, + image_path: String, + timeout_secs: Option, + collect_console: bool, + ) -> PyResult { + let mut client = self.clone(); + if let Some(timeout_secs) = timeout_secs { + client.timeout_secs = timeout_secs; + } + + py.detach(move || run_blocking(client.boot_image_async(board, image_path, collect_console))) + } + + #[pyo3(signature = (board, image_path, timeout_secs = None))] + fn boot_image_session( + &self, + py: Python<'_>, + board: String, + image_path: String, + timeout_secs: Option, + ) -> PyResult { + let mut client = self.clone(); + if let Some(timeout_secs) = timeout_secs { + client.timeout_secs = timeout_secs; + } + + client.boot_image_session_blocking(py, board, image_path) + } + + fn power_off(&self, py: Python<'_>, board: String) -> PyResult<()> { + let client = self.clone(); + py.detach(move || run_blocking(client.send_board_ack_async(board, Sk8brdMsgs::MsgPowerOff))) + } + + fn power_on(&self, py: Python<'_>, board: String) -> PyResult<()> { + let client = self.clone(); + py.detach(move || run_blocking(client.send_board_ack_async(board, Sk8brdMsgs::MsgPowerOn))) + } +} + +#[pyclass] +struct BootResult { + #[pyo3(get)] + image_sent: bool, + #[pyo3(get)] + console: String, + #[pyo3(get)] + status: String, +} + +#[pymethods] +impl BootResult { + fn __repr__(&self) -> String { + format!( + "BootResult(image_sent={}, console_len={}, status_len={})", + self.image_sent, + self.console.len(), + self.status.len() + ) + } +} + +#[pyclass] +struct CdbaSession { + shared: Arc, + tx: Sender, + worker: StdMutex>>, +} + +#[pymethods] +impl CdbaSession { + #[getter] + fn image_sent(&self) -> bool { + self.shared.image_sent.load(Ordering::SeqCst) + } + + #[getter] + fn closed(&self) -> bool { + self.shared.closed.load(Ordering::SeqCst) + } + + #[pyo3(signature = (max_bytes = None))] + fn read_console(&self, max_bytes: Option) -> PyResult { + self.shared.check_error()?; + drain_buffer(&self.shared.console, max_bytes) + } + + #[pyo3(signature = (max_bytes = None))] + fn read_status(&self, max_bytes: Option) -> PyResult { + self.shared.check_error()?; + drain_buffer(&self.shared.status, max_bytes) + } + + fn write_console(&self, data: String) -> PyResult<()> { + self.shared.check_error()?; + self.tx + .send(SessionCommand::Write(data.into_bytes())) + .map_err(py_runtime_err) + } + + fn power_off(&self) -> PyResult<()> { + self.shared.check_error()?; + self.tx + .send(SessionCommand::PowerOff) + .map_err(py_runtime_err) + } + + fn close(&self, py: Python<'_>) -> PyResult<()> { + let _ = self.tx.send(SessionCommand::Close); + py.detach(|| self.join_worker()) + } + + fn __repr__(&self) -> String { + format!( + "CdbaSession(image_sent={}, closed={}, console_len={}, status_len={})", + self.image_sent(), + self.closed(), + self.shared.console.lock().unwrap().len(), + self.shared.status.lock().unwrap().len() + ) + } +} + +impl Drop for CdbaSession { + fn drop(&mut self) { + let _ = self.tx.send(SessionCommand::Close); + } +} + +impl CdbaSession { + fn join_worker(&self) -> PyResult<()> { + if let Some(worker) = self.worker.lock().unwrap().take() { + worker + .join() + .map_err(|_| PyRuntimeError::new_err("cdba session worker panicked"))?; + } + + self.shared.check_error() + } +} + +struct SessionShared { + console: StdMutex, + status: StdMutex, + error: StdMutex>, + image_sent: AtomicBool, + closed: AtomicBool, +} + +impl SessionShared { + fn new() -> Self { + Self { + console: StdMutex::new(String::new()), + status: StdMutex::new(String::new()), + error: StdMutex::new(None), + image_sent: AtomicBool::new(false), + closed: AtomicBool::new(false), + } + } + + fn append_console(&self, data: &[u8]) { + self.console + .lock() + .unwrap() + .push_str(&String::from_utf8_lossy(data)); + } + + fn append_status(&self, data: &[u8]) { + self.status + .lock() + .unwrap() + .push_str(&String::from_utf8_lossy(data)); + } + + fn fail(&self, err: impl std::fmt::Display) { + *self.error.lock().unwrap() = Some(err.to_string()); + } + + fn check_error(&self) -> PyResult<()> { + match self.error.lock().unwrap().as_ref() { + Some(err) => Err(PyRuntimeError::new_err(err.clone())), + None => Ok(()), + } + } +} + +enum SessionCommand { + Write(Vec), + PowerOff, + Close, +} + +impl CdbaClient { + fn address(&self) -> String { + format!("{}:{}", self.host, self.port) + } + + async fn list_devices_async(self) -> anyhow::Result { + let chan = Arc::new(TokioMutex::new( + sk8brd::ssh::ssh_connect(&self.address(), self.user.clone()).await?, + )); + (*chan.lock().await) + .exec(true, CDBA_SERVER_BIN_NAME) + .await + .with_context(|| { + format!("could not execute {CDBA_SERVER_BIN_NAME} on remote server") + })?; + let mut server_stdin = Arc::new(TokioMutex::new((*chan.lock().await).make_writer())); + let (mut server_stdout, mut server_stderr) = sk8brd::ssh::into_streams::(chan).await; + send_ack(&mut server_stdin, Sk8brdMsgs::MsgListDevices).await?; + + let deadline = Instant::now() + Duration::from_secs(self.timeout_secs); + let mut stdout_buf = Vec::new(); + let mut stdout_chunk = [0u8; STREAM_CHUNK_SIZE]; + let mut stderr_chunk = [0u8; STREAM_CHUNK_SIZE]; + let mut stderr_open = true; + let mut devices = String::new(); + + while Instant::now() < deadline { + tokio::select! { + _ = sleep_until(deadline) => break, + read = server_stderr.read(&mut stderr_chunk), if stderr_open => { + if let Ok(bytes_read) = read { + stderr_open = bytes_read != 0; + } + } + read = server_stdout.read(&mut stdout_chunk) => { + let bytes_read = read?; + if bytes_read == 0 { + break; + } + + stdout_buf.extend_from_slice(&stdout_chunk[..bytes_read]); + while let Some((msg, payload)) = next_frame(&mut stdout_buf)? { + if msg == Sk8brdMsgs::MsgListDevices { + devices.push_str(&String::from_utf8_lossy(&payload)); + return Ok(devices); + } + } + } + } + } + + bail!("timed out waiting for device list") + } + + async fn boot_image_async( + self, + board: String, + image_path: String, + collect_console: bool, + ) -> anyhow::Result { + let image = + fs::read(&image_path).with_context(|| format!("could not read {image_path}"))?; + let chan = Arc::new(TokioMutex::new( + sk8brd::ssh::ssh_connect(&self.address(), self.user.clone()).await?, + )); + (*chan.lock().await) + .exec(true, CDBA_SERVER_BIN_NAME) + .await + .with_context(|| { + format!("could not execute {CDBA_SERVER_BIN_NAME} on remote server") + })?; + let mut server_stdin = Arc::new(TokioMutex::new((*chan.lock().await).make_writer())); + let (mut server_stdout, mut server_stderr) = sk8brd::ssh::into_streams::(chan).await; + select_brd(&mut server_stdin, &board).await?; + + let mut deadline = Instant::now() + Duration::from_secs(self.timeout_secs); + let mut stdout_buf = Vec::new(); + let mut stdout_chunk = [0u8; STREAM_CHUNK_SIZE]; + let mut stderr_chunk = [0u8; STREAM_CHUNK_SIZE]; + let mut stderr_open = true; + let mut image_sent = false; + let mut console = String::new(); + let mut status = String::new(); + + while Instant::now() < deadline { + tokio::select! { + _ = sleep_until(deadline) => break, + read = server_stderr.read(&mut stderr_chunk), if stderr_open => { + let bytes_read = read?; + if bytes_read == 0 { + stderr_open = false; + } else { + status.push_str(&String::from_utf8_lossy(&stderr_chunk[..bytes_read])); + } + } + read = server_stdout.read(&mut stdout_chunk) => { + let bytes_read = read?; + if bytes_read == 0 { + break; + } + + stdout_buf.extend_from_slice(&stdout_chunk[..bytes_read]); + while let Some((msg, payload)) = next_frame(&mut stdout_buf)? { + match msg { + Sk8brdMsgs::MsgSelectBoard => { + send_ack(&mut server_stdin, Sk8brdMsgs::MsgPowerOn).await?; + } + Sk8brdMsgs::MsgPowerOn => { + deadline = Instant::now() + Duration::from_secs(self.timeout_secs); + } + Sk8brdMsgs::MsgFastbootPresent => { + if !payload.is_empty() && payload[0] != 0 { + send_image_quiet(&mut server_stdin, &image).await?; + image_sent = true; + } + } + Sk8brdMsgs::MsgConsole + if collect_console => { + console.push_str(&String::from_utf8_lossy(&payload)); + } + _ => (), + } + } + } + } + } + + if !image_sent { + bail!("timed out waiting for fastboot image transfer"); + } + + Ok(BootResult { + image_sent, + console, + status, + }) + } + + fn boot_image_session_blocking( + self, + py: Python<'_>, + board: String, + image_path: String, + ) -> PyResult { + let image = fs::read(&image_path) + .with_context(|| format!("could not read {image_path}")) + .map_err(py_runtime_err)?; + let shared = Arc::new(SessionShared::new()); + let (tx, rx) = mpsc::channel(); + let (ready_tx, ready_rx) = mpsc::channel(); + let worker_shared = shared.clone(); + let ready_tx_for_error = ready_tx.clone(); + let startup_timeout = Duration::from_secs(self.timeout_secs); + + let worker = thread::spawn(move || { + let result = run_blocking_result(session_worker( + self, + board, + image, + worker_shared.clone(), + rx, + ready_tx, + )); + + if let Err(err) = result { + worker_shared.fail(&err); + let _ = ready_tx_for_error.send(Err(err.to_string())); + } + + worker_shared.closed.store(true, Ordering::SeqCst); + }); + + let startup_deadline = std::time::Instant::now() + startup_timeout; + loop { + if let Err(err) = py.check_signals() { + let _ = tx.send(SessionCommand::Close); + let _ = worker.join(); + return Err(err); + } + + match ready_rx.recv_timeout(Duration::from_millis(50)) { + Ok(Ok(())) => { + return Ok(CdbaSession { + shared, + tx, + worker: StdMutex::new(Some(worker)), + }); + } + Ok(Err(err)) => { + let _ = worker.join(); + return Err(PyRuntimeError::new_err(err)); + } + Err(RecvTimeoutError::Timeout) if std::time::Instant::now() < startup_deadline => {} + Err(err) => { + let _ = tx.send(SessionCommand::Close); + let _ = worker.join(); + return Err(PyRuntimeError::new_err(format!( + "timed out waiting for image upload: {err}" + ))); + } + } + } + } + + async fn send_board_ack_async(self, board: String, msg: Sk8brdMsgs) -> anyhow::Result<()> { + let chan = Arc::new(TokioMutex::new( + sk8brd::ssh::ssh_connect(&self.address(), self.user.clone()).await?, + )); + (*chan.lock().await) + .exec(true, CDBA_SERVER_BIN_NAME) + .await + .with_context(|| { + format!("could not execute {CDBA_SERVER_BIN_NAME} on remote server") + })?; + let mut server_stdin = Arc::new(TokioMutex::new((*chan.lock().await).make_writer())); + select_brd(&mut server_stdin, &board).await?; + send_ack(&mut server_stdin, msg).await + } +} + +async fn session_worker( + client: CdbaClient, + board: String, + image: Vec, + shared: Arc, + rx: Receiver, + ready_tx: Sender>, +) -> anyhow::Result<()> { + let chan = Arc::new(TokioMutex::new( + sk8brd::ssh::ssh_connect(&client.address(), client.user.clone()).await?, + )); + (*chan.lock().await) + .exec(true, CDBA_SERVER_BIN_NAME) + .await + .with_context(|| format!("could not execute {CDBA_SERVER_BIN_NAME} on remote server"))?; + let mut server_stdin = Arc::new(TokioMutex::new((*chan.lock().await).make_writer())); + let (mut server_stdout, mut server_stderr) = sk8brd::ssh::into_streams::(chan).await; + select_brd(&mut server_stdin, &board).await?; + + let mut deadline = Instant::now() + Duration::from_secs(client.timeout_secs); + let mut stdout_buf = Vec::new(); + let mut stdout_chunk = [0u8; STREAM_CHUNK_SIZE]; + let mut stderr_chunk = [0u8; STREAM_CHUNK_SIZE]; + let mut stderr_open = true; + let mut image_sent = false; + let mut ready_tx = Some(ready_tx); + + while Instant::now() < deadline || image_sent { + tokio::select! { + _ = sleep_until(deadline), if !image_sent => break, + _ = tokio::time::sleep(SESSION_POLL_INTERVAL) => { + if handle_session_commands(&mut server_stdin, &rx).await? { + return Ok(()); + } + } + read = server_stderr.read(&mut stderr_chunk), if stderr_open => { + let bytes_read = read?; + if bytes_read == 0 { + stderr_open = false; + } else { + shared.append_status(&stderr_chunk[..bytes_read]); + } + } + read = server_stdout.read(&mut stdout_chunk) => { + let bytes_read = read?; + if bytes_read == 0 { + break; + } + + stdout_buf.extend_from_slice(&stdout_chunk[..bytes_read]); + while let Some((msg, payload)) = next_frame(&mut stdout_buf)? { + match msg { + Sk8brdMsgs::MsgSelectBoard => { + send_ack(&mut server_stdin, Sk8brdMsgs::MsgPowerOn).await?; + } + Sk8brdMsgs::MsgPowerOn => { + deadline = Instant::now() + Duration::from_secs(client.timeout_secs); + } + Sk8brdMsgs::MsgFastbootPresent => { + if !payload.is_empty() && payload[0] != 0 { + if send_session_image(&mut server_stdin, &image, &rx).await? { + return Ok(()); + } + image_sent = true; + shared.image_sent.store(true, Ordering::SeqCst); + if let Some(ready_tx) = ready_tx.take() { + let _ = ready_tx.send(Ok(())); + } + deadline = Instant::now() + Duration::from_secs(client.timeout_secs); + } + } + Sk8brdMsgs::MsgConsole => shared.append_console(&payload), + _ => (), + } + } + } + } + } + + let _ = timeout( + Duration::from_secs(1), + send_ack(&mut server_stdin, Sk8brdMsgs::MsgPowerOff), + ) + .await; + + if !image_sent { + bail!("timed out waiting for fastboot image transfer"); + } + + Ok(()) +} + +async fn send_session_image( + server_stdin: &mut Arc>, + image: &[u8], + rx: &Receiver, +) -> anyhow::Result { + let mut server_stdin = server_stdin.lock().await; + + for chunk in image.chunks(IMAGE_CHUNK_SIZE) { + match rx.try_recv() { + Ok(SessionCommand::Close) | Err(TryRecvError::Disconnected) => { + let _ = write_msg(&mut *server_stdin, Sk8brdMsgs::MsgPowerOff, &[]).await; + return Ok(true); + } + Ok(SessionCommand::PowerOff) => { + write_msg(&mut *server_stdin, Sk8brdMsgs::MsgPowerOff, &[]).await?; + } + Ok(SessionCommand::Write(_)) | Err(TryRecvError::Empty) => (), + } + + write_msg(&mut *server_stdin, Sk8brdMsgs::MsgFastbootDownload, chunk).await?; + } + + write_msg(&mut *server_stdin, Sk8brdMsgs::MsgFastbootDownload, &[]).await?; + Ok(false) +} + +async fn handle_session_commands( + server_stdin: &mut Arc>, + rx: &Receiver, +) -> anyhow::Result { + loop { + match rx.try_recv() { + Ok(SessionCommand::Write(data)) => send_console(server_stdin, &data).await?, + Ok(SessionCommand::PowerOff) => send_ack(server_stdin, Sk8brdMsgs::MsgPowerOff).await?, + Ok(SessionCommand::Close) => { + let _ = timeout( + Duration::from_secs(1), + send_ack(server_stdin, Sk8brdMsgs::MsgPowerOff), + ) + .await; + return Ok(true); + } + Err(TryRecvError::Empty) => return Ok(false), + Err(TryRecvError::Disconnected) => { + let _ = timeout( + Duration::from_secs(1), + send_ack(server_stdin, Sk8brdMsgs::MsgPowerOff), + ) + .await; + return Ok(true); + } + } + } +} + +fn run_blocking(future: F) -> PyResult +where + F: Future>, +{ + run_blocking_result(future).map_err(py_runtime_err) +} + +fn run_blocking_result(future: F) -> anyhow::Result +where + F: Future>, +{ + let runtime = tokio::runtime::Runtime::new()?; + runtime.block_on(future) +} + +fn next_frame(buf: &mut Vec) -> anyhow::Result)>> { + loop { + if buf.len() < MSG_HDR_SIZE { + return Ok(None); + } + + let msg = parse_recv_msg(&buf[..MSG_HDR_SIZE]); + let Ok(msg_type) = Sk8brdMsgs::try_from(msg.r#type) else { + buf.drain(..1); + continue; + }; + + let payload_len = msg.len as usize; + if payload_len > MAX_MSG_LEN { + buf.drain(..1); + continue; + } + + let total_len = MSG_HDR_SIZE + payload_len; + if buf.len() < total_len { + return Ok(None); + } + + let payload = buf[MSG_HDR_SIZE..total_len].to_vec(); + buf.drain(..total_len); + return Ok(Some((msg_type, payload))); + } +} + +fn drain_buffer(buf: &StdMutex, max_bytes: Option) -> PyResult { + let mut buf = buf.lock().unwrap(); + let Some(max_bytes) = max_bytes else { + return Ok(std::mem::take(&mut *buf)); + }; + + if max_bytes >= buf.len() { + return Ok(std::mem::take(&mut *buf)); + } + + let mut split = max_bytes; + while split > 0 && !buf.is_char_boundary(split) { + split -= 1; + } + + if split == 0 { + return Ok(String::new()); + } + + Ok(buf.drain(..split).collect()) +} + +fn py_runtime_err(err: impl std::fmt::Display) -> PyErr { + PyRuntimeError::new_err(err.to_string()) +} + +#[pymodule] +fn sk8brd_cdba(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) +} From 2c39d69821fae7c3fd0dd75988d0c46411fc8986 Mon Sep 17 00:00:00 2001 From: Bjorn Andersson Date: Fri, 3 Jul 2026 01:51:14 +0000 Subject: [PATCH 5/5] python: Add python example Add a small Python example application that exercises the module through environment variables and documents the packaging workflow. Signed-off-by: Bjorn Andersson --- README.md | 4 ++ examples/python_cdba/README.md | 28 ++++++++ examples/python_cdba/boot_linux.py | 108 +++++++++++++++++++++++++++++ 3 files changed, 140 insertions(+) create mode 100644 examples/python_cdba/README.md create mode 100755 examples/python_cdba/boot_linux.py diff --git a/README.md b/README.md index 5ae6b2e..c8e7382 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,10 @@ The `sk8brd-cdba` package exposes the sk8brd cdba client as `sk8brd_cdba`: ```sh python -m pip install maturin maturin develop +python examples/python_cdba/boot_linux.py \ + --host \ + --board \ + --image ``` ## License diff --git a/examples/python_cdba/README.md b/examples/python_cdba/README.md new file mode 100644 index 0000000..9c40097 --- /dev/null +++ b/examples/python_cdba/README.md @@ -0,0 +1,28 @@ +# Python cdba example + +This is a small Python application that uses the `sk8brd_cdba` module to list +available devices, boot a Linux boot image on a selected board, wait for a +console prompt, run a command, collect its output, and disconnect. + +```sh +python -m pip install maturin +maturin develop +python examples/python_cdba/boot_linux.py \ + --host \ + --board \ + --image \ + --prompt 'root@qcom-armv8a:~#' \ + --command 'uname -a; id' +``` + +Useful options: + +- `--port`, defaults to `22` +- `--user`, defaults to `cdba` +- `--timeout`, defaults to `120` +- `--prompt`, defaults to `root@qcom-armv8a:~#` +- `--command`, defaults to `uname -a; id` + +The example intentionally drives the console through `read_console()` and +`write_console()` so it can be used as a starting point for more involved +Python tests. diff --git a/examples/python_cdba/boot_linux.py b/examples/python_cdba/boot_linux.py new file mode 100755 index 0000000..fa23d6d --- /dev/null +++ b/examples/python_cdba/boot_linux.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +import argparse +import sys +import time +import uuid + + +def marker_exit_code(output, marker): + for line in output.splitlines(): + if marker not in line: + continue + _, _, suffix = line.rpartition(f"{marker}:") + try: + return int(suffix) + except ValueError: + continue + return None + + +def read_until(session, done, description, timeout_secs): + deadline = time.monotonic() + timeout_secs + output = "" + + while time.monotonic() < deadline: + chunk = session.read_console() + if chunk: + print(chunk, end="", flush=True) + output += chunk + result = done(output) + if result is not None: + return output, result + time.sleep(0.1) + + raise TimeoutError(f"timed out waiting for {description}") + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Boot an image through CDBA and run one console command.", + ) + parser.add_argument( + "--host", + required=True, + help="CDBA server hostname or IP address", + ) + parser.add_argument("--board", required=True, help="CDBA board name") + parser.add_argument("--image", required=True, help="boot image to upload") + parser.add_argument("--port", type=int, default=22, help="CDBA SSH port") + parser.add_argument("--user", default="cdba", help="CDBA SSH user") + parser.add_argument( + "--timeout", + type=int, + default=120, + help="timeout in seconds for boot prompt and command completion", + ) + parser.add_argument( + "--prompt", + default="root@qcom-armv8a:~#", + help="shell prompt to wait for after boot", + ) + parser.add_argument( + "--command", + default="uname -a; id", + help="shell command to run after boot", + ) + return parser.parse_args() + + +def main(): + args = parse_args() + import sk8brd_cdba + + timeout_secs = args.timeout + prompt = args.prompt + command = args.command + client = sk8brd_cdba.CdbaClient( + args.host, + port=args.port, + user=args.user, + timeout_secs=timeout_secs, + ) + + print(client.list_devices(), end="") + session = client.boot_image_session(args.board, args.image) + try: + read_until(session, lambda output: True if prompt in output else None, prompt, timeout_secs) + + marker = f"__BOOT_LINUX_DONE_{uuid.uuid4().hex}__" + session.write_console(f"\n{command}\nprintf '\\n{marker}:%s\\n' \"$?\"\n") + _, exit_code = read_until( + session, + lambda output: marker_exit_code(output, marker), + f"{marker}:", + timeout_secs, + ) + finally: + session.close() + + print(f"\ncommand exit code: {exit_code}") + return exit_code + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except Exception as err: + print(err, file=sys.stderr) + raise SystemExit(1)