asahi_bless/
lib.rs

1// SPDX-License-Identifier: MIT
2#![allow(dead_code)]
3use apple_nvram::{nvram_parse, VarType};
4use gpt::{disk::LogicalBlockSize, GptConfig};
5use std::{
6    borrow::Cow,
7    collections::HashMap,
8    fs::{File, OpenOptions},
9    io::{self, Read, Seek, SeekFrom},
10    ops::Deref,
11};
12use uuid::Uuid;
13
14struct NxSuperblock([u8; NxSuperblock::SIZE]);
15
16impl NxSuperblock {
17    const SIZE: usize = 1408;
18    const MAGIC: u32 = 1112758350; //'BSXN'
19    const MAX_FILE_SYSTEMS: usize = 100;
20    fn get_buf(&mut self) -> &mut [u8] {
21        &mut self.0
22    }
23    fn new() -> Self {
24        NxSuperblock([0; NxSuperblock::SIZE])
25    }
26    fn magic(&self) -> u32 {
27        u32::from_le_bytes(self.0[32..32 + 4].try_into().unwrap())
28    }
29    fn block_size(&self) -> u32 {
30        u32::from_le_bytes(self.0[36..36 + 4].try_into().unwrap())
31    }
32    fn xid(&self) -> u64 {
33        u64::from_le_bytes(self.0[16..16 + 8].try_into().unwrap())
34    }
35    fn omap_oid(&self) -> u64 {
36        u64::from_le_bytes(self.0[160..160 + 8].try_into().unwrap())
37    }
38    fn xp_desc_blocks(&self) -> u32 {
39        u32::from_le_bytes(self.0[104..104 + 4].try_into().unwrap())
40    }
41    fn xp_desc_base(&self) -> u64 {
42        u64::from_le_bytes(self.0[112..112 + 8].try_into().unwrap())
43    }
44    fn fs_oid(&self, i: usize) -> u64 {
45        let at = 184 + 8 * i;
46        u64::from_le_bytes(self.0[at..at + 8].try_into().unwrap())
47    }
48}
49
50struct OmapPhys<'a>(&'a [u8]);
51impl OmapPhys<'_> {
52    const SIZE: usize = 88;
53    fn tree_oid(&self) -> u64 {
54        u64::from_le_bytes(self.0[48..48 + 8].try_into().unwrap())
55    }
56}
57
58struct NLoc<'a>(&'a [u8]);
59
60impl NLoc<'_> {
61    fn off(&self) -> u16 {
62        u16::from_le_bytes(self.0[0..2].try_into().unwrap())
63    }
64    fn len(&self) -> u16 {
65        u16::from_le_bytes(self.0[2..2 + 2].try_into().unwrap())
66    }
67}
68
69struct KVOff<'a>(&'a [u8]);
70impl KVOff<'_> {
71    const SIZE: usize = 4;
72    fn k(&self) -> u16 {
73        u16::from_le_bytes(self.0[0..2].try_into().unwrap())
74    }
75    fn v(&self) -> u16 {
76        u16::from_le_bytes(self.0[2..2 + 2].try_into().unwrap())
77    }
78}
79
80struct OmapKey<'a>(&'a [u8]);
81impl OmapKey<'_> {
82    fn oid(&self) -> u64 {
83        u64::from_le_bytes(self.0[0..8].try_into().unwrap())
84    }
85    fn xid(&self) -> u64 {
86        u64::from_le_bytes(self.0[8..8 + 8].try_into().unwrap())
87    }
88}
89
90struct OmapVal<'a>(&'a [u8]);
91impl OmapVal<'_> {
92    fn flags(&self) -> u32 {
93        u32::from_le_bytes(self.0[0..4].try_into().unwrap())
94    }
95    fn size(&self) -> u32 {
96        u32::from_le_bytes(self.0[4..4 + 4].try_into().unwrap())
97    }
98    fn paddr(&self) -> u64 {
99        u64::from_le_bytes(self.0[8..8 + 8].try_into().unwrap())
100    }
101}
102
103struct BTreeInfo;
104impl BTreeInfo {
105    const SIZE: usize = 40;
106}
107
108struct BTreeNodePhys<'a>(&'a [u8]);
109impl BTreeNodePhys<'_> {
110    const FIXED_KV_SIZE: u16 = 0x4;
111    const ROOT: u16 = 0x1;
112    const SIZE: usize = 56;
113    fn flags(&self) -> u16 {
114        u16::from_le_bytes(self.0[32..32 + 2].try_into().unwrap())
115    }
116    fn level(&self) -> u16 {
117        u16::from_le_bytes(self.0[34..34 + 2].try_into().unwrap())
118    }
119    fn table_space(&self) -> NLoc<'_> {
120        NLoc(&self.0[40..])
121    }
122    fn nkeys(&self) -> u32 {
123        u32::from_le_bytes(self.0[36..36 + 4].try_into().unwrap())
124    }
125}
126
127struct ApfsSuperblock<'a>(&'a [u8]);
128impl ApfsSuperblock<'_> {
129    fn volname(&self) -> &[u8] {
130        &self.0[704..704 + 128]
131    }
132    fn vol_uuid(&self) -> Uuid {
133        Uuid::from_slice(&self.0[240..240 + 16]).unwrap()
134    }
135    fn volume_group_id(&self) -> Uuid {
136        Uuid::from_slice(&self.0[1008..1008 + 16]).unwrap()
137    }
138    fn role(&self) -> u16 {
139        u16::from_le_bytes(self.0[964..964 + 2].try_into().unwrap())
140    }
141}
142
143const VOL_ROLE_SYSTEM: u16 = 1;
144
145fn pread<T: Read + Seek>(file: &mut T, pos: u64, target: &mut [u8]) -> io::Result<()> {
146    file.seek(SeekFrom::Start(pos))?;
147    file.read_exact(target)
148}
149
150// should probably fix xids here
151fn lookup(_disk: &mut File, cur_node: &BTreeNodePhys, key: u64) -> Option<u64> {
152    if cur_node.level() != 0 {
153        unimplemented!();
154    }
155    if cur_node.flags() & BTreeNodePhys::FIXED_KV_SIZE != 0 {
156        let toc_off = cur_node.table_space().off() as usize + BTreeNodePhys::SIZE;
157        let key_start = toc_off + cur_node.table_space().len() as usize;
158        let val_end = cur_node.0.len()
159            - if cur_node.flags() & BTreeNodePhys::ROOT == 0 {
160                0
161            } else {
162                BTreeInfo::SIZE
163            };
164        for i in 0..cur_node.nkeys() as usize {
165            let entry = KVOff(&cur_node.0[(toc_off + i * KVOff::SIZE)..]);
166            let key_off = entry.k() as usize + key_start;
167            let map_key = OmapKey(&cur_node.0[key_off..]);
168            if map_key.oid() == key {
169                let val_off = val_end - entry.v() as usize;
170                let val = OmapVal(&cur_node.0[val_off..]);
171                return Some(val.paddr());
172            }
173        }
174        None
175    } else {
176        unimplemented!();
177    }
178}
179
180fn trim_zeroes(s: &[u8]) -> &[u8] {
181    for i in 0..s.len() {
182        if s[i] == 0 {
183            return &s[..i];
184        }
185    }
186    s
187}
188
189fn scan_volume(disk: &mut File) -> io::Result<HashMap<Uuid, Vec<Volume>>> {
190    let mut sb = NxSuperblock::new();
191    disk.read_exact(sb.get_buf())?;
192    if sb.magic() != NxSuperblock::MAGIC {
193        return Ok(HashMap::new());
194    }
195    let block_size = sb.block_size() as u64;
196    for i in 0..sb.xp_desc_blocks() {
197        let mut sbc = NxSuperblock::new();
198        pread(
199            disk,
200            (sb.xp_desc_base() + i as u64) * block_size,
201            sbc.get_buf(),
202        )?;
203        if sbc.magic() == NxSuperblock::MAGIC {
204            if sbc.xid() > sb.xid() {
205                sb = sbc;
206            }
207        }
208    }
209    let mut omap_bytes = vec![0; OmapPhys::SIZE];
210    pread(disk, sb.omap_oid() * block_size, &mut omap_bytes)?;
211    let omap = OmapPhys(&omap_bytes);
212    let mut node_bytes = vec![0; sb.block_size() as usize];
213    pread(disk, omap.tree_oid() * block_size, &mut node_bytes)?;
214    let node = BTreeNodePhys(&node_bytes);
215    let mut vgs_found = HashMap::<Uuid, Vec<Volume>>::new();
216    for i in 0..NxSuperblock::MAX_FILE_SYSTEMS {
217        let fs_id = sb.fs_oid(i);
218        if fs_id == 0 {
219            continue;
220        }
221        let vsb = lookup(disk, &node, fs_id);
222        let mut asb_bytes = vec![0; sb.block_size() as usize];
223        if vsb.is_none() {
224            continue;
225        }
226        pread(disk, vsb.unwrap() * sb.block_size() as u64, &mut asb_bytes)?;
227        let asb = ApfsSuperblock(&asb_bytes);
228        if asb.volume_group_id().is_nil() {
229            continue;
230        }
231        if let Ok(name) = std::str::from_utf8(trim_zeroes(asb.volname())) {
232            vgs_found
233                .entry(asb.volume_group_id())
234                .or_default()
235                .push(Volume {
236                    name: name.to_owned(),
237                    is_system: asb.role() == VOL_ROLE_SYSTEM,
238                });
239        }
240    }
241    Ok(vgs_found)
242}
243
244#[derive(Debug)]
245pub struct Volume {
246    pub name: String,
247    pub is_system: bool,
248}
249
250#[derive(Debug)]
251pub struct BootCandidate {
252    pub part_uuid: Uuid,
253    pub vg_uuid: Uuid,
254    pub volumes: Vec<Volume>,
255}
256
257fn swap_uuid(u: &Uuid) -> Uuid {
258    let (a, b, c, d) = u.as_fields();
259    Uuid::from_fields(a.swap_bytes(), b.swap_bytes(), c.swap_bytes(), d)
260}
261
262#[derive(Debug)]
263pub enum Error {
264    Parse,
265    SectionTooBig,
266    ApplyError(std::io::Error),
267    OutOfRange,
268    Ambiguous,
269    NvramReadError(std::io::Error),
270    DiskReadError(std::io::Error),
271    VolumeNotFound,
272}
273
274impl From<apple_nvram::Error> for Error {
275    fn from(e: apple_nvram::Error) -> Self {
276        match e {
277            apple_nvram::Error::ParseError => Error::Parse,
278            apple_nvram::Error::SectionTooBig => Error::SectionTooBig,
279            apple_nvram::Error::ApplyError(e) => Error::ApplyError(e),
280        }
281    }
282}
283
284type Result<T> = std::result::Result<T, Error>;
285
286pub fn get_boot_candidates() -> Result<Vec<BootCandidate>> {
287    let disk = GptConfig::new()
288        .writable(false)
289        .logical_block_size(LogicalBlockSize::Lb4096)
290        .open("/dev/nvme0n1")
291        .map_err(Error::DiskReadError)?;
292    let mut cands = Vec::new();
293    for (i, v) in disk.partitions() {
294        if v.part_type_guid.guid != "7C3457EF-0000-11AA-AA11-00306543ECAC" {
295            continue;
296        }
297        let mut part = File::open(format!("/dev/nvme0n1p{i}")).map_err(Error::DiskReadError)?;
298        for (vg_uuid, volumes) in scan_volume(&mut part).unwrap_or_default() {
299            cands.push(BootCandidate {
300                vg_uuid,
301                volumes,
302                part_uuid: swap_uuid(&v.part_guid),
303            });
304        }
305    }
306
307    return Ok(cands);
308}
309
310const ALT_BOOT_VAR: &'static [u8] = b"alt-boot-volume";
311
312pub fn get_boot_volume(device: &str, next: bool) -> Result<BootCandidate> {
313    let mut file = OpenOptions::new()
314        .read(true)
315        .write(true)
316        .open(device)
317        .map_err(Error::NvramReadError)?;
318    let mut data = Vec::new();
319    file.read_to_end(&mut data).map_err(Error::NvramReadError)?;
320    let mut nv = nvram_parse(&data)?;
321
322    let active = nv.active_part_mut();
323    let v;
324    if next {
325        v = active
326            .get_variable(ALT_BOOT_VAR, VarType::System)
327            .or(active.get_variable(b"boot-volume", VarType::System))
328            .ok_or(Error::Parse);
329    } else {
330        v = active
331            .get_variable(b"boot-volume", VarType::System)
332            .ok_or(Error::Parse);
333    }
334    let data = String::from_utf8(v?.value().deref().to_vec()).unwrap();
335    let [_, part_uuid, part_vg_uuid]: [&str; 3] =
336        data.split(":").collect::<Vec<&str>>().try_into().unwrap();
337
338    Ok(BootCandidate {
339        volumes: Vec::new(),
340        part_uuid: Uuid::parse_str(part_uuid).unwrap(),
341        vg_uuid: Uuid::parse_str(part_vg_uuid).unwrap(),
342    })
343}
344
345pub fn clear_next_boot(device: &str) -> Result<bool> {
346    let mut file = OpenOptions::new()
347        .read(true)
348        .write(true)
349        .open(device)
350        .map_err(Error::ApplyError)?;
351    let mut data = Vec::new();
352    file.read_to_end(&mut data).map_err(Error::ApplyError)?;
353    let mut nv = nvram_parse(&data)?;
354    nv.prepare_for_write();
355    if nv.active_part_mut().get_variable(ALT_BOOT_VAR, VarType::System).is_none() {
356        return Ok(false);
357    }
358    nv.active_part_mut().remove_variable(
359        ALT_BOOT_VAR,
360        VarType::System,
361    );
362    nv.apply(&mut file)?;
363    Ok(true)
364}
365
366pub fn set_boot_volume(device: &str, cand: &BootCandidate, next: bool) -> Result<()> {
367    let mut nvram_key: &[u8] = b"boot-volume".as_ref();
368    if next {
369        nvram_key = ALT_BOOT_VAR.as_ref();
370    }
371
372    let boot_str = format!(
373        "EF57347C-0000-AA11-AA11-00306543ECAC:{}:{}",
374        cand.part_uuid
375            .hyphenated()
376            .encode_upper(&mut Uuid::encode_buffer()),
377        cand.vg_uuid
378            .hyphenated()
379            .encode_upper(&mut Uuid::encode_buffer())
380    );
381    let mut file = OpenOptions::new()
382        .read(true)
383        .write(true)
384        .open(device)
385        .map_err(Error::ApplyError)?;
386    let mut data = Vec::new();
387    file.read_to_end(&mut data).map_err(Error::ApplyError)?;
388    let mut nv = nvram_parse(&data)?;
389    nv.prepare_for_write();
390    nv.active_part_mut().insert_variable(
391        nvram_key,
392        Cow::Owned(boot_str.into_bytes()),
393        VarType::System,
394    );
395    nv.apply(&mut file)?;
396    Ok(())
397}
asahi_bless/lib.rs

asahi_bless/
lib.rs
