ubifs

src/formats/ubi/mod.rs

@@ -0,0 +1,366 @@
+//! UBI / UBIFS NAND rootfs extractor.
+//!
+//! Handles raw UBI images as commonly carved out of MStar/Toshiba NAND dumps
+//! (e.g. `rootfs_ubi.bin` produced by the `mstar_unfd` format). These images
+//! frequently still carry the NAND spare/OOB bytes interleaved after every
+//! page, and use a vendor-quirked UBI layout (EC header `data_offset` field
+//! reads 0). This module:
+//!
+//!   1. Detects the `UBI#` erase-counter header at offset 0.
+//!   2. Auto-detects and strips interleaved OOB (page + spare geometry).
+//!   3. Parses per-PEB EC/VID headers (big-endian) and rebuilds each logical
+//!      volume by ordering LEBs and keeping the copy with the highest sqnum.
+//!   4. Reads the UBI volume table (layout volume) to recover volume names.
+//!   5. Hands each data volume's reconstructed image to the UBIFS walker
+//!      ([`ubifs`]) which rebuilds the file tree.
+
+use std::any::Any;
+use std::collections::HashMap;
+use std::fs::{self, File};
+use std::io::{Read, Seek, SeekFrom};
+
+use log::{info, warn};
+
+use crate::AppContext;
+
+pub mod ubifs;
+
+/// UBI erase-counter header magic ("UBI#").
+const EC_MAGIC: &[u8; 4] = b"UBI#";
+/// UBI volume-identifier header magic ("UBI!").
+const VID_MAGIC: &[u8; 4] = b"UBI!";
+/// vol_id of the internal layout volume that holds the volume table.
+const UBI_LAYOUT_VOLUME_ID: u32 = 0x7FFF_EFFF;
+/// Size of a single volume-table record.
+const UBI_VTBL_RECORD_SIZE: usize = 172;
+/// NAND main page size we assume for OOB detection.
+const NAND_PAGE_SIZE: usize = 2048;
+
+/// Geometry / layout information for a UBI image.
+#[derive(Debug, Clone)]
+pub struct UbiCtx {
+    /// Physical erase block size as stored in the file (may include OOB).
+    pub phys_peb_size: usize,
+    /// OOB/spare bytes interleaved after each `NAND_PAGE_SIZE` page (0 = none).
+    pub oob_size: usize,
+    /// Clean PEB size after OOB has been removed.
+    pub clean_peb_size: usize,
+    /// Number of physical erase blocks in the image.
+    pub peb_count: usize,
+}
+
+/// Detect a UBI image: `UBI#` at offset 0, and derive PEB geometry (including
+/// any interleaved NAND OOB) from the spacing of consecutive EC headers.
+pub fn is_ubi_file(
+    app_ctx: &AppContext,
+) -> Result<Option<Box<dyn Any>>, Box<dyn std::error::Error>> {
+    let file = match app_ctx.file() {
+        Some(f) => f,
+        None => return Ok(None),
+    };
+
+    let mut head = [0u8; 4];
+    {
+        let mut f = file;
+        f.seek(SeekFrom::Start(0))?;
+        if f.read_exact(&mut head).is_err() {
+            return Ok(None);
+        }
+    }
+    if &head != EC_MAGIC {
+        return Ok(None);
+    }
+
+    let file_size = file.metadata()?.len() as usize;
+
+    // Find the spacing between the first two EC headers to learn the physical
+    // PEB size (OOB-interleaved or not). Scan a bounded prefix so detection
+    // stays cheap.
+    let scan_len = file_size.min(4 * 1024 * 1024);
+    let prefix = crate::utils::common::read_file(&file, 0, scan_len)?;
+    let phys_peb_size = match second_magic_offset(&prefix, EC_MAGIC) {
+        Some(d) if d > 0 => d,
+        // Only one EC header visible in the prefix — fall back to a common
+        // clean PEB size guess.
+        _ => guess_single_peb(file_size),
+    };
+
+    let oob_size = detect_oob(phys_peb_size);
+    let pages = phys_peb_size / (NAND_PAGE_SIZE + oob_size);
+    let clean_peb_size = pages * NAND_PAGE_SIZE;
+    let peb_count = file_size / phys_peb_size;
+
+    if clean_peb_size == 0 || peb_count == 0 {
+        return Ok(None);
+    }
+
+    info!("- Detected UBI image");
+    info!(
+        "  physical PEB: {} bytes, OOB/page: {} bytes, clean PEB: {} bytes, PEBs: {}",
+        phys_peb_size, oob_size, clean_peb_size, peb_count
+    );
+
+    Ok(Some(Box::new(UbiCtx {
+        phys_peb_size,
+        oob_size,
+        clean_peb_size,
+        peb_count,
+    })))
+}
+
+/// A reconstructed logical volume: its LEBs concatenated in `lnum` order.
+struct Volume {
+    /// lnum -> (sqnum, leb data)
+    lebs: HashMap<u32, (u64, Vec<u8>)>,
+}
+
+pub fn extract_ubi(
+    app_ctx: &AppContext,
+    ctx: Box<dyn Any>,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let ctx = ctx.downcast::<UbiCtx>().map_err(|_| "Invalid UBI context")?;
+
+    let input_path = app_ctx
+        .input_path()
+        .ok_or("UBI extractor requires an input path")?;
+    let mut file = File::open(input_path)?;
+
+    fs::create_dir_all(&app_ctx.output_dir)?;
+
+    // --- Parse every physical erase block ---
+    let mut volumes: HashMap<u32, Volume> = HashMap::new();
+    let mut buf = vec![0u8; ctx.phys_peb_size];
+
+    for peb in 0..ctx.peb_count {
+        file.seek(SeekFrom::Start((peb * ctx.phys_peb_size) as u64))?;
+        if file.read_exact(&mut buf).is_err() {
+            break;
+        }
+
+        let clean = strip_oob(&buf, ctx.oob_size, ctx.clean_peb_size);
+
+        if &clean[0..4] != EC_MAGIC {
+            continue; // erased / non-UBI block
+        }
+
+        // EC header (big-endian). vid_hdr_offset @0x10, data_offset @0x14.
+        let mut vid_hdr_offset = be32(&clean, 0x10) as usize;
+        let mut data_offset = be32(&clean, 0x14) as usize;
+
+        // Vendor quirk: fields may read 0. Fall back to page-aligned defaults.
+        if vid_hdr_offset == 0 || vid_hdr_offset + 64 > ctx.clean_peb_size {
+            vid_hdr_offset = NAND_PAGE_SIZE;
+        }
+        if data_offset == 0 || data_offset >= ctx.clean_peb_size {
+            data_offset = vid_hdr_offset + NAND_PAGE_SIZE;
+        }
+
+        if vid_hdr_offset + 64 > clean.len() || &clean[vid_hdr_offset..vid_hdr_offset + 4] != VID_MAGIC
+        {
+            continue; // no valid VID header -> unmapped PEB
+        }
+
+        // VID header (big-endian). vol_id @0x08, lnum @0x0C, sqnum @0x28.
+        let vol_id = be32(&clean, vid_hdr_offset + 0x08);
+        let lnum = be32(&clean, vid_hdr_offset + 0x0C);
+        let sqnum = be64(&clean, vid_hdr_offset + 0x28);
+
+        if data_offset >= clean.len() {
+            continue;
+        }
+        let leb_data = clean[data_offset..].to_vec();
+
+        let vol = volumes.entry(vol_id).or_insert_with(|| Volume {
+            lebs: HashMap::new(),
+        });
+        // Keep the newest copy (highest sqnum) of each LEB.
+        match vol.lebs.get(&lnum) {
+            Some((prev_sq, _)) if *prev_sq >= sqnum => {}
+            _ => {
+                vol.lebs.insert(lnum, (sqnum, leb_data));
+            }
+        }
+    }
+
+    if volumes.is_empty() {
+        warn!("  No valid UBI volumes found");
+        return Ok(());
+    }
+
+    // --- Recover volume names from the layout volume's volume table ---
+    let names = volumes
+        .get(&UBI_LAYOUT_VOLUME_ID)
+        .map(|v| parse_volume_table(v))
+        .unwrap_or_default();
+
+    info!("  Found {} UBI volume(s)", volumes.len());
+
+    // --- Reconstruct and extract each data volume ---
+    let mut vol_ids: Vec<u32> = volumes.keys().cloned().collect();
+    vol_ids.sort_unstable();
+
+    for vol_id in vol_ids {
+        if vol_id == UBI_LAYOUT_VOLUME_ID {
+            continue; // internal layout volume, not a filesystem
+        }
+        let vol = &volumes[&vol_id];
+
+        let name = names
+            .get(&vol_id)
+            .cloned()
+            .unwrap_or_else(|| format!("vol_{vol_id}"));
+
+        let image = reconstruct_volume(vol);
+        info!(
+            "  Volume '{}' (id {}): {} LEBs, {} bytes",
+            name,
+            vol_id,
+            vol.lebs.len(),
+            image.len()
+        );
+
+        // Always dump the raw reconstructed volume image alongside the tree.
+        let img_path = app_ctx.output_dir.join(format!("{name}.ubifs"));
+        if let Err(e) = fs::write(&img_path, &image) {
+            warn!("    Could not write {}: {}", img_path.display(), e);
+        }
+
+        // Walk the UBIFS filesystem and rebuild files.
+        let out_dir = app_ctx.output_dir.join(&name);
+        match ubifs::extract_ubifs(&image, &out_dir) {
+            Ok(n) => info!("    Extracted {} file(s) from '{}'", n, name),
+            Err(e) => warn!("    UBIFS extraction of '{}' failed: {}", name, e),
+        }
+    }
+
+    Ok(())
+}
+
+/// Concatenate a volume's LEBs in ascending `lnum` order.
+fn reconstruct_volume(vol: &Volume) -> Vec<u8> {
+    let mut lnums: Vec<u32> = vol.lebs.keys().cloned().collect();
+    lnums.sort_unstable();
+    let mut out = Vec::new();
+    for lnum in lnums {
+        out.extend_from_slice(&vol.lebs[&lnum].1);
+    }
+    out
+}
+
+/// Parse the UBI volume table (stored in the layout volume) into a
+/// `vol_id -> name` map. Each record is `UBI_VTBL_RECORD_SIZE` bytes; the
+/// record index equals the volume id.
+fn parse_volume_table(layout: &Volume) -> HashMap<u32, String> {
+    let mut names = HashMap::new();
+    // The layout volume mirrors the table across its LEBs; LEB 0 is enough.
+    let Some((_, data)) = layout.lebs.get(&0).or_else(|| layout.lebs.values().next()) else {
+        return names;
+    };
+
+    let count = data.len() / UBI_VTBL_RECORD_SIZE;
+    for idx in 0..count {
+        let base = idx * UBI_VTBL_RECORD_SIZE;
+        let rec = &data[base..base + UBI_VTBL_RECORD_SIZE];
+        let reserved_pebs = be32(rec, 0x00);
+        if reserved_pebs == 0 {
+            continue; // unused record
+        }
+        let name_len = be16(rec, 0x0E) as usize;
+        if name_len == 0 || name_len > 128 {
+            continue;
+        }
+        let name = String::from_utf8_lossy(&rec[0x10..0x10 + name_len]).to_string();
+        if !name.is_empty() {
+            names.insert(idx as u32, sanitize_name(&name));
+        }
+    }
+    names
+}
+
+/// Remove path separators / control characters from a volume name so it is
+/// safe to use as a directory name.
+fn sanitize_name(name: &str) -> String {
+    name.chars()
+        .map(|c| {
+            if c.is_control() || c == '/' || c == '\\' || c == ':' {
+                '_'
+            } else {
+                c
+            }
+        })
+        .collect::<String>()
+        .trim()
+        .to_string()
+}
+
+/// Remove interleaved OOB from a physical PEB, returning the clean main data.
+fn strip_oob(peb: &[u8], oob_size: usize, clean_peb_size: usize) -> Vec<u8> {
+    if oob_size == 0 {
+        return peb[..clean_peb_size.min(peb.len())].to_vec();
+    }
+    let step = NAND_PAGE_SIZE + oob_size;
+    let mut out = Vec::with_capacity(clean_peb_size);
+    let mut pos = 0;
+    while pos + NAND_PAGE_SIZE <= peb.len() && out.len() < clean_peb_size {
+        out.extend_from_slice(&peb[pos..pos + NAND_PAGE_SIZE]);
+        pos += step;
+    }
+    out.truncate(clean_peb_size);
+    out
+}
+
+/// Offset of the second occurrence of `magic` (i.e. the distance from the
+/// first), searched only on page-aligned boundaries for speed.
+fn second_magic_offset(data: &[u8], magic: &[u8; 4]) -> Option<usize> {
+    let mut off = NAND_PAGE_SIZE;
+    while off + 4 <= data.len() {
+        if &data[off..off + 4] == magic {
+            return Some(off);
+        }
+        off += NAND_PAGE_SIZE;
+    }
+    None
+}
+
+/// Determine the OOB size per page given a physical PEB size. Prefers no OOB;
+/// otherwise picks the first spare size that divides the PEB into a power-of-two
+/// page count.
+fn detect_oob(phys_peb_size: usize) -> usize {
+    for oob in [0usize, 16, 32, 64, 128, 218, 224, 256] {
+        let step = NAND_PAGE_SIZE + oob;
+        if phys_peb_size % step != 0 {
+            continue;
+        }
+        let pages = phys_peb_size / step;
+        if pages.is_power_of_two() && (16..=4096).contains(&pages) {
+            return oob;
+        }
+    }
+    0
+}
+
+/// Fallback PEB size when only one EC header is present.
+fn guess_single_peb(file_size: usize) -> usize {
+    for peb in [131072usize, 262144, 126976, 524288, 65536] {
+        if file_size % peb == 0 {
+            return peb;
+        }
+    }
+    131072
+}
+
+// --- endian helpers ---------------------------------------------------------
+
+fn be32(b: &[u8], off: usize) -> u32 {
+    u32::from_be_bytes([b[off], b[off + 1], b[off + 2], b[off + 3]])
+}
+
+fn be16(b: &[u8], off: usize) -> u16 {
+    u16::from_be_bytes([b[off], b[off + 1]])
+}
+
+fn be64(b: &[u8], off: usize) -> u64 {
+    u64::from_be_bytes([
+        b[off], b[off + 1], b[off + 2], b[off + 3], b[off + 4], b[off + 5], b[off + 6], b[off + 7],
+    ])
+}