ntfs_reader/
aligned_reader.rs

1// Copyright (c) 2022, Matteo Bernacchia <dev@kikijiki.com>. All rights reserved.
2// This project is dual licensed under the Apache License 2.0 and the MIT license.
3// See the LICENSE files in the project root for details.
4
5use std::fs::File;
6use std::io::{self, BufReader};
7use std::io::{Read, Seek, SeekFrom};
8use std::path::Path;
9
10pub struct AlignedReader<R>
11where
12    R: Read + Seek,
13{
14    inner: R,
15    alignment: u64,
16    position: u64,
17
18    buffer_pos: u64,
19    buffer_size: usize,
20    buffer: Vec<u8>,
21}
22
23impl<R> AlignedReader<R>
24where
25    R: Read + Seek,
26{
27    pub fn new(inner: R, alignment: u64) -> io::Result<Self> {
28        if !alignment.is_power_of_two() {
29            return Err(io::Error::new(
30                io::ErrorKind::InvalidInput,
31                "alignment must be a power of two",
32            ));
33        }
34
35        Ok(Self {
36            inner,
37            alignment,
38            position: 0,
39            buffer_pos: 0,
40            buffer_size: 0,
41            buffer: Vec::with_capacity(alignment as usize),
42        })
43    }
44
45    fn round_down(&self, n: u64) -> u64 {
46        n / self.alignment * self.alignment
47    }
48
49    fn round_up(&self, n: u64) -> u64 {
50        if n.is_multiple_of(self.alignment) {
51            n
52        } else {
53            self.round_down(n) + self.alignment
54        }
55    }
56}
57
58impl<R> Read for AlignedReader<R>
59where
60    R: Read + Seek,
61{
62    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
63        let aligned_position = self.round_down(self.position);
64
65        let start = (self.position - aligned_position) as usize;
66        let to_read = buf.len().min(self.alignment as usize - start);
67        let end = start + to_read;
68        let size = self.round_up(end as u64) as usize;
69
70        if aligned_position != self.buffer_pos || size > self.buffer_size {
71            self.inner.seek(SeekFrom::Start(aligned_position))?;
72            self.buffer.resize(size, 0u8);
73            self.inner.read_exact(&mut self.buffer)?;
74            self.buffer_pos = aligned_position;
75            self.buffer_size = size;
76        }
77
78        buf[..to_read].copy_from_slice(&self.buffer[start..end]);
79
80        self.position += to_read as u64;
81        Ok(to_read)
82    }
83}
84
85impl<R> Seek for AlignedReader<R>
86where
87    R: Read + Seek,
88{
89    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
90        let raw_position = match pos {
91            SeekFrom::Start(n) => Some(n),
92            SeekFrom::End(_) => {
93                return Err(io::Error::other("unsupported"));
94            }
95            SeekFrom::Current(n) => {
96                if n >= 0 {
97                    self.position.checked_add(n as u64)
98                } else {
99                    self.position.checked_sub(n.wrapping_neg() as u64)
100                }
101            }
102        };
103
104        match raw_position {
105            Some(n) => {
106                let aligned_position = self.round_down(n);
107                self.inner.seek(SeekFrom::Start(aligned_position))?;
108                self.position = n;
109                Ok(n)
110            }
111            None => Err(io::Error::new(
112                io::ErrorKind::InvalidInput,
113                "invalid position",
114            )),
115        }
116    }
117}
118
119pub fn open_volume(path: &Path) -> std::io::Result<BufReader<AlignedReader<File>>> {
120    let file = File::open(path)?;
121    let sr = AlignedReader::new(file, 4096u64)?;
122    let mut reader = BufReader::new(sr);
123
124    reader.seek(SeekFrom::Start(0))?;
125    Ok(reader)
126}
ntfs_reader/aligned_reader.rs

ntfs_reader/
aligned_reader.rs
