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..c8e7382 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,18 @@ 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 +python examples/python_cdba/boot_linux.py \ + --host \ + --board \ + --image +``` + ## License `BSD-3-Clause` @@ -30,4 +42,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/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/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) diff --git a/proto/src/lib.rs b/proto/src/lib.rs index 0e448da..cf62127 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::(); +pub 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]; +pub 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( diff --git a/proto/src/ssh.rs b/proto/src/ssh.rs index 94f8d3b..fab0fc1 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,51 +31,114 @@ 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"); } let chan = sess .channel_open_session() .await - .expect("Couldn't open session"); + .context("could not open SSH session")?; Ok(chan) } +fn candidate_private_keys() -> 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 { @@ -106,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, _ => (), } @@ -131,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, } } } diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..ddb2e6f --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[build-system] +requires = ["maturin>=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(()) +}