macos-catalog/src/restore.rs (a74924e) - apple-crates · Git

restore.rs

//! Apple-silicon Macs, their boards and names. Apple's restore catalog keys each
//! model identifier by its chip and board ID; the universal restore image's
//! `BuildManifest.plist` names the board and its device tree for each such pair,
//! and the device tree names the Mac.

use crate::Error;
use futures_util::{StreamExt, TryStreamExt};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::io::Read;

const RESTORE_CATALOG: &str =
    "https://mesu.apple.com/assets/macos/com_apple_macOSIPSW/com_apple_macOSIPSW.xml";

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct Model {
    /// For example `Mac16,8`.
    pub identifier: String,
    /// The board as release catalogs list it, for example `J614sAP`, in any case.
    pub board: String,
    /// What About This Mac calls it, in English, for example `MacBook Pro
    /// (14-inch, Nov 2024)`.
    #[serde(default)]
    pub name: String,
}

#[derive(Deserialize)]
struct RestoreCatalog {
    #[serde(rename = "MobileDeviceProductTypes")]
    product_types: ProductTypes,
    #[serde(rename = "MobileDeviceSoftwareVersionsByVersion")]
    versions: HashMap<String, Versions>,
}

#[derive(Deserialize)]
struct ProductTypes {
    // Keyed by (board << 24) | chip, in decimal.
    #[serde(rename = "DFU")]
    dfu: HashMap<String, String>,
}

#[derive(Deserialize)]
struct Versions {
    #[serde(rename = "MobileDeviceSoftwareVersions")]
    models: HashMap<String, HashMap<String, Build>>,
}

#[derive(Deserialize)]
struct Build {
    #[serde(rename = "Restore")]
    restore: Option<Restore>,
}

#[derive(Deserialize)]
struct Restore {
    #[serde(rename = "FirmwareURL")]
    url: String,
    #[serde(rename = "ProductVersion")]
    version: String,
}

#[derive(Deserialize)]
struct Manifest {
    #[serde(rename = "BuildIdentities")]
    identities: Vec<Identity>,
}

#[derive(Deserialize)]
struct Identity {
    #[serde(rename = "ApChipID")]
    chip: String,
    #[serde(rename = "ApBoardID")]
    board: String,
    #[serde(rename = "Info")]
    info: IdentityInfo,
    #[serde(rename = "Manifest")]
    manifest: Components,
}

#[derive(Deserialize)]
struct IdentityInfo {
    #[serde(rename = "DeviceClass")]
    device_class: String,
}

#[derive(Deserialize)]
struct Components {
    #[serde(rename = "DeviceTree")]
    device_tree: Component,
}

#[derive(Deserialize)]
struct Component {
    #[serde(rename = "Info")]
    info: ComponentInfo,
}

#[derive(Deserialize)]
struct ComponentInfo {
    #[serde(rename = "Path")]
    path: String,
}

/// Every Apple-silicon Mac the newest restore image supports, by identifier.
pub async fn apple_silicon_models(client: &reqwest::Client) -> Result<Vec<Model>, Error> {
    let catalog = client
        .get(RESTORE_CATALOG)
        .send()
        .await?
        .error_for_status()?
        .bytes()
        .await?;
    let catalog: RestoreCatalog = plist::from_bytes(&catalog).map_err(|_| Error::InvalidCatalog)?;
    // One universal image restores every Mac; the newest names the newest boards.
    let image = catalog
        .versions
        .values()
        .flat_map(|versions| versions.models.values())
        .flat_map(|builds| builds.values())
        .filter_map(|build| build.restore.as_ref())
        .max_by_key(|restore| crate::version_key(&restore.version))
        .ok_or(Error::InvalidCatalog)?;
    let image = RemoteZip::open(client, &image.url).await?;
    let manifest: Manifest = plist::from_bytes(&image.file("BuildManifest.plist").await?)
        .map_err(|_| Error::InvalidCatalog)?;

let hex = |text: &str| u32::from_str_radix(text.trim_start_matches("0x"), 16).ok();
    let boards: HashMap<(u32, u32), (String, String)> = manifest
        .identities
        .into_iter()
        .filter_map(|identity| {
            Some((
                (hex(&identity.chip)?, hex(&identity.board)?),
                (
                    identity.info.device_class,
                    identity.manifest.device_tree.info.path,
                ),
            ))
        })
        .collect();
    // Each identifier with its board and device tree.
    let mut macs: Vec<(String, &str, &str)> = catalog
        .product_types
        .dfu
        .into_iter()
        .filter(|(_, identifier)| !identifier.starts_with("VirtualMac"))
        .filter_map(|(key, identifier)| {
            let key: u32 = key.parse().ok()?;
            let (board, tree) = boards.get(&(key & 0xFFFF, key >> 24))?;
            Some((identifier, board.as_str(), tree.as_str()))
        })
        .collect();
    macs.sort_by(|a, b| a.0.cmp(&b.0));
    macs.dedup_by(|a, b| a.0 == b.0);

// About 60 boards, each tree well under 100 kB; a few at a time.
    // Owned paths keep the future `Send`.
    let trees: Vec<String> = macs.iter().map(|(_, _, tree)| tree.to_string()).collect();
    let image = &image;
    let names: Vec<String> =
        futures_util::stream::iter(trees)
            .map(|tree| async move {
                crate::device_tree::product_description(&image.file(&tree).await?)
            })
            .buffered(8)
            .try_collect()
            .await?;
    Ok(macs
        .into_iter()
        .zip(names)
        .map(|((identifier, board, _), name)| Model {
            identifier,
            board: board.into(),
            name,
        })
        .collect())
}

