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
10pub trait Image: Send + Sync {
16 fn read_at(&self, offset: u64, output_buf: &mut [u8]) -> io::Result<usize>;
21
22 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 fn len(&self) -> u64;
47
48 fn is_empty(&self) -> bool {
50 self.len() == 0
51 }
52}
53
54pub(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
70pub(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
104pub(crate) struct EncryptedSlice<'a> {
106 data: &'a [u8],
107 decryptor: Xts128<Aes128>,
108 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 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 if block >= self.encrypted_start {
149 let tweak = get_tweak_default(block as _);
150 self.decryptor.decrypt_sector(&mut scratch, tweak);
151 }
152
153 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}