term query: refactor, add --beginning flag (#14446)

# Description

- Refactor code to be simpler.
- Make the mentioned changes.
- `scopeguard` is added as a direct dependency. Helps simplify the code.
Rather than roll an ad-hoc version of it myself, I thought it would be
better to use `scopeguard` as it was already an indirect dependency.

# User-Facing Changes

- Add `--beginning` flag, which is used to validate the response and
provide early errors in case of unexpected inputs.
- Both `terminator` and `beginning` sequences (when provided) are not
included in the command's output. Turns out they are almost always
removed from the output, and because they are known beforehand they can
be added back by the user.

Cargo.lock

@@ -3421,6 +3421,7 @@ dependencies = [
  "roxmltree",
  "rstest",
  "rusqlite",
+ "scopeguard",
  "serde",
  "serde_json",
  "serde_urlencoded",

Cargo.toml

@@ -147,6 +147,7 @@ roxmltree = "0.19"
 rstest = { version = "0.23", default-features = false }
 rusqlite = "0.31"
 rust-embed = "8.5.0"
+scopeguard = { version = "1.2.0" }
 serde = { version = "1.0" }
 serde_json = "1.0"
 serde_urlencoded = "0.7.1"

crates/nu-command/Cargo.toml

@@ -80,6 +80,7 @@ regex = { workspace = true }
 roxmltree = { workspace = true }
 rusqlite = { workspace = true, features = ["bundled", "backup", "chrono"], optional = true }
 rmp = { workspace = true }
+scopeguard = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true, features = ["preserve_order"] }
 serde_urlencoded = { workspace = true }
@@ -196,4 +197,4 @@ quickcheck_macros = { workspace = true }
 rstest = { workspace = true, default-features = false }
 pretty_assertions = { workspace = true }
 tempfile = { workspace = true }
-rand_chacha = { workspace = true }
+rand_chacha = { workspace = true }

crates/nu-command/src/platform/term/term_query.rs

@@ -23,9 +23,12 @@ impl Command for TermQuery {
         "Print the given query, and read the immediate result from stdin.
 
 The standard input will be read right after `query` is printed, and consumed until the `terminator`
-sequence is encountered. The `terminator` is not removed from the output.
+sequence is encountered. The `terminator` is not included in the output.
 
-If `terminator` is not supplied, input will be read until Ctrl-C is pressed."
+If `terminator` is not supplied, input will be read until Ctrl-C is pressed.
+
+If `prefix` is supplied, input's initial bytes will be validated against it.
+The `prefix` is not included in the output."
     }
 
     fn signature(&self) -> Signature {
@@ -38,29 +41,45 @@ If `terminator` is not supplied, input will be read until Ctrl-C is pressed."
                 SyntaxShape::OneOf(vec![SyntaxShape::Binary, SyntaxShape::String]),
                 "The query that will be printed to stdout.",
             )
+            .named(
+                "prefix",
+                SyntaxShape::OneOf(vec![SyntaxShape::Binary, SyntaxShape::String]),
+                "Prefix sequence for the expected reply.",
+                Some('p'),
+            )
             .named(
                 "terminator",
                 SyntaxShape::OneOf(vec![SyntaxShape::Binary, SyntaxShape::String]),
                 "Terminator sequence for the expected reply.",
                 Some('t'),
             )
+            .switch(
+                "keep",
+                "Include prefix and terminator in the output.",
+                Some('k'),
+            )
     }
 
     fn examples(&self) -> Vec<Example> {
         vec![
             Example {
                 description: "Get cursor position.",
-                example: r#"term query (ansi cursor_position) --terminator 'R'"#,
+                example: r#"term query (ansi cursor_position) --prefix (ansi csi) --terminator 'R'"#,
                 result: None,
             },
             Example {
                 description: "Get terminal background color.",
-                example: r#"term query $'(ansi osc)10;?(ansi st)' --terminator (ansi st)"#,
+                example: r#"term query $'(ansi osc)10;?(ansi st)' --prefix $'(ansi osc)10;' --terminator (ansi st)"#,
+                result: None,
+            },
+            Example {
+                description: "Get terminal background color. (some terminals prefer `char bel` rather than `ansi st` as string terminator)",
+                example: r#"term query $'(ansi osc)10;?(char bel)' --prefix $'(ansi osc)10;' --terminator (char bel)"#,
                 result: None,
             },
             Example {
                 description: "Read clipboard content on terminals supporting OSC-52.",
-                example: r#"term query $'(ansi osc)52;c;?(ansi st)' --terminator (ansi st)"#,
+                example: r#"term query $'(ansi osc)52;c;?(ansi st)' --prefix $'(ansi osc)52;c;' --terminator (ansi st)"#,
                 result: None,
             },
         ]
@@ -74,9 +93,15 @@ If `terminator` is not supplied, input will be read until Ctrl-C is pressed."
         _input: PipelineData,
     ) -> Result<PipelineData, ShellError> {
         let query: Vec<u8> = call.req(engine_state, stack, 0)?;
+        let keep = call.has_flag(engine_state, stack, "keep")?;
+        let prefix: Option<Vec<u8>> = call.get_flag(engine_state, stack, "prefix")?;
+        let prefix = prefix.unwrap_or_default();
         let terminator: Option<Vec<u8>> = call.get_flag(engine_state, stack, "terminator")?;
 
         crossterm::terminal::enable_raw_mode()?;
+        scopeguard::defer! {
+            let _ = crossterm::terminal::disable_raw_mode();
+        }
 
         // clear terminal events
         while crossterm::event::poll(Duration::from_secs(0))? {
@@ -94,44 +119,55 @@ If `terminator` is not supplied, input will be read until Ctrl-C is pressed."
             stdout.flush()?;
         }
 
-        let out = if let Some(terminator) = terminator {
+        // Validate and skip prefix
+        for bc in prefix {
+            stdin.read_exact(&mut b)?;
+            if b[0] != bc {
+                return Err(ShellError::GenericError {
+                    error: "Input did not begin with expected sequence".into(),
+                    msg: "".into(),
+                    span: None,
+                    help: Some("Try running without `--prefix` and inspecting the output.".into()),
+                    inner: vec![],
+                });
+            }
+            if keep {
+                buf.push(b[0]);
+            }
+        }
+
+        if let Some(terminator) = terminator {
             loop {
-                if let Err(err) = stdin.read_exact(&mut b) {
-                    break Err(ShellError::from(err));
-                }
+                stdin.read_exact(&mut b)?;
 
                 if b[0] == CTRL_C {
-                    break Err(ShellError::Interrupted { span: call.head });
+                    return Err(ShellError::InterruptedByUser {
+                        span: Some(call.head),
+                    });
                 }
 
                 buf.push(b[0]);
 
                 if buf.ends_with(&terminator) {
-                    break Ok(Value::Binary {
-                        val: buf,
-                        internal_span: call.head,
+                    if !keep {
+                        // Remove terminator
+                        buf.drain((buf.len() - terminator.len())..);
                     }
-                    .into_pipeline_data());
+                    break;
                 }
             }
         } else {
             loop {
-                if let Err(err) = stdin.read_exact(&mut b) {
-                    break Err(ShellError::from(err));
-                }
+                stdin.read_exact(&mut b)?;
 
                 if b[0] == CTRL_C {
-                    break Ok(Value::Binary {
-                        val: buf,
-                        internal_span: call.head,
-                    }
-                    .into_pipeline_data());
+                    break;
                 }
 
                 buf.push(b[0]);
             }
         };
-        crossterm::terminal::disable_raw_mode()?;
-        out
+
+        Ok(Value::binary(buf, call.head).into_pipeline_data())
     }
 }