1use crate::{
12 raw,
13 util::{read_le_u16, read_le_u64, read_uuid},
14};
15use std::{collections::BTreeMap, mem};
16use uuid::Uuid;
17
18#[derive(Debug, Clone)]
20pub struct Stripe {
21 pub devid: u64,
22 pub offset: u64,
23 pub dev_uuid: Uuid,
24}
25
26#[derive(Debug, Clone)]
28pub struct ChunkMapping {
29 pub logical: u64,
30 pub length: u64,
31 pub stripe_len: u64,
32 pub chunk_type: u64,
33 pub num_stripes: u16,
34 pub sub_stripes: u16,
35 pub stripes: Vec<Stripe>,
36}
37
38#[derive(Debug, Default)]
42pub struct ChunkTreeCache {
43 inner: BTreeMap<u64, ChunkMapping>,
44}
45
46impl ChunkTreeCache {
47 pub fn insert(&mut self, mapping: ChunkMapping) {
49 self.inner.insert(mapping.logical, mapping);
50 }
51
52 pub fn lookup(&self, logical: u64) -> Option<&ChunkMapping> {
54 self.inner
56 .range(..=logical)
57 .next_back()
58 .map(|(_, mapping)| mapping)
59 .filter(|mapping| logical < mapping.logical + mapping.length)
60 }
61
62 pub fn resolve(&self, logical: u64) -> Option<u64> {
68 let mapping = self.lookup(logical)?;
69 let offset_within_chunk = logical - mapping.logical;
70 Some(mapping.stripes[0].offset + offset_within_chunk)
71 }
72
73 pub fn len(&self) -> usize {
75 self.inner.len()
76 }
77
78 pub fn is_empty(&self) -> bool {
80 self.inner.is_empty()
81 }
82}
83
84pub fn parse_chunk_item(
89 buf: &[u8],
90 logical: u64,
91) -> Option<(ChunkMapping, usize)> {
92 let chunk_base_size = mem::offset_of!(raw::btrfs_chunk, stripe);
93 let stripe_size = mem::size_of::<raw::btrfs_stripe>();
94
95 if buf.len() < chunk_base_size {
96 return None;
97 }
98
99 let length = read_le_u64(buf, 0);
100 let stripe_len = read_le_u64(buf, 16);
101 let chunk_type = read_le_u64(buf, 24);
102 let num_stripes = read_le_u16(buf, 44);
103 let sub_stripes = read_le_u16(buf, 46);
104
105 let total_size = chunk_base_size + num_stripes as usize * stripe_size;
106 if buf.len() < total_size {
107 return None;
108 }
109
110 let mut stripes = Vec::with_capacity(num_stripes as usize);
111 for i in 0..num_stripes as usize {
112 let s_off = chunk_base_size + i * stripe_size;
113 stripes.push(Stripe {
114 devid: read_le_u64(buf, s_off),
115 offset: read_le_u64(buf, s_off + 8),
116 dev_uuid: read_uuid(buf, s_off + 16),
117 });
118 }
119
120 let mapping = ChunkMapping {
121 logical,
122 length,
123 stripe_len,
124 chunk_type,
125 num_stripes,
126 sub_stripes,
127 stripes,
128 };
129
130 Some((mapping, total_size))
131}
132
133pub fn seed_from_sys_chunk_array(array: &[u8], size: u32) -> ChunkTreeCache {
138 let array = &array[..size as usize];
139 let mut cache = ChunkTreeCache::default();
140
141 let disk_key_size = mem::size_of::<raw::btrfs_disk_key>();
142 let mut offset = 0usize;
143
144 while offset + disk_key_size <= array.len() {
145 let key_offset = read_le_u64(array, offset + 9);
146 offset += disk_key_size;
147
148 if let Some((mapping, consumed)) =
149 parse_chunk_item(&array[offset..], key_offset)
150 {
151 cache.insert(mapping);
152 offset += consumed;
153 } else {
154 break;
155 }
156 }
157
158 cache
159}
160
161#[cfg(test)]
162mod tests {
163 use super::*;
164
165 fn make_mapping(logical: u64, length: u64, physical: u64) -> ChunkMapping {
166 ChunkMapping {
167 logical,
168 length,
169 stripe_len: 65536,
170 chunk_type: 0,
171 num_stripes: 1,
172 sub_stripes: 0,
173 stripes: vec![Stripe {
174 devid: 1,
175 offset: physical,
176 dev_uuid: Uuid::nil(),
177 }],
178 }
179 }
180
181 #[test]
182 fn empty_cache() {
183 let cache = ChunkTreeCache::default();
184 assert!(cache.is_empty());
185 assert_eq!(cache.resolve(0), None);
186 }
187
188 #[test]
189 fn single_mapping() {
190 let mut cache = ChunkTreeCache::default();
191 cache.insert(make_mapping(1000, 500, 2000));
192 assert_eq!(cache.len(), 1);
193
194 assert_eq!(cache.resolve(1000), Some(2000));
195 assert_eq!(cache.resolve(1100), Some(2100));
196 assert_eq!(cache.resolve(1499), Some(2499));
197 assert_eq!(cache.resolve(1500), None); assert_eq!(cache.resolve(999), None); }
200
201 #[test]
202 fn multiple_mappings() {
203 let mut cache = ChunkTreeCache::default();
204 cache.insert(make_mapping(0, 1000, 5000));
205 cache.insert(make_mapping(1000, 1000, 6000));
206 cache.insert(make_mapping(5000, 2000, 10000));
207
208 assert_eq!(cache.resolve(0), Some(5000));
209 assert_eq!(cache.resolve(500), Some(5500));
210 assert_eq!(cache.resolve(1000), Some(6000));
211 assert_eq!(cache.resolve(1999), Some(6999));
212 assert_eq!(cache.resolve(2000), None); assert_eq!(cache.resolve(5000), Some(10000));
214 assert_eq!(cache.resolve(6999), Some(11999));
215 assert_eq!(cache.resolve(7000), None);
216 }
217
218 #[test]
219 fn lookup_returns_mapping() {
220 let mut cache = ChunkTreeCache::default();
221 cache.insert(make_mapping(1000, 500, 2000));
222
223 let m = cache.lookup(1100).unwrap();
224 assert_eq!(m.logical, 1000);
225 assert_eq!(m.length, 500);
226 assert!(cache.lookup(500).is_none());
227 }
228
229 #[test]
230 fn seed_from_empty_array() {
231 let array = [0u8; 2048];
232 let cache = seed_from_sys_chunk_array(&array, 0);
233 assert!(cache.is_empty());
234 }
235
236 #[test]
237 fn parse_chunk_item_basic() {
238 let chunk_base_size = mem::offset_of!(raw::btrfs_chunk, stripe);
239 let stripe_size = mem::size_of::<raw::btrfs_stripe>();
240 let total = chunk_base_size + stripe_size;
241 let mut buf = vec![0u8; total];
242
243 buf[0..8].copy_from_slice(&1000u64.to_le_bytes());
245 buf[8..16].copy_from_slice(&2u64.to_le_bytes());
247 buf[16..24].copy_from_slice(&65536u64.to_le_bytes());
249 buf[24..32].copy_from_slice(&1u64.to_le_bytes());
251 buf[44..46].copy_from_slice(&1u16.to_le_bytes());
253
254 buf[chunk_base_size..chunk_base_size + 8]
256 .copy_from_slice(&1u64.to_le_bytes());
257 buf[chunk_base_size + 8..chunk_base_size + 16]
258 .copy_from_slice(&5000u64.to_le_bytes());
259
260 let (mapping, consumed) = parse_chunk_item(&buf, 0).unwrap();
261 assert_eq!(consumed, total);
262 assert_eq!(mapping.logical, 0);
263 assert_eq!(mapping.length, 1000);
264 assert_eq!(mapping.num_stripes, 1);
265 assert_eq!(mapping.stripes[0].devid, 1);
266 assert_eq!(mapping.stripes[0].offset, 5000);
267 }
268
269 #[test]
270 fn parse_chunk_item_too_short() {
271 let buf = [0u8; 10];
272 assert!(parse_chunk_item(&buf, 0).is_none());
273 }
274}