diff --git a/src/uu/dd/src/dd.rs b/src/uu/dd/src/dd.rs index 163d10c942..f17d3cc1b7 100644 --- a/src/uu/dd/src/dd.rs +++ b/src/uu/dd/src/dd.rs @@ -1232,7 +1232,7 @@ fn dd_copy(mut i: Input, o: Output) -> io::Result<()> { // As an optimization, make an educated guess about the // best buffer size for reading based on the number of // blocks already read and the number of blocks remaining. - let loop_bsize = calc_loop_bsize(i.settings.count, &rstat, &wstat, i.settings.ibs, bsize); + let loop_bsize = calc_loop_bsize(i.settings.count, &rstat, i.settings.ibs, bsize); let rstat_update = read_helper(&mut i, &mut buf, loop_bsize)?; if rstat_update.is_empty() { if input_nocache { @@ -1426,13 +1426,7 @@ fn calc_bsize(ibs: usize, obs: usize) -> usize { /// Calculate the buffer size appropriate for this loop iteration, respecting /// a `count=N` if present. -fn calc_loop_bsize( - count: Option, - rstat: &ReadStat, - wstat: &WriteStat, - ibs: usize, - ideal_bsize: usize, -) -> usize { +fn calc_loop_bsize(count: Option, rstat: &ReadStat, ibs: usize, ideal_bsize: usize) -> usize { match count { Some(Num::Blocks(rmax)) => { let rsofar = rstat.reads_complete + rstat.reads_partial; @@ -1440,9 +1434,9 @@ fn calc_loop_bsize( cmp::min(ideal_bsize as u64, rremain * ibs as u64) as usize } Some(Num::Bytes(bmax)) => { - let bmax: u128 = bmax.into(); - let bremain: u128 = bmax - wstat.bytes_total; - cmp::min(ideal_bsize as u128, bremain) as usize + // `iflag=count_bytes` limits input, so use bytes read. + let bremain = bmax.saturating_sub(rstat.bytes_total); + cmp::min(ideal_bsize as u64, bremain) as usize } None => ideal_bsize, } @@ -1680,4 +1674,21 @@ mod tests { Output::new_file(Path::new(settings.outfile.as_ref().unwrap()), &settings).is_err() ); } + + #[test] + fn test_calc_loop_bsize_count_bytes() { + use crate::progress::ReadStat; + use crate::{Num, calc_loop_bsize}; + + for (bytes_read, expected) in [(512, 488), (1000, 0), (1001, 0)] { + let rstat = ReadStat { + bytes_total: bytes_read, + ..ReadStat::default() + }; + assert_eq!( + calc_loop_bsize(Some(Num::Bytes(1000)), &rstat, 512, 512), + expected + ); + } + } } diff --git a/tests/by-util/test_dd.rs b/tests/by-util/test_dd.rs index 6ba3bb8257..9c6d1d7b23 100644 --- a/tests/by-util/test_dd.rs +++ b/tests/by-util/test_dd.rs @@ -2189,3 +2189,23 @@ fn test_bs_not_positive() { } } } + +#[test] +fn test_count_bytes_with_expanding_block_conv() { + let (at, mut ucmd) = at_and_ucmd!(); + let mut input = vec![b'a'; 1000]; + input.extend([b'Z'; 24]); + at.write_bytes("input.txt", &input); + ucmd.args(&[ + "if=input.txt", + "of=output.bin", + "conv=block", + "cbs=1024", + "count=1000", + "iflag=count_bytes", + ]) + .succeeds(); + let output = at.read_bytes("output.bin"); + assert_eq!(bytecount::count(&output, b'a'), 1000); + assert!(!output.contains(&b'Z')); +}