Make the mechanism for escaping discoverable and customizable

Author: lif

bin/propolis-cli/src/main.rs

@@ -1,9 +1,10 @@
+use std::ffi::OsString;
 use std::fs::File;
 use std::io::BufReader;
 use std::path::{Path, PathBuf};
 use std::{
     net::{IpAddr, SocketAddr, ToSocketAddrs},
-    os::unix::prelude::AsRawFd,
+    os::unix::prelude::{AsRawFd, OsStringExt},
     time::Duration,
 };
 
@@ -90,6 +91,15 @@ enum Command {
         /// Defaults to the most recent 16 KiB of console output (-16384).
         #[clap(long, short)]
         byte_offset: Option<i64>,
+
+        /// If this sequence of bytes is typed, the client will exit.
+        /// Defaults to ^]^C (Ctrl+], Ctrl+C).
+        #[clap(long, short, default_value = "\x1d\x03")]
+        escape_string: OsString,
+
+        /// Disable escape string altogether (to exit, use pkill or similar).
+        #[clap(long, short = 'E')]
+        no_escape: bool,
     },
 
     /// Migrate instance to new propolis-server
@@ -221,60 +231,56 @@ async fn put_instance(
 async fn stdin_to_websockets_task(
     mut stdinrx: tokio::sync::mpsc::Receiver<Vec<u8>>,
     wstx: tokio::sync::mpsc::Sender<Vec<u8>>,
+    escape_vector: Option<Vec<u8>>,
 ) {
-    // next_raw must live outside loop, because Ctrl-A should work across
-    // multiple inbuf reads.
-    let mut next_raw = false;
+    if let Some(esc_sequence) = &escape_vector {
+        // esc_pos must live outside loop, because escape string should work
+        // across multiple inbuf reads.
+        let mut esc_pos = 0;
 
-    loop {
-        let inbuf = if let Some(inbuf) = stdinrx.recv().await {
-            inbuf
-        } else {
-            continue;
-        };
+        loop {
+            let inbuf = if let Some(inbuf) = stdinrx.recv().await {
+                inbuf
+            } else {
+                continue;
+            };
 
-        // Put bytes from inbuf to outbuf, but don't send Ctrl-A unless
-        // next_raw is true.
-        let mut outbuf = Vec::with_capacity(inbuf.len());
-
-        let mut exit = false;
-        for c in inbuf {
-            match c {
-                // Ctrl-A means send next one raw
-                b'\x01' => {
-                    if next_raw {
-                        // Ctrl-A Ctrl-A should be sent as Ctrl-A
-                        outbuf.push(c);
-                        next_raw = false;
-                    } else {
-                        next_raw = true;
-                    }
-                }
-                b'\x03' => {
-                    if !next_raw {
-                        // Exit on non-raw Ctrl-C
+            // Put bytes from inbuf to outbuf, but don't send characters in the
+            // escape string sequence unless we bail.
+            let mut outbuf = Vec::with_capacity(inbuf.len());
+
+            let mut exit = false;
+            for c in inbuf {
+                if c == esc_sequence[esc_pos] {
+                    esc_pos += 1;
+                    if esc_pos == esc_sequence.len() {
+                        // Exit on completed escape string
                         exit = true;
                         break;
-                    } else {
-                        // Otherwise send Ctrl-C
-                        outbuf.push(c);
-                        next_raw = false;
                     }
-                }
-                _ => {
+                } else {
+                    // they bailed from the sequence,
+                    // feed everything that matched so far through
+                    if esc_pos != 0 {
+                        outbuf.extend(&esc_sequence[..esc_pos])
+                    }
+                    esc_pos = 0;
                     outbuf.push(c);
-                    next_raw = false;
                 }
             }
-        }
 
-        // Send what we have, even if there's a Ctrl-C at the end.
-        if !outbuf.is_empty() {
-            wstx.send(outbuf).await.unwrap();
-        }
+            // Send what we have, even if we're about to exit.
+            if !outbuf.is_empty() {
+                wstx.send(outbuf).await.unwrap();
+            }
 
-        if exit {
-            break;
+            if exit {
+                break;
+            }
+        }
+    } else {
+        while let Some(buf) = stdinrx.recv().await {
+            wstx.send(buf).await.unwrap();
         }
     }
 }
@@ -286,7 +292,10 @@ async fn test_stdin_to_websockets_task() {
     let (stdintx, stdinrx) = tokio::sync::mpsc::channel(16);
     let (wstx, mut wsrx) = tokio::sync::mpsc::channel(16);
 
-    tokio::spawn(async move { stdin_to_websockets_task(stdinrx, wstx).await });
+    let escape_vector = Some(vec![0x1d, 0x03]);
+    tokio::spawn(async move {
+        stdin_to_websockets_task(stdinrx, wstx, escape_vector).await
+    });
 
     // send characters, receive characters
     stdintx
@@ -296,33 +305,22 @@ async fn test_stdin_to_websockets_task() {
     let actual = wsrx.recv().await.unwrap();
     assert_eq!(String::from_utf8(actual).unwrap(), "test post please ignore");
 
-    // don't send ctrl-a
-    stdintx.send("\x01".chars().map(|c| c as u8).collect()).await.unwrap();
+    // don't send a started escape sequence
+    stdintx.send("\x1d".chars().map(|c| c as u8).collect()).await.unwrap();
     assert_eq!(wsrx.try_recv(), Err(TryRecvError::Empty));
 
-    // the "t" here is sent "raw" because of last ctrl-a but that doesn't change anything
+    // since we didn't enter the \x03, the previous \x1d shows up here
     stdintx.send("test".chars().map(|c| c as u8).collect()).await.unwrap();
     let actual = wsrx.recv().await.unwrap();
-    assert_eq!(String::from_utf8(actual).unwrap(), "test");
-
-    // ctrl-a ctrl-c = only ctrl-c sent
-    stdintx.send("\x01\x03".chars().map(|c| c as u8).collect()).await.unwrap();
-    let actual = wsrx.recv().await.unwrap();
-    assert_eq!(String::from_utf8(actual).unwrap(), "\x03");
+    assert_eq!(String::from_utf8(actual).unwrap(), "\x1dtest");
 
-    // same as above, across two messages
-    stdintx.send("\x01".chars().map(|c| c as u8).collect()).await.unwrap();
+    // \x03 gets sent if not preceded by \x1d
     stdintx.send("\x03".chars().map(|c| c as u8).collect()).await.unwrap();
-    assert_eq!(wsrx.try_recv(), Err(TryRecvError::Empty));
     let actual = wsrx.recv().await.unwrap();
     assert_eq!(String::from_utf8(actual).unwrap(), "\x03");
 
-    // ctrl-a ctrl-a = only ctrl-a sent
-    stdintx.send("\x01\x01".chars().map(|c| c as u8).collect()).await.unwrap();
-    let actual = wsrx.recv().await.unwrap();
-    assert_eq!(String::from_utf8(actual).unwrap(), "\x01");
-
-    // ctrl-c on its own means exit
+    // \x1d followed by \x03 means exit, even if they're separate messages
+    stdintx.send("\x1d".chars().map(|c| c as u8).collect()).await.unwrap();
     stdintx.send("\x03".chars().map(|c| c as u8).collect()).await.unwrap();
     assert_eq!(wsrx.try_recv(), Err(TryRecvError::Empty));
 
@@ -333,6 +331,7 @@ async fn test_stdin_to_websockets_task() {
 async fn serial(
     addr: SocketAddr,
     byte_offset: Option<i64>,
+    escape_vector: Option<Vec<u8>>,
 ) -> anyhow::Result<()> {
     let client = propolis_client::Client::new(&format!("http://{}", addr));
     let mut req = client.instance_serial();
@@ -375,7 +374,9 @@ async fn serial(
         }
     });
 
-    tokio::spawn(async move { stdin_to_websockets_task(stdinrx, wstx).await });
+    tokio::spawn(async move {
+        stdin_to_websockets_task(stdinrx, wstx, escape_vector).await
+    });
 
     loop {
         tokio::select! {
@@ -569,7 +570,14 @@ async fn main() -> anyhow::Result<()> {
         }
         Command::Get => get_instance(&client).await?,
         Command::State { state } => put_instance(&client, state).await?,
-        Command::Serial { byte_offset } => serial(addr, byte_offset).await?,
+        Command::Serial { byte_offset, escape_string, no_escape } => {
+            let escape_vector = if no_escape || escape_string.is_empty() {
+                None
+            } else {
+                Some(escape_string.into_vec())
+            };
+            serial(addr, byte_offset, escape_vector).await?
+        }
         Command::Migrate { dst_server, dst_port, dst_uuid } => {
             let dst_addr = SocketAddr::new(dst_server, dst_port);
             let dst_client = Client::new(dst_addr, log.clone());