dd_lib/io/
read.rs

1use opts::{self, Error::ConflictingOption, IFlag};
2use results::{Error, Result};
3use std::{
4    fs::File,
5    io::{BufRead, Read, Seek, SeekFrom, Stdin, StdinLock},
6};
7
8/// Readers read blocks from the specified file or [Stdin][std::io::Stdin]
9pub struct Reader<'a, R: Read> {
10    inner: InnerReader<'a, R>,
11    /// internal buffer to implement BufRead
12    pub buf: Vec<u8>,
13    /// the front of the buffer; beyond here is free space
14    pub front: usize,
15    /// the remaining bytes to read. None means to continue reading until EOF.
16    pub remaining_bytes: Option<usize>,
17    /// the mode (`Standard`, `Block`, or `Unblock`) that dd will run in
18    pub mode: opts::Mode,
19    /// if true, pad every read blocks to self.buf.len() size
20    pub sync: bool,
21}
22
23#[derive(Debug)]
24pub enum InnerReader<'a, R: Read> {
25    Stdin(StdinLock<'a>),
26    R(R),
27}
28
29impl<'a> Reader<'a, File> {
30    pub fn new(stdin: &'a Stdin, o: &opts::Opts) -> Result<Self> {
31        // let fullblock = opts.iflags.contains(IFlag::FULLBLOCK),
32
33        Ok(Reader {
34            buf: vec![0; o.input_block_size],
35            front: 0,
36            mode: o.mode,
37            inner: InnerReader::new(stdin, o)?,
38            remaining_bytes: match (o.count, o.iflag(IFlag::COUNT_BYTES)) {
39                (Some(bytes), true) => Some(bytes),
40                (Some(blocks), false) => Some(blocks * o.input_block_size),
41                _ => None,
42            },
43            sync: o.iflag(IFlag::SYNC),
44        })
45    }
46}
47impl<'a, R: Read> Reader<'a, R> {
48    pub fn basic(inner: InnerReader<'a, R>, block_size: usize) -> Self {
49        Reader {
50            buf: vec![0; block_size],
51            front: 0,
52            mode: opts::Mode::Standard,
53            inner,
54            remaining_bytes: None,
55            sync: false,
56        }
57    }
58}
59
60impl<'a, R: Read> Read for Reader<'a, R> {
61    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> { self.inner.read(buf) }
62}
63
64impl<'a, R: Read> BufRead for Reader<'a, R> {
65    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
66        self.front += match self.remaining_bytes {
67            Some(0) => return Ok(&self.buf[0..0]),
68            Some(n) => {
69                let back = usize::min(n, self.buf.len());
70                self.inner.read(&mut self.buf[self.front..back])?
71            },
72            None => self.inner.read(&mut self.buf[self.front..])?,
73        };
74        match (self.sync, self.mode) {
75            (true, opts::Mode::Standard) => {
76                self.buf.iter_mut().for_each(|b| *b = 0);
77                Ok(&self.buf)
78            },
79            (true, _) => {
80                self.buf.iter_mut().for_each(|b| *b = b' ');
81                Ok(&self.buf)
82            },
83            (false, _) => Ok(&self.buf[..self.front]),
84        }
85    }
86
87    fn consume(&mut self, bytes: usize) { self.front -= bytes }
88}
89
90impl<'a> InnerReader<'a, File> {
91    fn new(stdin: &'a Stdin, o: &opts::Opts) -> Result<Self> {
92        let ibs = o.input_block_size;
93        match &o.input_file {
94            None if o.iflag(IFlag::APPEND) => {
95                return Err(Error::from(ConflictingOption(
96                    "iflag 'append' is only usable with an input file",
97                )))
98            },
99            None if o.input_seek.is_some() || o.iflag(IFlag::SKIP_BYTES) => {
100                return Err(Error::from(ConflictingOption(
101                    "options 'input_seek' and 'seek_bytes' are only usable with an input file",
102                )))
103            },
104
105            // valid stdin
106            None => Ok(InnerReader::Stdin(stdin.lock())),
107            Some(path) => {
108                let mut f = File::open(path)?;
109                match (o.input_seek, o.iflag(IFlag::SKIP_BYTES)) {
110                    (None, true) => {
111                        return Err(Error::from(ConflictingOption(
112                            "cannot use 'seek_bytes' with no 'input_seek' distance selected",
113                        )))
114                    },
115                    (None, false) => {},
116                    (Some(bytes), true) => {
117                        f.seek(SeekFrom::Start(bytes as u64))?;
118                    },
119                    (Some(blocks), false) => {
120                        f.seek(SeekFrom::Start(blocks as u64 * ibs as u64))?;
121                    },
122                };
123                Ok(InnerReader::R(f))
124            },
125        }
126    }
127}
128
129impl<'a, T: BufRead> From<StdinLock<'a>> for InnerReader<'a, T> {
130    fn from(s: StdinLock) -> InnerReader<T> { InnerReader::Stdin(s) }
131}
132
133impl<'a> From<File> for InnerReader<'a, File> {
134    fn from(f: File) -> InnerReader<'a, File> { InnerReader::R(f) }
135}
136
137impl<'a, T: Read> Read for InnerReader<'a, T> {
138    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
139        match self {
140            InnerReader::R(f) => f.read(buf),
141            InnerReader::Stdin(s) => s.read(buf),
142        }
143    }
144}