/// A remote ZIP64 archive read with range requests: restore images are many
/// gigabytes, the files read here well under a megabyte.
struct RemoteZip<'a> {
    client: &'a reqwest::Client,
    url: &'a str,
    directory: Vec<u8>,
}

impl<'a> RemoteZip<'a> {
    async fn open(client: &'a reqwest::Client, url: &'a str) -> Result<Self, Error> {
        // A HEAD response has no body, so read its length header rather than the body's.
        let size: u64 = client
            .head(url)
            .send()
            .await?
            .error_for_status()?
            .headers()
            .get(reqwest::header::CONTENT_LENGTH)
            .and_then(|length| length.to_str().ok()?.parse().ok())
            .ok_or(Error::InvalidCatalog)?;
        let mut zip = Self {
            client,
            url,
            directory: Vec::new(),
        };
        // The ZIP64 end-of-central-directory record sits within the last few bytes.
        let tail_length = size.min(64 * 1024);
        let tail = zip.range(size - tail_length, tail_length).await?;
        let record = find(&tail, b"PK\x06\x06").ok_or(Error::InvalidCatalog)?;
        let directory_length = u64_at(&tail, record + 40)?;
        let directory_start = u64_at(&tail, record + 48)?;
        zip.directory = zip.range(directory_start, directory_length).await?;
        Ok(zip)
    }

async fn file(&self, name: &str) -> Result<Vec<u8>, Error> {
        let entry = central_entry(&self.directory, name)?;
        let header = self.range(entry.offset, 30).await?;
        let data_start =
            entry.offset + 30 + u16_at(&header, 26)? as u64 + u16_at(&header, 28)? as u64;
        let compressed = self.range(data_start, entry.compressed).await?;
        match entry.method {
            0 => Ok(compressed),
            8 => {
                let mut data = Vec::new();
                flate2::read::DeflateDecoder::new(&compressed[..])
                    .read_to_end(&mut data)
                    .map_err(|_| Error::InvalidCatalog)?;
                Ok(data)
            }
            _ => Err(Error::InvalidCatalog),
        }
    }

async fn range(&self, start: u64, length: u64) -> Result<Vec<u8>, Error> {
        let bytes = self
            .client
            .get(self.url)
            .header(
                reqwest::header::RANGE,
                format!("bytes={start}-{}", start + length - 1),
            )
            .send()
            .await?
            .error_for_status()?
            .bytes()
            .await?;
        match bytes.len() as u64 == length {
            true => Ok(bytes.to_vec()),
            false => Err(Error::InvalidCatalog),
        }
    }
}

struct Entry {
    method: u16,
    compressed: u64,
    offset: u64,
}

