Make the mechanism for escaping discoverable and customizable

lif · lif · commit ec2a48729585 · 2023-01-19T21:04:26.000Z
This changes the method of entering an escape sequence:
- raw Ctrl+C gets sent to the VM unimpeded.
- by default, the sequence Ctrl+], Ctrl+C is used to quit the program
  (`^]^C`)
- this can be customized or removed via CLI flags, allowing the string
  be of arbitrary length.
  - i.e. if you `propolis-cli serial -e "beans"` and then type "bea",
    nothing gets sent to the VM after the "b" yet. and then if you type:
    1. "k", the VM gets sent "beak"
    2. '"ns", the VM doesn't get sent anything else, and the client
       exits.
- the client can be configured to pass through an arbitrary prefix
  length of the escape string before it starts suppressing inputs, such
  that you can, for example, mimic ssh's Enter-tilde-dot sequence
  without temporarily suppressing Enter presses not intended to
  start an escape sequence, which would interfere with function:
  `-e '^M~.' --escape-prefix-length=1` (this also works around ANSI
  escape sequences being sent by xterm-like emulators when Enter is
  pressed in a shell that sends a request for such)
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -96,6 +96,7 @@ propolis-server-config = { path = "crates/propolis-server-config" }
 propolis_types = { path = "crates/propolis-types" }
 quote = "1.0"
 rand = "0.8"
+regex = "1.7.1"
 reqwest = "0.11.12"
 rfb = { git = "https://github.com/oxidecomputer/rfb", rev = "0cac8d9c25eb27acfa35df80f3b9d371de98ab3b" }
 ring = "0.16"
diff --git a/bin/propolis-cli/Cargo.toml b/bin/propolis-cli/Cargo.toml
@@ -11,6 +11,7 @@ anyhow.workspace = true
 clap = { workspace = true, features = ["derive"] }
 futures.workspace = true
 libc.workspace = true
+regex.workspace = true
 propolis-client = { workspace = true, features = ["generated"] }
 slog.workspace = true
 slog-async.workspace = true
diff --git a/bin/propolis-cli/src/main.rs b/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,
 };
 
@@ -17,6 +18,7 @@ use propolis_client::handmade::{
     },
     Client,
 };
+use regex::bytes::Regex;
 use slog::{o, Drain, Level, Logger};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio_tungstenite::tungstenite::protocol::Role;
@@ -90,6 +92,29 @@ 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). Note that the string passed
+        /// for this argument is used verbatim without any parsing; in most
+        /// shells, if you wish to include a special character (such as Enter
+        /// or a Ctrl+letter combo), you can insert the character by preceding
+        /// it with Ctrl+V at the command line.
+        #[clap(long, short, default_value = "\x1d\x03")]
+        escape_string: OsString,
+
+        /// The number of bytes from the beginning of the escape string to pass
+        /// to the VM before beginning to buffer inputs until a mismatch.
+        /// Defaults to 0, such that input matching the escape string does not
+        /// get sent to the VM at all until a non-matching character is typed.
+        /// To mimic the escape sequence for exiting SSH (Enter, tilde, dot),
+        /// you may pass `-e '^M~.' --escape-prefix-length=1` such that normal
+        /// Enter presses are sent to the VM immediately.
+        #[clap(long, default_value = "0")]
+        escape_prefix_length: usize,
+
+        /// Disable escape string altogether (to exit, use pkill or similar).
+        #[clap(long, short = 'E')]
+        no_escape: bool,
     },
 
     /// Migrate instance to new propolis-server
@@ -225,60 +250,89 @@ 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>>,
+    escape_prefix_length: usize,
 ) {
-    // 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;
-        };
+        // matches partial increments of "\x1b[14;30R"
+        let ansi_curs_pat =
+            Regex::new("^\x1b(\\[([0-9]+(;([0-9]+R?)?)?)?)?$").unwrap();
+        let mut ansi_curs_check = Vec::new();
 
-        // 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;
+        loop {
+            let inbuf = if let Some(inbuf) = stdinrx.recv().await {
+                inbuf
+            } else {
+                continue;
+            };
+
+            // 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 {
+                // ignore ANSI escape sequence for the cursor position
+                // response sent by xterm-alikes in response to shells
+                // requesting one after receiving a newline.
+                if esc_pos > 0
+                    && esc_pos <= escape_prefix_length
+                    && b"\r\n".contains(&esc_sequence[esc_pos - 1])
+                {
+                    ansi_curs_check.push(c);
+                    if ansi_curs_pat.is_match(&ansi_curs_check) {
+                        // end of the sequence?
+                        if c == b'R' {
+                            outbuf.extend(&ansi_curs_check);
+                            ansi_curs_check.clear();
+                        }
+                        continue;
                     } else {
-                        next_raw = true;
+                        ansi_curs_check.pop(); // we're not `continue`ing
+                        outbuf.extend(&ansi_curs_check);
+                        ansi_curs_check.clear();
                     }
                 }
-                b'\x03' => {
-                    if !next_raw {
-                        // Exit on non-raw Ctrl-C
+
+                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
+                    } else if esc_pos <= escape_prefix_length {
+                        // let through incomplete prefix up to the given limit
                         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[escape_prefix_length..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();
         }
     }
 }
@@ -290,7 +344,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, 0).await
+    });
 
     // send characters, receive characters
     stdintx
@@ -300,33 +357,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));
 
@@ -337,6 +383,8 @@ async fn test_stdin_to_websockets_task() {
 async fn serial(
     addr: SocketAddr,
     byte_offset: Option<i64>,
+    escape_vector: Option<Vec<u8>>,
+    escape_prefix_length: usize,
 ) -> anyhow::Result<()> {
     let client = propolis_client::Client::new(&format!("http://{}", addr));
     let mut req = client.instance_serial();
@@ -379,7 +427,23 @@ async fn serial(
         }
     });
 
-    tokio::spawn(async move { stdin_to_websockets_task(stdinrx, wstx).await });
+    let escape_len = escape_vector.as_ref().map(|x| x.len()).unwrap_or(0);
+    if escape_prefix_length > escape_len {
+        anyhow::bail!(
+            "prefix length {} is greater than length of escape string ({})",
+            escape_prefix_length,
+            escape_len
+        );
+    }
+    tokio::spawn(async move {
+        stdin_to_websockets_task(
+            stdinrx,
+            wstx,
+            escape_vector,
+            escape_prefix_length,
+        )
+        .await
+    });
 
     loop {
         tokio::select! {
@@ -574,7 +638,20 @@ 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,
+            escape_prefix_length,
+            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, escape_prefix_length)
+                .await?
+        }
         Command::Migrate { dst_server, dst_port, dst_uuid, crucible_disks } => {
             let dst_addr = SocketAddr::new(dst_server, dst_port);
             let dst_client = Client::new(dst_addr, log.clone());
