dd_lib/io/
write.rs

1use opts::{CFlag, Error::ConflictingOption, OFlag, Opts};
2use results::{Error, Result};
3use std::{
4    fs::{File, OpenOptions},
5    io::{self, BufWriter, Seek, SeekFrom, Stdout, StdoutLock, Write},
6    path::Path,
7};
8/// Writers write to the specified [File][std::fs::File] or
9/// [Stdout][std::io::Stdout]
10pub struct Writer<'a, W: Write> {
11    w: Inner<'a, W>,
12    sparse: bool,
13}
14
15enum Inner<'a, W: Write> {
16    W(BufWriter<W>),
17    Stdout(BufWriter<StdoutLock<'a>>),
18}
19
20impl<'a, W: Write + Seek> Write for Writer<'a, W> {
21    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
22        match (&mut self.w, self.sparse && is_null(buf)) {
23            (Inner::W(w), true) => {
24                w.seek(SeekFrom::Current(buf.len() as i64))?;
25                Ok(buf.len())
26            },
27            // "skip" ahead by not not writing to stdout at all
28            (Inner::Stdout(_), true) => Ok(buf.len()),
29            (w, _) => w.write(buf),
30        }
31    }
32
33    fn flush(&mut self) -> std::io::Result<()> { self.w.flush() }
34}
35
36impl<'a, W: Write> Writer<'a, W> {
37    fn from_stdout(s: &'a Stdout, cap: usize, sparse: bool) -> Self {
38        let w = Inner::Stdout(BufWriter::with_capacity(cap, s.lock()));
39        Writer { w, sparse }
40    }
41
42    fn from_w(w: W, cap: usize, sparse: bool) -> Self {
43        let w = Inner::W(BufWriter::with_capacity(cap, w));
44        Writer { w, sparse }
45    }
46}
47
48impl<'a> Writer<'a, File> {
49    pub fn new(stdout: &'a Stdout, o: &Opts) -> Result<Self> {
50        match (o.output_file.clone(), o.output_seek) {
51            (Some(ref path), _) if o.oflag(OFlag::DIRECTORY) && !Path::is_dir(path.as_ref()) => Err(Error::from(
52                ConflictingOption("output flag 'directory' specified, but the path was not a directory"),
53            )),
54            (None, Some(_)) => Err(Error::from(ConflictingOption("cannot 'seek' on stdout"))),
55            (None, None) if o.oflag(OFlag::APPEND) => Err(Error::from(ConflictingOption("cannot 'append' on stdout"))),
56
57            (None, None) => Ok(Writer::from_stdout(stdout, o.input_block_size, o.cflag(CFlag::SPARSE))),
58            (Some(path), None) => {
59                let file = open(&path, o.oflag(OFlag::APPEND), !o.cflag(CFlag::NOTRUNC))?;
60                Ok(Writer::from_w(file, o.input_block_size, o.cflag(CFlag::SPARSE)))
61            },
62            (Some(path), Some(blocks)) => {
63                let mut file = open(&path, o.oflag(OFlag::APPEND), !o.cflag(CFlag::NOTRUNC))?;
64                file.seek(SeekFrom::Start(blocks as u64 * o.input_block_size as u64))?;
65                Ok(Writer::from_w(file, o.input_block_size, o.cflag(CFlag::SPARSE)))
66            },
67        }
68    }
69
70    #[allow(dead_code)]
71    // get a reference to the file inside. this isn't safe for all varieties!
72    pub unsafe fn get_file(self) -> File {
73        match self.w {
74            Inner::W(file) => file.into_inner().unwrap(),
75            _ => panic!("shouldn't have called this!"),
76        }
77    }
78}
79
80/// open a file for creation and writing, appending and/or truncating if
81/// specified
82fn open(path: &str, append: bool, truncate: bool) -> Result<File> {
83    let file = OpenOptions::new()
84        .create(true)
85        .write(true)
86        .append(append)
87        .truncate(truncate)
88        .open(path)?;
89    Ok(file)
90}
91
92impl<'a, W: Write> Write for Inner<'a, W> {
93    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
94        match self {
95            Inner::W(w) => w.write(buf),
96            Inner::Stdout(stdout) => stdout.write(buf),
97        }
98    }
99
100    fn flush(&mut self) -> io::Result<()> {
101        match self {
102            Inner::W(w) => w.flush(),
103            Inner::Stdout(stdout) => stdout.flush(),
104        }
105    }
106}
107/// all bytes are zero
108fn is_null(b: &[u8]) -> bool { b.iter().all(|b| *b == 0) }
dd_lib/io/write.rs

dd_lib/io/
write.rs
