Skip to main content

orbis_pfs/
image.rs

1use aes::Aes128;
2use hmac::{Hmac, Mac};
3use sha2::Sha256;
4use std::cmp::min;
5use std::io;
6use xts_mode::{Xts128, get_tweak_default};
7
8pub(crate) const XTS_BLOCK_SIZE: usize = 0x1000;
9
10/// Encapsulates a PFS image with positional read support.
11///
12/// This trait provides thread-safe, stateless access to PFS image data.
13/// Unlike `Read + Seek`, each call specifies its own offset, enabling
14/// concurrent reads from multiple threads without synchronization.
15pub trait Image: Send + Sync {
16    /// Reads bytes from the image at the given offset into `buf`.
17    ///
18    /// Returns the number of bytes actually read. A short read indicates
19    /// the end of the image was reached.
20    fn read_at(&self, offset: u64, output_buf: &mut [u8]) -> io::Result<usize>;
21
22    /// Reads exactly `buf.len()` bytes from `image` at `offset`.
23    ///
24    /// Returns [`io::ErrorKind::UnexpectedEof`] if the image ends before the buffer
25    /// is filled.
26    fn read_exact_at(&self, offset: u64, output_buf: &mut [u8]) -> io::Result<()> {
27        let mut total = 0;
28
29        while total < output_buf.len() {
30            let n = self.read_at(offset + total as u64, &mut output_buf[total..])?;
31
32            if n == 0 {
33                return Err(io::Error::new(
34                    io::ErrorKind::UnexpectedEof,
35                    "unexpected EOF in image",
36                ));
37            }
38
39            total += n;
40        }
41
42        Ok(())
43    }
44
45    /// Returns the total length of the image in bytes.
46    fn len(&self) -> u64;
47
48    /// Returns `true` if the image is empty.
49    fn is_empty(&self) -> bool {
50        self.len() == 0
51    }
52}
53
54/// Gets data key and tweak key from EKPFS and seed.
55pub(crate) fn get_xts_keys(ekpfs: &[u8], seed: &[u8; 16]) -> ([u8; 16], [u8; 16]) {
56    let mut hmac = Hmac::<Sha256>::new_from_slice(ekpfs).unwrap();
57    hmac.update(&[0x01, 0x00, 0x00, 0x00]);
58    hmac.update(seed);
59
60    let secret = hmac.finalize().into_bytes();
61    let mut data_key: [u8; 16] = Default::default();
62    let mut tweak_key: [u8; 16] = Default::default();
63
64    tweak_key.copy_from_slice(&secret[..16]);
65    data_key.copy_from_slice(&secret[16..]);
66
67    (data_key, tweak_key)
68}
69
70/// Unencrypted PFS image backed by a byte slice.
71///
72/// Reads are pure slice indexing — no locks, no allocation, no state.
73pub(crate) struct UnencryptedSlice<'a> {
74    data: &'a [u8],
75}
76
77impl<'a> UnencryptedSlice<'a> {
78    pub fn new(data: &'a [u8]) -> Self {
79        Self { data }
80    }
81}
82
83impl Image for UnencryptedSlice<'_> {
84    fn read_at(&self, offset: u64, output_buf: &mut [u8]) -> io::Result<usize> {
85        let start = offset as usize;
86
87        if start >= self.data.len() {
88            return Ok(0);
89        }
90
91        let available = self.data.len() - start;
92        let n = min(output_buf.len(), available);
93
94        output_buf[..n].copy_from_slice(&self.data[start..start + n]);
95
96        Ok(n)
97    }
98
99    fn len(&self) -> u64 {
100        self.data.len() as u64
101    }
102}
103
104/// Encrypted PFS image backed by a byte slice.
105pub(crate) struct EncryptedSlice<'a> {
106    data: &'a [u8],
107    decryptor: Xts128<Aes128>,
108    /// XTS block index where encryption begins.
109    encrypted_start: usize,
110}
111
112impl<'a> EncryptedSlice<'a> {
113    pub fn new(data: &'a [u8], decryptor: Xts128<Aes128>, encrypted_start: usize) -> Self {
114        Self {
115            data,
116            decryptor,
117            encrypted_start,
118        }
119    }
120}
121
122impl Image for EncryptedSlice<'_> {
123    fn read_at(&self, offset: u64, output_buf: &mut [u8]) -> io::Result<usize> {
124        let len = self.data.len() as u64;
125
126        if output_buf.is_empty() || offset >= len {
127            return Ok(0);
128        }
129
130        let mut copied = 0;
131        let mut pos = offset;
132        let mut scratch = vec![0u8; XTS_BLOCK_SIZE];
133
134        while copied < output_buf.len() && pos < len {
135            let block = (pos as usize) / XTS_BLOCK_SIZE;
136            let offset_in_block = (pos as usize) % XTS_BLOCK_SIZE;
137            let block_start = block * XTS_BLOCK_SIZE;
138
139            // Copy XTS block from backing slice into scratch buffer.
140            let src = self
141                .data
142                .get(block_start..block_start + XTS_BLOCK_SIZE)
143                .ok_or_else(|| io::Error::other(format!("XTS block #{} out of bounds", block)))?;
144
145            scratch.copy_from_slice(src);
146
147            // Decrypt if in encrypted region.
148            if block >= self.encrypted_start {
149                let tweak = get_tweak_default(block as _);
150                self.decryptor.decrypt_sector(&mut scratch, tweak);
151            }
152
153            // Copy the relevant portion to the output buffer.
154            let available = XTS_BLOCK_SIZE - offset_in_block;
155            let remaining_file = (len - pos) as usize;
156            let n = min(min(available, remaining_file), output_buf.len() - copied);
157
158            output_buf[copied..copied + n]
159                .copy_from_slice(&scratch[offset_in_block..offset_in_block + n]);
160
161            copied += n;
162            pos += n as u64;
163        }
164
165        Ok(copied)
166    }
167
168    fn len(&self) -> u64 {
169        self.data.len() as u64
170    }
171}