For msm/python

Cargo.toml

@@ -1,3 +1,3 @@
 [workspace]
-members = ["client", "cli", "proto"]
+members = ["client", "cli", "proto", "python"]
 resolver = "2"

README.md

@@ -23,11 +23,19 @@ 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`
 
 ## Credits
 ```
 Author: Konrad Dybcio <konradybcio@kernel.org>
 cdba contributors (for the original cdba implementation)
-```
+```

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<u8> = 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::<Msg>(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::<Msg>(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::<Vec<_>>().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::<Vec<_>>().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?;