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        assert!(alignment.is_power_of_two());
29
30        Ok(Self {
31            inner,
32            alignment,
33            position: 0,
34            buffer_pos: 0,
35            buffer_size: 0,
36            buffer: Vec::with_capacity(alignment as usize),
37        })
38    }
39
40    fn round_down(&self, n: u64) -> u64 {
41        n / self.alignment * self.alignment
42    }
43
44    fn round_up(&self, n: u64) -> u64 {
45        self.round_down(n) + self.alignment
46    }
47}
48
49impl<R> Read for AlignedReader<R>
50where
51    R: Read + Seek,
52{
53    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
54        let aligned_position = self.round_down(self.position);
55
56        let start = self.position as usize - aligned_position as usize;
57        let end = start + buf.len();
58        let size = self.round_up(end as u64) as usize;
59
60        if aligned_position != self.buffer_pos || size > self.buffer_size {
61            self.buffer.resize(size, 0u8);
62            self.inner.read_exact(&mut self.buffer)?;
63            self.buffer_pos = aligned_position;
64            self.buffer_size = size;
65        }
66
67        buf.copy_from_slice(&self.buffer[start..end]);
68
69        self.position += buf.len() as u64;
70        Ok(buf.len())
71    }
72}
73
74impl<R> Seek for AlignedReader<R>
75where
76    R: Read + Seek,
77{
78    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
79        let raw_position = match pos {
80            SeekFrom::Start(n) => Some(n),
81            SeekFrom::End(_) => {
82                return Err(io::Error::other("unsupported"));
83            }
84            SeekFrom::Current(n) => {
85                if n >= 0 {
86                    self.position.checked_add(n as u64)
87                } else {
88                    self.position.checked_sub(n.wrapping_neg() as u64)
89                }
90            }
91        };
92
93        match raw_position {
94            Some(n) => {
95                let aligned_position = self.round_down(n);
96                self.inner.seek(SeekFrom::Start(aligned_position))?;
97                self.position = n;
98                Ok(n)
99            }
100            None => Err(io::Error::new(
101                io::ErrorKind::InvalidInput,
102                "invalid position",
103            )),
104        }
105    }
106}
107
108pub fn open_volume(path: &Path) -> std::io::Result<BufReader<AlignedReader<File>>> {
109    let file = File::open(path)?;
110    let sr = AlignedReader::new(file, 4096u64)?;
111    let mut reader = BufReader::new(sr);
112
113    reader.seek(SeekFrom::Start(0))?;
114    Ok(reader)
115}
ntfs_reader/aligned_reader.rs

ntfs_reader/
aligned_reader.rs
