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

ntfs_reader/
aligned_reader.rs