fn central_entry(directory: &[u8], name: &str) -> Result<Entry, Error> {
    let mut at = 0;
    while directory.get(at..at + 4) == Some(b"PK\x01\x02") {
        let name_length = u16_at(directory, at + 28)? as usize;
        let extra_length = u16_at(directory, at + 30)? as usize;
        let comment_length = u16_at(directory, at + 32)? as usize;
        let names = at + 46..at + 46 + name_length;
        if directory.get(names.clone()) == Some(name.as_bytes()) {
            let mut compressed = u32_at(directory, at + 20)? as u64;
            let uncompressed = u32_at(directory, at + 24)?;
            let mut offset = u32_at(directory, at + 42)? as u64;
            // ZIP64 keeps values that do not fit in 32 bits in an extra field, in
            // this order, and only those.
            let mut extra = names.end;
            while extra + 4 <= names.end + extra_length {
                let (id, length) = (
                    u16_at(directory, extra)?,
                    u16_at(directory, extra + 2)? as usize,
                );
                if id == 1 {
                    let mut field = extra + 4;
                    let mut next = || {
                        let value = u64_at(directory, field);
                        field += 8;
                        value
                    };
                    if uncompressed == u32::MAX {
                        next()?;
                    }
                    if compressed == u32::MAX as u64 {
                        compressed = next()?;
                    }
                    if offset == u32::MAX as u64 {
                        offset = next()?;
                    }
                }
                extra += 4 + length;
            }
            return Ok(Entry {
                method: u16_at(directory, at + 10)?,
                compressed,
                offset,
            });
        }
        at = names.end + extra_length + comment_length;
    }
    Err(Error::InvalidCatalog)
}

fn find(bytes: &[u8], signature: &[u8]) -> Option<usize> {
    bytes
        .windows(signature.len())
        .rposition(|window| window == signature)
}

fn u16_at(bytes: &[u8], at: usize) -> Result<u16, Error> {
    Ok(u16::from_le_bytes(field(bytes, at)?))
}

fn u32_at(bytes: &[u8], at: usize) -> Result<u32, Error> {
    Ok(u32::from_le_bytes(field(bytes, at)?))
}

fn u64_at(bytes: &[u8], at: usize) -> Result<u64, Error> {
    Ok(u64::from_le_bytes(field(bytes, at)?))
}

fn field<const N: usize>(bytes: &[u8], at: usize) -> Result<[u8; N], Error> {
    bytes
        .get(at..at + N)
        .and_then(|bytes| bytes.try_into().ok())
        .ok_or(Error::InvalidCatalog)
}

#[cfg(test)]
mod tests {
    use super::*;

#[test]
    fn finds_a_zip64_entry_and_its_extra_fields() {
        let name = b"BuildManifest.plist";
        let mut entry = Vec::new();
        entry.extend(b"PK\x01\x02");
        entry.extend([0; 6]);
        entry.extend(8u16.to_le_bytes()); // method
        entry.extend([0; 8]);
        entry.extend(u32::MAX.to_le_bytes()); // compressed
        entry.extend(u32::MAX.to_le_bytes()); // uncompressed
        entry.extend((name.len() as u16).to_le_bytes());
        entry.extend(28u16.to_le_bytes()); // extra
        entry.extend(0u16.to_le_bytes()); // comment
        entry.extend([0; 8]);
        entry.extend(u32::MAX.to_le_bytes()); // offset
        entry.extend(name);
        entry.extend(1u16.to_le_bytes());
        entry.extend(24u16.to_le_bytes());
        entry.extend(5000u64.to_le_bytes()); // uncompressed
        entry.extend(1234u64.to_le_bytes()); // compressed
        entry.extend((7u64 << 32).to_le_bytes()); // offset
        let mut directory = Vec::new();
        directory.extend(&entry[..46]);
        directory.extend(b"other.bin");
        directory.extend(&entry[46 + name.len()..]);
        directory[28..30].copy_from_slice(&9u16.to_le_bytes());
        directory.extend(&entry);

let found = central_entry(&directory, "BuildManifest.plist").unwrap();
        assert_eq!(found.method, 8);
        assert_eq!(found.compressed, 1234);
        assert_eq!(found.offset, 7 << 32);
        assert!(central_entry(&directory, "missing").is_err());
    }

#[tokio::test]
    #[ignore = "contacts Apple's restore catalog and reads part of a restore image"]
    async fn maps_models_to_boards() {
        let models = apple_silicon_models(&reqwest::Client::new()).await.unwrap();
        let board = |identifier: &str| {
            models
                .iter()
                .find(|model| model.identifier == identifier)
                .map(|model| model.board.to_ascii_lowercase())
        };
        assert_eq!(board("MacBookPro17,1").as_deref(), Some("j293ap"));
        assert_eq!(board("Mac16,8").as_deref(), Some("j614sap"));
        assert!(models.len() > 50);
        let name = |identifier: &str| {
            let model = models.iter().find(|model| model.identifier == identifier);
            model.map(|model| model.name.as_str())
        };
        assert_eq!(name("Mac16,8"), Some("MacBook Pro (14-inch, Nov 2024)"));
        assert!(models.iter().all(|model| !model.name.is_empty()));
    }
}

apple-crates

Refactor and unify a lot of things for another project and implement airdrop support