console input, QoL

src/lib.rs

@@ -1,6 +1,8 @@
 use serde::{Deserialize, Serialize};
 use std::io::{stdout, Write};
 use std::mem::size_of;
+use std::sync::Arc;
+use tokio::sync::Mutex;
 
 #[repr(u8)]
 #[derive(Debug, PartialEq)]
@@ -59,7 +61,14 @@ pub struct Sk8brdMsg {
 }
 pub const MSG_HDR_SIZE: usize = size_of::<Sk8brdMsg>();
 
-pub fn send_msg(write_sink: &mut impl Write, r#type: Sk8brdMsgs, buf: &[u8]) -> anyhow::Result<()> {
+pub async fn send_msg(
+    write_sink: &mut Arc<Mutex<impl Write>>,
+    r#type: Sk8brdMsgs,
+    buf: &[u8],
+) -> anyhow::Result<()> {
+    // Make sure we're not trying to send two messages at once
+    let mut write_sink = write_sink.lock().await;
+
     let len = buf.len();
     let hdr = [r#type as u8, (len & 0xff) as u8, ((len >> 8) & 0xff) as u8];
 
@@ -68,8 +77,11 @@ pub fn send_msg(write_sink: &mut impl Write, r#type: Sk8brdMsgs, buf: &[u8]) ->
     Ok(())
 }
 
-pub fn send_ack(write_sink: &mut impl Write, r#type: Sk8brdMsgs) -> anyhow::Result<()> {
-    send_msg(write_sink, r#type, &[])
+pub async fn send_ack(
+    write_sink: &mut Arc<Mutex<impl Write>>,
+    r#type: Sk8brdMsgs,
+) -> anyhow::Result<()> {
+    send_msg(write_sink, r#type, &[]).await
 }
 
 pub fn parse_recv_msg(buf: &[u8]) -> Sk8brdMsg {
@@ -83,13 +95,13 @@ pub fn parse_recv_msg(buf: &[u8]) -> Sk8brdMsg {
     msg
 }
 
-pub fn console_print(buf: &[u8]) {
+pub async fn console_print(buf: &[u8]) {
     print!("{}", String::from_utf8_lossy(buf));
     stdout().flush().unwrap();
 }
 
 #[allow(clippy::explicit_write)]
-pub fn send_image(write_sink: &mut impl Write, buf: &[u8]) -> anyhow::Result<()> {
+pub async fn send_image(write_sink: &mut Arc<Mutex<impl Write>>, buf: &[u8]) -> anyhow::Result<()> {
     let mut last_percent_done: usize = 0;
     let mut bytes_sent = 0;
 
@@ -100,20 +112,23 @@ pub fn send_image(write_sink: &mut impl Write, buf: &[u8]) -> anyhow::Result<()>
             write!(stdout(), " Sending image: {}%\r", percent_done).unwrap();
         }
 
-        send_msg(write_sink, Sk8brdMsgs::MsgFastbootDownload, chunk)?;
+        send_msg(write_sink, Sk8brdMsgs::MsgFastbootDownload, chunk).await?;
 
         bytes_sent += chunk.len();
         last_percent_done = percent_done;
     }
 
-    send_ack(write_sink, Sk8brdMsgs::MsgFastbootDownload)
+    send_ack(write_sink, Sk8brdMsgs::MsgFastbootDownload).await
 }
 
-pub fn select_brd(write_sink: &mut impl Write, name: &str) -> anyhow::Result<()> {
-    send_msg(write_sink, Sk8brdMsgs::MsgSelectBoard, name.as_bytes())
+pub async fn select_brd(write_sink: &mut Arc<Mutex<impl Write>>, name: &str) -> anyhow::Result<()> {
+    send_msg(write_sink, Sk8brdMsgs::MsgSelectBoard, name.as_bytes()).await
 }
 
-pub fn send_vbus_ctrl(write_sink: &mut impl Write, en: bool) -> anyhow::Result<()> {
+pub async fn send_vbus_ctrl(
+    write_sink: &mut Arc<Mutex<impl Write>>,
+    en: bool,
+) -> anyhow::Result<()> {
     send_ack(
         write_sink,
         if en {
@@ -122,6 +137,7 @@ pub fn send_vbus_ctrl(write_sink: &mut impl Write, en: bool) -> anyhow::Result<(
             Sk8brdMsgs::MsgVbusOff
         },
     )
+    .await
 }
 
 #[allow(clippy::explicit_write)]

src/main.rs

@@ -38,13 +38,27 @@ struct Args {
     power_cycle: bool,
 }
 
-async fn handle_keypress(c: char, quit: &mut Arc<Mutex<bool>>) {
-    if true {
+async fn handle_keypress(
+    c: char,
+    quit: &mut Arc<Mutex<bool>>,
+    special: &mut bool,
+    message_sink: &mut Arc<Mutex<impl Write>>,
+) {
+    if *special {
+        *special = false;
         match c {
             'q' => *quit.lock().await = true,
             'Q' => *quit.lock().await = true,
             _ => (),
         }
+    } else {
+        match c.try_into() {
+            Ok(1u8) => *special = true, // CTRL-A, TODO: configurable?
+            Ok(_) => send_msg(message_sink, Sk8brdMsgs::MsgConsole, &[c as u8])
+                .await
+                .unwrap(),
+            Err(_) => (),
+        }
     }
 }
 
@@ -55,6 +69,12 @@ macro_rules! todo {
     }};
 }
 
+macro_rules! get_arc {
+    ($a: expr) => {{
+        $a.lock().await
+    }};
+}
+
 // For raw mode TTY
 #[allow(clippy::explicit_write)]
 #[tokio::main]
@@ -80,62 +100,67 @@ async fn main() -> anyhow::Result<()> {
     sess.userauth_agent(USERNAME)
         .with_context(|| format!("Couldn't authenticate as {USERNAME}"))?;
 
-    let mut chan = sess.channel_session()?;
-    chan.exec(CDBA_SERVER_BIN_NAME)
+    let mut chan = Arc::new(Mutex::new(sess.channel_session()?));
+    (*get_arc!(chan))
+        .exec(CDBA_SERVER_BIN_NAME)
         .with_context(|| format!("Couldn't execute {CDBA_SERVER_BIN_NAME} on remote host"))?;
 
     sess.set_blocking(false);
 
-    send_ack(&mut chan, Sk8brdMsgs::MsgListDevices)?;
-    select_brd(&mut chan, &args.board)?;
+    send_ack(&mut chan, Sk8brdMsgs::MsgListDevices).await?;
+    select_brd(&mut chan, &args.board).await?;
     if args.power_cycle {
         println!("Powering off the board first");
-        send_ack(&mut chan, Sk8brdMsgs::MsgPowerOff)?;
+        send_ack(&mut chan, Sk8brdMsgs::MsgPowerOff).await?;
     }
 
     crossterm::terminal::enable_raw_mode()?;
 
     let mut quit2 = Arc::clone(&quit);
+    let mut chan2 = Arc::clone(&chan);
     let stdin_handler = tokio::spawn(async move {
         let mut stdin = os_pipe::dup_stdin().expect("Couldn't dup stdin");
+        let mut ctrl_a_pressed = false;
 
         while !*quit2.lock().await {
             if let Ok(len) = stdin.read(&mut key_buf) {
                 for c in key_buf[0..len].iter() {
-                    handle_keypress(*c as char, &mut quit2).await;
+                    handle_keypress(*c as char, &mut quit2, &mut ctrl_a_pressed, &mut chan2).await;
                 }
             };
         }
     });
 
     while !*quit.lock().await {
         // Stream of "blue text" - status updates from the server
-        if let Ok(bytes_read) = chan.stderr().read(&mut buf) {
+        if let Ok(bytes_read) = (*get_arc!(chan)).stderr().read(&mut buf) {
             let s = String::from_utf8_lossy(&buf[..bytes_read]);
             writeln!(stdout(), "{}\r", s.blue())?;
             stdout().flush()?;
         }
 
         // Msg handler
         // Read the message header first
-        if chan.read_exact(&mut hdr_buf).is_ok() {
+        if (*get_arc!(chan)).read_exact(&mut hdr_buf).is_ok() {
             sess.set_blocking(true);
             let msg = parse_recv_msg(&hdr_buf);
             let mut msgbuf = vec![0u8; msg.len as usize];
 
             // Now read the actual data...
-            chan.read_exact(&mut msgbuf)?;
+            (*get_arc!(chan)).read_exact(&mut msgbuf)?;
 
             // ..and process it
             match msg.r#type.try_into() {
-                Ok(Sk8brdMsgs::MsgSelectBoard) => send_msg(&mut chan, Sk8brdMsgs::MsgPowerOn, &[])?,
-                Ok(Sk8brdMsgs::MsgConsole) => console_print(&msgbuf),
+                Ok(Sk8brdMsgs::MsgSelectBoard) => {
+                    send_msg(&mut chan, 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 chan, &fastboot_image)?
+                        send_image(&mut chan, &fastboot_image).await?
                     }
                 }
                 Ok(Sk8brdMsgs::MsgFastbootDownload) => (),
@@ -163,7 +188,7 @@ async fn main() -> anyhow::Result<()> {
     crossterm::terminal::disable_raw_mode()?;
 
     // Power off the board on goodbye
-    send_ack(&mut chan, Sk8brdMsgs::MsgPowerOff)?;
+    send_ack(&mut chan, Sk8brdMsgs::MsgPowerOff).await?;
 
     sess.disconnect(
         Option::Some(ssh2::DisconnectCode::ConnectionLost),