Skip to content

Commit

Permalink
add b3sum --seek
Browse files Browse the repository at this point in the history
  • Loading branch information
oconnor663 committed Jun 7, 2023
1 parent 76f9339 commit 4108923
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 21 deletions.
38 changes: 25 additions & 13 deletions b3sum/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,16 @@ struct Inner {
/// When no file is given, or when - is given, read standard input.
file: Vec<PathBuf>,

/// Use the keyed mode, reading the 32-byte key from stdin
#[arg(long, requires("file"))]
keyed: bool,

/// Use the key derivation mode, with the given context string
///
/// Cannot be used with --keyed.
#[arg(long, value_name("CONTEXT"), conflicts_with(KEYED_ARG))]
derive_key: Option<String>,

/// The number of output bytes, before hex encoding
#[arg(
short,
Expand All @@ -35,23 +45,17 @@ struct Inner {
)]
length: u64,

/// The starting output byte offset, before hex encoding
#[arg(long, default_value_t = 0, value_name("SEEK"))]
seek: u64,

/// The maximum number of threads to use
///
/// By default, this is the number of logical cores. If this flag is
/// omitted, or if its value is 0, RAYON_NUM_THREADS is also respected.
#[arg(long, value_name("NUM"))]
num_threads: Option<usize>,

/// Use the keyed mode, reading the 32-byte key from stdin
#[arg(long, requires("file"))]
keyed: bool,

/// Use the key derivation mode, with the given context string
///
/// Cannot be used with --keyed.
#[arg(long, value_name("CONTEXT"), conflicts_with(KEYED_ARG))]
derive_key: Option<String>,

/// Disable memory mapping
///
/// Currently this also disables multithreading.
Expand Down Expand Up @@ -80,7 +84,7 @@ struct Inner {
)]
check: bool,

/// Skip printing OK for each successfully verified file
/// Skip printing OK for each checked file
///
/// Must be used with --check.
#[arg(long, requires(CHECK_ARG))]
Expand Down Expand Up @@ -146,6 +150,10 @@ impl Args {
self.inner.length
}

fn seek(&self) -> u64 {
self.inner.seek
}

fn keyed(&self) -> bool {
self.inner.keyed
}
Expand Down Expand Up @@ -208,7 +216,9 @@ impl Input {
copy_wide(lock, &mut hasher)?;
}
}
Ok(hasher.finalize_xof())
let mut output_reader = hasher.finalize_xof();
output_reader.set_position(args.seek());
Ok(output_reader)
}
}

Expand Down Expand Up @@ -275,7 +285,9 @@ fn maybe_memmap_file(file: &File) -> Result<Option<memmap2::Mmap>> {
}

fn write_hex_output(mut output: blake3::OutputReader, args: &Args) -> Result<()> {
// Encoding multiples of the block size is most efficient.
// Encoding multiples of the 64 bytes is most efficient.
// TODO: This computes each output block twice when the --seek argument isn't a multiple of 64.
// We'll refactor all of this soon anyway, once SIMD optimizations are available for the XOF.
let mut len = args.len();
let mut block = [0; blake3::guts::BLOCK_LEN];
while len > 0 {
Expand Down
41 changes: 33 additions & 8 deletions b3sum/tests/cli_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -87,18 +87,43 @@ fn test_missing_files() {
}

#[test]
fn test_hash_length() {
let mut buf = [0; 100];
fn test_hash_length_and_seek() {
let mut expected = [0; 100];
blake3::Hasher::new()
.update(b"foo")
.finalize_xof()
.fill(&mut buf);
let expected = format!("{} -", hex::encode(&buf[..]));
let output = cmd!(b3sum_exe(), "--length=100")
.fill(&mut expected);
let output = cmd!(b3sum_exe(), "--raw", "--length=100")
.stdin_bytes("foo")
.read()
.unwrap();
assert_eq!(&*expected, &*output);
.stdout_capture()
.run()
.unwrap()
.stdout;
assert_eq!(expected[..], output);

let short_output = cmd!(b3sum_exe(), "--raw", "--length=99")
.stdin_bytes("foo")
.stdout_capture()
.run()
.unwrap()
.stdout;
assert_eq!(expected[..99], short_output);

let seek1_output = cmd!(b3sum_exe(), "--raw", "--length=99", "--seek=1")
.stdin_bytes("foo")
.stdout_capture()
.run()
.unwrap()
.stdout;
assert_eq!(expected[1..], seek1_output);

let seek99_output = cmd!(b3sum_exe(), "--raw", "--length=1", "--seek=99")
.stdin_bytes("foo")
.stdout_capture()
.run()
.unwrap()
.stdout;
assert_eq!(expected[99..], seek99_output);
}

#[test]
Expand Down

0 comments on commit 4108923

Please sign in to comment.