diff --git a/crates/nu-command/src/bytes/starts_with.rs b/crates/nu-command/src/bytes/starts_with.rs index 16504100d3..0de2ce9105 100644 --- a/crates/nu-command/src/bytes/starts_with.rs +++ b/crates/nu-command/src/bytes/starts_with.rs @@ -4,6 +4,7 @@ use nu_protocol::ast::Call; use nu_protocol::ast::CellPath; use nu_protocol::engine::{Command, EngineState, Stack}; use nu_protocol::Category; +use nu_protocol::IntoPipelineData; use nu_protocol::{Example, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, Value}; struct Arguments { @@ -60,13 +61,63 @@ impl Command for BytesStartsWith { pattern, cell_paths, }; - operate( - starts_with, - arg, - input, - call.head, - engine_state.ctrlc.clone(), - ) + + match input { + PipelineData::ExternalStream { + stdout: Some(stream), + span, + .. + } => { + let mut i = 0; + + for item in stream { + let byte_slice = match &item { + // String and binary data are valid byte patterns + Ok(Value::String { val, .. }) => val.as_bytes(), + Ok(Value::Binary { val, .. }) => val, + // If any Error value is output, echo it back + Ok(v @ Value::Error { .. }) => return Ok(v.clone().into_pipeline_data()), + // Unsupported data + Ok(other) => { + return Ok(Value::Error { + error: ShellError::OnlySupportsThisInputType( + "string and binary".into(), + other.get_type().to_string(), + span, + // This line requires the Value::Error match above. + other.expect_span(), + ), + } + .into_pipeline_data()); + } + Err(err) => return Err(err.to_owned()), + }; + + let max = byte_slice.len().min(arg.pattern.len() - i); + + if byte_slice[..max] == arg.pattern[i..i + max] { + i += max; + + if i >= arg.pattern.len() { + return Ok(Value::boolean(true, span).into_pipeline_data()); + } + } else { + return Ok(Value::boolean(false, span).into_pipeline_data()); + } + } + + // We reached the end of the stream and never returned, + // the pattern wasn't exhausted so it probably doesn't match + Ok(Value::boolean(false, span).into_pipeline_data()) + } + _ => operate( + starts_with, + arg, + input, + call.head, + engine_state.ctrlc.clone(), + ), + } } fn examples(&self) -> Vec { diff --git a/crates/nu-command/tests/commands/bytes/mod.rs b/crates/nu-command/tests/commands/bytes/mod.rs new file mode 100644 index 0000000000..b5517bdacb --- /dev/null +++ b/crates/nu-command/tests/commands/bytes/mod.rs @@ -0,0 +1 @@ +mod starts_with; diff --git a/crates/nu-command/tests/commands/bytes/starts_with.rs b/crates/nu-command/tests/commands/bytes/starts_with.rs new file mode 100644 index 0000000000..7e80ce15ef --- /dev/null +++ b/crates/nu-command/tests/commands/bytes/starts_with.rs @@ -0,0 +1,153 @@ +use nu_test_support::nu; + +#[test] +fn basic_binary_starts_with() { + let actual = nu!( + cwd: ".", + r#" + "hello world" | into binary | bytes starts-with 0x[68 65 6c 6c 6f] + "# + ); + + assert_eq!(actual.out, "true"); +} + +#[test] +fn basic_string_fails() { + let actual = nu!( + cwd: ".", + r#" + "hello world" | bytes starts-with 0x[68 65 6c 6c 6f] + "# + ); + + assert!(actual.err.contains("Input type not supported")); + assert_eq!(actual.out, ""); +} + +#[test] +fn short_stream_binary() { + let actual = nu!( + cwd: ".", + r#" + nu --testbin repeater (0x[01]) 5 | bytes starts-with 0x[010101] + "# + ); + + assert_eq!(actual.out, "true"); +} + +#[test] +fn short_stream_mismatch() { + let actual = nu!( + cwd: ".", + r#" + nu --testbin repeater (0x[010203]) 5 | bytes starts-with 0x[010204] + "# + ); + + assert_eq!(actual.out, "false"); +} + +#[test] +fn short_stream_binary_overflow() { + let actual = nu!( + cwd: ".", + r#" + nu --testbin repeater (0x[01]) 5 | bytes starts-with 0x[010101010101] + "# + ); + + assert_eq!(actual.out, "false"); +} + +#[test] +fn long_stream_binary() { + let actual = nu!( + cwd: ".", + r#" + nu --testbin repeater (0x[01]) 32768 | bytes starts-with 0x[010101] + "# + ); + + assert_eq!(actual.out, "true"); +} + +#[test] +fn long_stream_binary_overflow() { + // .. ranges are inclusive..inclusive, so we don't need to +1 to check for an overflow + let actual = nu!( + cwd: ".", + r#" + nu --testbin repeater (0x[01]) 32768 | bytes starts-with (0..32768 | each { 0x[01] } | bytes collect) + "# + ); + + assert_eq!(actual.out, "false"); +} + +#[test] +fn long_stream_binary_exact() { + // ranges are inclusive..inclusive, so we don't need to +1 to check for an overflow + let actual = nu!( + cwd: ".", + r#" + nu --testbin repeater (0x[01020304]) 8192 | bytes starts-with (0..<8192 | each { 0x[01020304] } | bytes collect) + "# + ); + + assert_eq!(actual.out, "true"); +} + +#[test] +fn long_stream_string_exact() { + // ranges are inclusive..inclusive, so we don't need to +1 to check for an overflow + let actual = nu!( + cwd: ".", + r#" + nu --testbin repeater hell 8192 | bytes starts-with (0..<8192 | each { "hell" | into binary } | bytes collect) + "# + ); + + assert_eq!(actual.out, "true"); +} + +#[test] +fn long_stream_mixed_exact() { + // ranges are inclusive..inclusive, so we don't need to +1 to check for an overflow + let actual = nu!( + cwd: ".", + r#" + let binseg = (0..<2048 | each { 0x[003d9fbf] } | bytes collect) + let strseg = (0..<2048 | each { "hell" | into binary } | bytes collect) + + nu --testbin repeat_bytes 003d9fbf 2048 68656c6c 2048 | bytes starts-with (bytes build $binseg $strseg) + "# + ); + + assert_eq!( + actual.err, "", + "invocation failed. command line limit likely reached" + ); + assert_eq!(actual.out, "true"); +} + +#[test] +fn long_stream_mixed_overflow() { + // ranges are inclusive..inclusive, so we don't need to +1 to check for an overflow + let actual = nu!( + cwd: ".", + r#" + let binseg = (0..<2048 | each { 0x[003d9fbf] } | bytes collect) + let strseg = (0..<2048 | each { "hell" | into binary } | bytes collect) + + nu --testbin repeat_bytes 003d9fbf 2048 68656c6c 2048 | bytes starts-with (bytes build $binseg $strseg 0x[01]) + "# + ); + + assert_eq!( + actual.err, "", + "invocation failed. command line limit likely reached" + ); + assert_eq!(actual.out, "false"); +} diff --git a/crates/nu-command/tests/commands/mod.rs b/crates/nu-command/tests/commands/mod.rs index 3e95115df7..872d2e9f3e 100644 --- a/crates/nu-command/tests/commands/mod.rs +++ b/crates/nu-command/tests/commands/mod.rs @@ -4,6 +4,7 @@ mod any; mod append; mod assignment; mod break_; +mod bytes; mod cal; mod cd; mod compact; diff --git a/src/main.rs b/src/main.rs index e0f0fcef33..813a71791c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -168,6 +168,7 @@ fn main() -> Result<()> { "nonu" => test_bins::nonu(), "chop" => test_bins::chop(), "repeater" => test_bins::repeater(), + "repeat_bytes" => test_bins::repeat_bytes(), "nu_repl" => test_bins::nu_repl(), "input_bytes_length" => test_bins::input_bytes_length(), _ => std::process::exit(1), diff --git a/src/test_bins.rs b/src/test_bins.rs index 6bbbc29918..0a0ded7db9 100644 --- a/src/test_bins.rs +++ b/src/test_bins.rs @@ -84,6 +84,32 @@ pub fn repeater() { let _ = stdout.flush(); } +// A version of repeater that can output binary data, even null bytes +pub fn repeat_bytes() { + let mut stdout = io::stdout(); + let args = args(); + let mut args = args.iter().skip(1); + + while let (Some(binary), Some(count)) = (args.next(), args.next()) { + let bytes: Vec = (0..binary.len()) + .step_by(2) + .map(|i| { + u8::from_str_radix(&binary[i..i + 2], 16) + .expect("binary string is valid hexadecimal") + }) + .collect(); + let count: u64 = count.parse().expect("repeat count must be a number"); + + for _ in 0..count { + stdout + .write_all(&bytes) + .expect("writing to stdout must not fail"); + } + } + + let _ = stdout.flush(); +} + pub fn iecho() { // println! panics if stdout gets closed, whereas writeln gives us an error let mut stdout = io::stdout();