tftpd/
window.rs

1use std::{
2    collections::VecDeque,
3    error::Error,
4    fs::File,
5    io::{BufRead, BufReader, Read, Write},
6};
7
8/// WindowRead `struct` is used to store chunks of data from a file. It is
9/// used to help store the data that is being sent for the
10/// [RFC 7440](https://www.rfc-editor.org/rfc/rfc7440) Windowsize option.
11///
12/// # Example
13/// ```rust
14/// use std::{fs::{self, OpenOptions, File}, io::Write};
15/// use tftpd::WindowRead;
16///
17/// let mut file = File::create("test.txt").unwrap();
18/// file.write_all(b"Hello, world!").unwrap();
19/// file.flush().unwrap();
20///
21/// let file = File::open("test.txt").unwrap();
22/// let mut window = WindowRead::new(5, 512, file);
23/// window.fill().unwrap();
24/// fs::remove_file("test.txt").unwrap();
25/// ```
26pub struct WindowRead {
27    elements: VecDeque<Vec<u8>>,
28    size: u16,
29    chunk_size: u16,
30    bufreader: BufReader<File>,
31}
32
33impl WindowRead {
34    /// Creates a new `Window` with the supplied size and chunk size.
35    pub fn new(size: u16, chunk_size: u16, file: File) -> WindowRead {
36        WindowRead {
37            elements: VecDeque::new(),
38            size,
39            chunk_size,
40            bufreader: BufReader::with_capacity(
41                2 * size as usize*chunk_size as usize,
42                file,
43            ),
44        }
45    }
46
47    /// Fills the `Window` with chunks of data from the file.
48    /// Returns `true` if the `Window` is full.
49    pub fn fill(&mut self) -> Result<bool, Box<dyn Error>> {
50        for _ in self.len()..self.size {
51            let mut chunk = vec![0; self.chunk_size as usize];
52            let size = self.bufreader.read(&mut chunk)?;
53            if size != self.chunk_size as usize {
54                chunk.truncate(size);
55                self.elements.push_back(chunk);
56                return Ok(false);
57            }
58
59            self.elements.push_back(chunk);
60        }
61
62        Ok(true)
63    }
64
65    /// Fill the read buffer to speed up next window fill
66    pub fn prefill(&mut self) -> Result<(), Box<dyn Error>> {
67        self.bufreader.fill_buf()?;
68        Ok(())
69    }
70
71    /// Removes the first `amount` of elements from the `Window`.
72    pub fn remove(&mut self, amount: u16) -> Result<(), &'static str> {
73        if amount > self.len() {
74            return Err("amount cannot be larger than length of window");
75        }
76
77        drop(self.elements.drain(0..amount as usize));
78
79        Ok(())
80    }
81
82    /// Returns a reference to the `VecDeque` containing the elements.
83    pub fn get_elements(&self) -> &VecDeque<Vec<u8>> {
84        &self.elements
85    }
86
87    /// Returns the length of the `Window`.
88    pub fn len(&self) -> u16 {
89        self.elements.len() as u16
90    }
91
92    /// Returns `true` if the `Window` is empty.
93    pub fn is_empty(&self) -> bool {
94        self.elements.is_empty()
95    }
96}
97
98
99/// WindowWrite `struct` is used to store data and write them in a file. 
100/// It is used to help store the data that is being received for the
101/// [RFC 7440](https://www.rfc-editor.org/rfc/rfc7440) Windowsize option.
102///
103/// # Example
104/// ```rust
105/// use std::{fs::{self, OpenOptions, File}, io::Write};
106/// use tftpd::WindowWrite;
107///
108/// let file = File::create("test.txt").unwrap();
109/// let mut window = WindowWrite::new(5, file);
110/// window.add(vec![0x1, 0x2, 0x3]).unwrap();
111/// window.add(vec![0x4, 0x5, 0x6]).unwrap();
112/// window.empty().unwrap();
113/// ```
114pub struct WindowWrite {
115    elements: VecDeque<Vec<u8>>,
116    size: u16,
117    file: File,
118}
119
120impl WindowWrite {
121    /// Creates a new `Window` with the supplied size and chunk size.
122    pub fn new(size: u16, file: File) -> WindowWrite {
123        WindowWrite {
124            elements: VecDeque::new(),
125            size,
126            file,
127        }
128    }
129
130    /// Empties the `Window` by writing the data to the file.
131    pub fn empty(&mut self) -> Result<(), Box<dyn Error>> {
132        for data in &self.elements {
133            self.file.write_all(data)?;
134        }
135
136        self.elements.clear();
137
138        Ok(())
139    }
140
141    /// Adds a data `Vec<u8>` to the `Window`.
142    pub fn add(&mut self, data: Vec<u8>) -> Result<(), &'static str> {
143        if self.len() == self.size {
144            return Err("cannot add to a full window");
145        }
146
147        self.elements.push_back(data);
148
149        Ok(())
150    }
151
152    /// Returns the length of the `Window`.
153    pub fn len(&self) -> u16 {
154        self.elements.len() as u16
155    }
156
157    /// Returns `true` if the `Window` is empty.
158    pub fn is_empty(&self) -> bool {
159        self.elements.is_empty()
160    }
161
162    /// Returns `true` if the `Window` is full.
163    pub fn is_full(&self) -> bool {
164        self.elements.len() as u16 == self.size
165    }
166
167    /// Returns the length of the file
168    pub fn file_len(&self) -> Result<u64, Box<dyn Error>> {
169        Ok(self.file.metadata()?.len())
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176    use std::{
177        fs::{self, OpenOptions},
178        io::Write,
179    };
180
181    const DIR_NAME: &str = "target/test";
182
183    #[test]
184    fn fills_and_removes_from_window() {
185        const FILENAME: &str = "fills_and_removes_from_window.txt";
186
187        let mut file = initialize(FILENAME);
188        file.write_all(b"Hello, world!").unwrap();
189        file.flush().unwrap();
190        drop(file);
191
192        file = open(FILENAME);
193
194        let mut window = WindowRead::new(2, 5, file);
195        window.fill().unwrap();
196        assert_eq!(window.elements.len(), 2);
197        assert_eq!(window.elements[0], b"Hello"[..]);
198        assert_eq!(window.elements[1], b", wor"[..]);
199
200        window.remove(1).unwrap();
201        assert_eq!(window.elements.len(), 1);
202        assert_eq!(window.elements[0], b", wor"[..]);
203
204        window.fill().unwrap();
205        assert_eq!(window.elements.len(), 2);
206        assert_eq!(window.elements[0], b", wor"[..]);
207        assert_eq!(window.elements[1], b"ld!"[..]);
208
209        clean(FILENAME);
210    }
211
212    #[test]
213    fn adds_to_and_empties_window() {
214        const FILENAME: &str = "adds_to_and_empties_window.txt";
215
216        let file = initialize(FILENAME);
217
218        let mut window = WindowWrite::new(3, file);
219        window.add(b"Hello".to_vec()).unwrap();
220        assert_eq!(window.elements.len(), 1);
221        assert_eq!(window.elements[0], b"Hello"[..]);
222
223        window.add(b", wor".to_vec()).unwrap();
224        assert_eq!(window.elements.len(), 2);
225        assert_eq!(window.elements[0], b"Hello"[..]);
226        assert_eq!(window.elements[1], b", wor"[..]);
227
228        window.add(b"ld!".to_vec()).unwrap();
229        assert_eq!(window.elements.len(), 3);
230        assert_eq!(window.elements[0], b"Hello"[..]);
231        assert_eq!(window.elements[1], b", wor"[..]);
232        assert_eq!(window.elements[2], b"ld!"[..]);
233
234        window.empty().unwrap();
235        assert_eq!(window.elements.len(), 0);
236
237        let mut contents = Default::default();
238        File::read_to_string(
239            &mut File::open(DIR_NAME.to_string() + "/" + FILENAME).unwrap(),
240            &mut contents,
241        )
242        .unwrap();
243        assert_eq!(contents, "Hello, world!");
244
245        clean(FILENAME);
246    }
247
248    fn initialize(filename: &str) -> File {
249        let filename = DIR_NAME.to_string() + "/" + filename;
250
251        let _ = fs::create_dir_all(DIR_NAME);
252
253        if File::open(&filename).is_ok() {
254            fs::remove_file(&filename).unwrap();
255        }
256
257        OpenOptions::new()
258            .read(true)
259            .append(true)
260            .create(true)
261            .open(&filename)
262            .unwrap()
263    }
264
265    fn open(filename: &str) -> File {
266        let filename = DIR_NAME.to_string() + "/" + filename;
267
268        OpenOptions::new()
269            .read(true)
270            .append(true)
271            .create(true)
272            .open(filename)
273            .unwrap()
274    }
275
276    fn clean(filename: &str) {
277        let filename = DIR_NAME.to_string() + "/" + filename;
278        fs::remove_file(filename).unwrap();
279        if fs::remove_dir(DIR_NAME).is_err() {
280            // ignore removing directory, as other tests are
281            // still running
282        }
283    }
284}
tftpd/window.rs

tftpd/
window.rs
