Mirrors/coreutils
2
0
Fork 0
mirror of https://github.com/uutils/coreutils synced 2024-12-13 23:02:38 +00:00
Code Issues Projects Releases Packages Wiki Activity

dd: create Dest enum and simpler Output struct

Browse source
This commit is contained in:
Jeffrey Finkelstein 2022-11-12 14:54:43 -05:00 committed by Sylvestre Ledru
parent 8e1742f45a
commit 4621557ce7
1 changed files with 158 additions and 132 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

290
src/uu/dd/src/dd.rs
View file

@ -25,7 +25,7 @@ use std::cmp;
use std::env;
use std::ffi::OsString;
use std::fs::{File, OpenOptions};
use std::io::{self, Read, Seek, SeekFrom, Write};
use std::io::{self, Read, Seek, SeekFrom, Stdout, Write};
#[cfg(any(target_os = "linux", target_os = "android"))]
use std::os::unix::fs::OpenOptionsExt;
use std::path::Path;
@ -280,51 +280,161 @@ impl<'a, R: Read> Input<'a, R> {
}
}
trait OutputTrait: Sized + Write {
fn fsync(&mut self) -> io::Result<()>;
fn fdatasync(&mut self) -> io::Result<()>;
enum Density {
Sparse,
Dense,
}
struct Output<'a, W: Write> {
dst: W,
settings: &'a Settings,
/// Data destinations.
enum Dest {
/// Output to stdout.
Stdout(Stdout),
/// Output to a file.
///
/// The [`Density`] component indicates whether to attempt to
/// write a sparse file when all-zero blocks are encountered.
File(File, Density),
}
impl<'a> Output<'a, io::Stdout> {
fn new(settings: &'a Settings) -> UResult<Self> {
let mut dst = io::stdout();
// stdout is not seekable, so we just write null bytes.
if settings.seek > 0 {
io::copy(&mut io::repeat(0u8).take(settings.seek as u64), &mut dst)
.map_err_context(|| String::from("write error"))?;
}
Ok(Self { dst, settings })
}
}
impl<'a> OutputTrait for Output<'a, io::Stdout> {
impl Dest {
fn fsync(&mut self) -> io::Result<()> {
self.dst.flush()
match self {
Self::Stdout(stdout) => stdout.flush(),
Self::File(f, _) => {
f.flush()?;
f.sync_all()
}
}
}
fn fdatasync(&mut self) -> io::Result<()> {
self.dst.flush()
match self {
Self::Stdout(stdout) => stdout.flush(),
Self::File(f, _) => {
f.flush()?;
f.sync_data()
}
}
}
fn seek(&mut self, n: u64) -> io::Result<u64> {
match self {
Self::Stdout(stdout) => io::copy(&mut io::repeat(0).take(n), stdout),
Self::File(f, _) => f.seek(io::SeekFrom::Start(n)),
}
}
}
impl<'a, W: Write> Output<'a, W>
where
Self: OutputTrait,
{
/// Decide whether the given buffer is all zeros.
fn is_sparse(buf: &[u8]) -> bool {
buf.iter().all(|&e| e == 0u8)
}
impl Write for Dest {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
match self {
Self::File(f, Density::Sparse) if is_sparse(buf) => {
let seek_amt: i64 = buf
.len()
.try_into()
.expect("Internal dd Error: Seek amount greater than signed 64-bit integer");
f.seek(io::SeekFrom::Current(seek_amt))?;
Ok(buf.len())
}
Self::File(f, _) => f.write(buf),
Self::Stdout(stdout) => stdout.write(buf),
}
}
fn flush(&mut self) -> io::Result<()> {
match self {
Self::Stdout(stdout) => stdout.flush(),
Self::File(f, _) => f.flush(),
}
}
}
/// The destination of the data, configured with the given settings.
///
/// Use the [`Output::new_stdout`] or [`Output::new_file`] functions
/// to construct a new instance of this struct. Then use the
/// [`Output::dd_out`] function to execute the main copy operation for
/// `dd`.
struct Output<'a> {
/// The destination to which bytes will be written.
dst: Dest,
/// Configuration settings for how to read and write the data.
settings: &'a Settings,
}
impl<'a> Output<'a> {
/// Instantiate this struct with stdout as a destination.
fn new_stdout(settings: &'a Settings) -> UResult<Self> {
let mut dst = Dest::Stdout(io::stdout());
dst.seek(settings.seek)
.map_err_context(|| "write error".to_string())?;
Ok(Self { dst, settings })
}
/// Instantiate this struct with the named file as a destination.
fn new_file(filename: &Path, settings: &'a Settings) -> UResult<Self> {
fn open_dst(path: &Path, cflags: &OConvFlags, oflags: &OFlags) -> Result<File, io::Error> {
let mut opts = OpenOptions::new();
opts.write(true)
.create(!cflags.nocreat)
.create_new(cflags.excl)
.append(oflags.append);
#[cfg(any(target_os = "linux", target_os = "android"))]
if let Some(libc_flags) = make_linux_oflags(oflags) {
opts.custom_flags(libc_flags);
}
opts.open(path)
}
let dst = open_dst(filename, &settings.oconv, &settings.oflags)
.map_err_context(|| format!("failed to open {}", filename.quote()))?;
// Seek to the index in the output file, truncating if requested.
//
// Calling `set_len()` may result in an error (for example,
// when calling it on `/dev/null`), but we don't want to
// terminate the process when that happens. Instead, we
// suppress the error by calling `Result::ok()`. This matches
// the behavior of GNU `dd` when given the command-line
// argument `of=/dev/null`.
if !settings.oconv.notrunc {
dst.set_len(settings.seek).ok();
}
let density = if settings.oconv.sparse {
Density::Sparse
} else {
Density::Dense
};
let mut dst = Dest::File(dst, density);
dst.seek(settings.seek)
.map_err_context(|| "failed to seek in output file".to_string())?;
Ok(Self { dst, settings })
}
/// Write the given bytes one block at a time.
///
/// This may write partial blocks (for example, if the underlying
/// call to [`Write::write`] writes fewer than `buf.len()`
/// bytes). The returned [`WriteStat`] object will include the
/// number of partial and complete blocks written during execution
/// of this function.
fn write_blocks(&mut self, buf: &[u8]) -> io::Result<WriteStat> {
let mut writes_complete = 0;
let mut writes_partial = 0;
let mut bytes_total = 0;
for chunk in buf.chunks(self.settings.obs) {
let wlen = self.write(chunk)?;
let wlen = self.dst.write(chunk)?;
if wlen < self.settings.obs {
writes_partial += 1;
} else {
@ -343,9 +453,9 @@ where
/// Flush the output to disk, if configured to do so.
fn sync(&mut self) -> std::io::Result<()> {
if self.settings.oconv.fsync {
self.fsync()
self.dst.fsync()
} else if self.settings.oconv.fdatasync {
self.fdatasync()
self.dst.fdatasync()
} else {
// Intentionally do nothing in this case.
Ok(())
@ -511,96 +621,6 @@ fn make_linux_oflags(oflags: &OFlags) -> Option<libc::c_int> {
}
}
impl<'a> Output<'a, File> {
fn new(filename: &Path, settings: &'a Settings) -> UResult<Self> {
fn open_dst(path: &Path, cflags: &OConvFlags, oflags: &OFlags) -> Result<File, io::Error> {
let mut opts = OpenOptions::new();
opts.write(true)
.create(!cflags.nocreat)
.create_new(cflags.excl)
.append(oflags.append);
#[cfg(any(target_os = "linux", target_os = "android"))]
if let Some(libc_flags) = make_linux_oflags(oflags) {
opts.custom_flags(libc_flags);
}
opts.open(path)
}
let mut dst = open_dst(filename, &settings.oconv, &settings.oflags)
.map_err_context(|| format!("failed to open {}", filename.quote()))?;
// Seek to the index in the output file, truncating if requested.
//
// Calling `set_len()` may result in an error (for example,
// when calling it on `/dev/null`), but we don't want to
// terminate the process when that happens. Instead, we
// suppress the error by calling `Result::ok()`. This matches
// the behavior of GNU `dd` when given the command-line
// argument `of=/dev/null`.
if !settings.oconv.notrunc {
dst.set_len(settings.seek).ok();
}
dst.seek(io::SeekFrom::Start(settings.seek))
.map_err_context(|| "failed to seek in output file".to_string())?;
Ok(Self { dst, settings })
}
}
impl<'a> OutputTrait for Output<'a, File> {
fn fsync(&mut self) -> io::Result<()> {
self.dst.flush()?;
self.dst.sync_all()
}
fn fdatasync(&mut self) -> io::Result<()> {
self.dst.flush()?;
self.dst.sync_data()
}
}
impl<'a> Seek for Output<'a, File> {
fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
self.dst.seek(pos)
}
}
impl<'a> Write for Output<'a, File> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
fn is_sparse(buf: &[u8]) -> bool {
buf.iter().all(|&e| e == 0u8)
}
// -----------------------------
if self.settings.oconv.sparse && is_sparse(buf) {
let seek_amt: i64 = buf
.len()
.try_into()
.expect("Internal dd Error: Seek amount greater than signed 64-bit integer");
self.dst.seek(io::SeekFrom::Current(seek_amt))?;
Ok(buf.len())
} else {
self.dst.write(buf)
}
}
fn flush(&mut self) -> io::Result<()> {
self.dst.flush()
}
}
impl<'a> Write for Output<'a, io::Stdout> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.dst.write(buf)
}
fn flush(&mut self) -> io::Result<()> {
self.dst.flush()
}
}
/// Read helper performs read operations common to all dd reads, and dispatches the buffer to relevant helper functions as dictated by the operations requested by the user.
fn read_helper<R: Read>(
i: &mut Input<R>,
@ -753,22 +773,22 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
match (&settings.infile, &settings.outfile) {
(Some(infile), Some(outfile)) => {
let i = Input::<File>::new(Path::new(&infile), &settings)?;
let o = Output::<File>::new(Path::new(&outfile), &settings)?;
let o = Output::new_file(Path::new(&outfile), &settings)?;
o.dd_out(i).map_err_context(|| "IO error".to_string())
}
(None, Some(outfile)) => {
let i = Input::<io::Stdin>::new(&settings)?;
let o = Output::<File>::new(Path::new(&outfile), &settings)?;
let o = Output::new_file(Path::new(&outfile), &settings)?;
o.dd_out(i).map_err_context(|| "IO error".to_string())
}
(Some(infile), None) => {
let i = Input::<File>::new(Path::new(&infile), &settings)?;
if is_stdout_redirected_to_seekable_file() {
let filename = stdout_canonicalized();
let o = Output::<File>::new(Path::new(&filename), &settings)?;
let o = Output::new_file(Path::new(&filename), &settings)?;
o.dd_out(i).map_err_context(|| "IO error".to_string())
} else {
let o = Output::<io::Stdout>::new(&settings)?;
let o = Output::new_stdout(&settings)?;
o.dd_out(i).map_err_context(|| "IO error".to_string())
}
}
@ -776,10 +796,10 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
let i = Input::<io::Stdin>::new(&settings)?;
if is_stdout_redirected_to_seekable_file() {
let filename = stdout_canonicalized();
let o = Output::<File>::new(Path::new(&filename), &settings)?;
let o = Output::new_file(Path::new(&filename), &settings)?;
o.dd_out(i).map_err_context(|| "IO error".to_string())
} else {
let o = Output::<io::Stdout>::new(&settings)?;
let o = Output::new_stdout(&settings)?;
o.dd_out(i).map_err_context(|| "IO error".to_string())
}
}
@ -798,7 +818,7 @@ pub fn uu_app() -> Command {
#[cfg(test)]
mod tests {
use crate::datastructures::{IConvFlags, IFlags};
use crate::{calc_bsize, Input, Output, Parser, Settings};
use crate::{calc_bsize, Density, Dest, Input, Output, Parser, Settings};
use std::cmp;
use std::fs;
@ -893,7 +913,7 @@ mod tests {
let args = &["conv=nocreat", "of=not-a-real.file"];
let settings = Parser::new().parse(args).unwrap();
assert!(
Output::<File>::new(Path::new(settings.outfile.as_ref().unwrap()), &settings).is_err()
Output::new_file(Path::new(settings.outfile.as_ref().unwrap()), &settings).is_err()
);
}
@ -917,7 +937,10 @@ mod tests {
};
let output = Output {
dst: File::create("./test-resources/FAILED-deadbeef-16-delayed.test").unwrap(),
dst: Dest::File(
File::create("./test-resources/FAILED-deadbeef-16-delayed.test").unwrap(),
Density::Dense,
),
settings: &settings,
};
@ -965,8 +988,11 @@ mod tests {
};
let output = Output {
dst: File::create("./test-resources/FAILED-random_73k_test_lazy_fullblock.test")
.unwrap(),
dst: Dest::File(
File::create("./test-resources/FAILED-random_73k_test_lazy_fullblock.test")
.unwrap(),
Density::Dense,
),
settings: &settings,
};