diff --git a/Cargo.lock b/Cargo.lock index a5399b9ddc..5e93f72a58 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3421,6 +3421,7 @@ dependencies = [ "roxmltree", "rstest", "rusqlite", + "scopeguard", "serde", "serde_json", "serde_urlencoded", diff --git a/Cargo.toml b/Cargo.toml index c7bafe89d6..102bd1246b 100644 --- a/Cargo.toml +++ b/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" diff --git a/crates/nu-command/Cargo.toml b/crates/nu-command/Cargo.toml index ce1c336153..ef88aa92e4 100644 --- a/crates/nu-command/Cargo.toml +++ b/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 } \ No newline at end of file +rand_chacha = { workspace = true } diff --git a/crates/nu-command/src/platform/term/term_query.rs b/crates/nu-command/src/platform/term/term_query.rs index a4966edde0..2cb11f1dc3 100644 --- a/crates/nu-command/src/platform/term/term_query.rs +++ b/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 { 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 { let query: Vec = call.req(engine_state, stack, 0)?; + let keep = call.has_flag(engine_state, stack, "keep")?; + let prefix: Option> = call.get_flag(engine_state, stack, "prefix")?; + let prefix = prefix.unwrap_or_default(); let terminator: Option> = 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()) } }