diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index b0caeb0..b328b81 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -70,16 +70,6 @@ jobs: with: key: ${{ matrix.target }} - - name: Install build deps (macOS) - if: runner.os == 'macOS' - run: brew install protobuf - - - name: Install build deps (Linux) - if: runner.os == 'Linux' - run: | - sudo apt-get update - sudo apt-get install -y protobuf-compiler - - name: Install zig if: runner.os == 'Linux' && !contains(matrix.target, 'android') uses: mlugg/setup-zig@v2 @@ -103,41 +93,35 @@ jobs: ndk-version: r23 local-cache: true - - name: Install protoc (Windows) - if: runner.os == 'Windows' - uses: arduino/setup-protoc@v1 - with: - version: "3.x" - - name: Build binary (Linux) if: runner.os == 'Linux' && !contains(matrix.target, 'android') run: | - cargo zigbuild --release --locked --target=${{ matrix.target }} --color=always --verbose + cargo zigbuild --profile release-cli --locked --target=${{ matrix.target }} --color=always --verbose --bin otadump - name: Build binary (Android) if: contains(matrix.target, 'android') run: | - cargo ndk --target=${{ matrix.target }} build --release --locked --color=always --verbose + cargo ndk --target=${{ matrix.target }} build --profile release-cli --locked --color=always --verbose --bin otadump - name: Build binary (macOS / Windows) if: runner.os == 'macOS' || runner.os == 'Windows' run: | - cargo build --release --locked --target=${{ matrix.target }} --color=always --verbose + cargo build --profile release-cli --locked --target=${{ matrix.target }} --color=always --verbose --bin otadump - name: Package (*nix) if: runner.os != 'Windows' run: > tar -cv - LICENSE README.md - -C target/${{ matrix.target }}/release/ otadump + LICENSE README.md src/puffin/LICENSE licenses/LICENSE.lz4 licenses/README.md + -C target/${{ matrix.target }}/release-cli/ otadump | gzip --best > 'otadump-${{ steps.get_version.outputs.value }}-${{ matrix.target }}.tar.gz' - name: Package (Windows) if: runner.os == 'Windows' run: > 7z a 'otadump-${{ steps.get_version.outputs.value }}-${{ matrix.target }}.zip' - LICENSE README.md - ./target/${{ matrix.target }}/release/otadump.exe + LICENSE README.md src/puffin/LICENSE licenses/LICENSE.lz4 licenses/README.md + ./target/${{ matrix.target }}/release-cli/otadump.exe - name: Upload artifact uses: actions/upload-artifact@v4.6.0 with: @@ -157,3 +141,79 @@ jobs: *.zip name: ${{ steps.get_version.outputs.value }} tag_name: "" + + wheel-linux-abi3: + name: wheel manylinux2014 x86_64 + runs-on: ubuntu-latest + container: + image: quay.io/pypa/manylinux2014_x86_64@sha256:493d2032114d757aaa761a9385ad8497f391503bf71acef9abeeb66682ca5d90 + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Install Rust + uses: dtolnay/rust-toolchain@stable + + - name: Install Python build tools + run: | + /opt/python/cp39-cp39/bin/python -m pip install --upgrade pip + /opt/python/cp39-cp39/bin/python -m pip install uv maturin auditwheel + + - name: Validate, test, and build wheel in pinned manylinux + env: + PATH: /opt/python/cp39-cp39/bin:${{ env.PATH }} + OTADUMP_NATIVE_CACHE: /tmp/otadump-native-cache + run: | + set -euxo pipefail + + cargo fmt --all -- --check + cargo check --locked + cargo test --locked --release --all-features + cargo clippy --locked --all-targets --all-features -- -D warnings + + rm -rf "$OTADUMP_NATIVE_CACHE" + python3 scripts/fetch-native-artifact.py \ + --bundle-lock native/zucchini/ARTIFACT_BUNDLE_LOCK.json \ + --checksum-lock native/zucchini/ARTIFACT_LOCK.sha256 \ + --out /tmp/otadump-native/linux-x86_64-gnu \ + --cache "$OTADUMP_NATIVE_CACHE" + python3 tests/native_artifact_fetch_test.py + + uv build --wheel + auditwheel show dist/*.whl + auditwheel repair --wheel-dir wheelhouse dist/*.whl + + python -m venv /tmp/otadump-wheel-smoke + source /tmp/otadump-wheel-smoke/bin/activate + pip install wheelhouse/*.whl + python tests/python_smoke.py + python - <<'PY' + import pathlib + import otadump + + extension = next(pathlib.Path(otadump.__file__).parent.glob("otadump*.so")) + print(extension) + PY + + SO_PATH=$(python - <<'PY' + import pathlib + import otadump + + extension = next(pathlib.Path(otadump.__file__).parent.glob("otadump*.so")) + print(extension) + PY + ) + nm -gD "$SO_PATH" >/tmp/otadump-wheel.nm.txt + readelf -d "$SO_PATH" >/tmp/otadump-wheel.readelf.txt + ldd "$SO_PATH" >/tmp/otadump-wheel.ldd.txt + + - name: Upload wheel artifacts + uses: actions/upload-artifact@v4.6.0 + with: + name: python-wheel-manylinux2014-x86_64 + path: | + dist/*.whl + wheelhouse/*.whl + /tmp/otadump-wheel.nm.txt + /tmp/otadump-wheel.readelf.txt + /tmp/otadump-wheel.ldd.txt diff --git a/Cargo.lock b/Cargo.lock index 370ea0e..09d0807 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,6 +11,21 @@ dependencies = [ "memchr", ] +[[package]] +name = "alloc-no-stdlib" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" + +[[package]] +name = "alloc-stdlib" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195" +dependencies = [ + "alloc-no-stdlib", +] + [[package]] name = "anstream" version = "0.6.19" @@ -82,6 +97,36 @@ version = "2.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1b8e56985ec62d17e9c1001dc89c88ecd7dc08e47eba5ec7c29c7b5eeecde967" +[[package]] +name = "block2" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdeb9d870516001442e364c5220d3574d2da8dc765554b4a617230d33fa58ef5" +dependencies = [ + "objc2", +] + +[[package]] +name = "brotli" +version = "8.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] + +[[package]] +name = "brotli-decompressor" +version = "5.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + [[package]] name = "bumpalo" version = "3.18.1" @@ -96,9 +141,9 @@ checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" [[package]] name = "bzip2" -version = "0.6.0" +version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bea8dcd42434048e4f7a304411d9273a411f647446c1234a65ce0554923f4cff" +checksum = "f3a53fac24f34a81bc9954b5d6cfce0c21e18ec6959f44f56e8e90e4bb7c346c" dependencies = [ "libbz2-rs-sys", ] @@ -120,6 +165,12 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9555578bc9e57714c812a1f84e4fc5b4d21fcb063490c624de019f7464c91268" +[[package]] +name = "cfg_aliases" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" + [[package]] name = "clap" version = "4.5.40" @@ -234,6 +285,17 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "ctrlc" +version = "3.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0b1fab2ae45819af2d0731d60f2afe17227ebb1a1538a236da84c93e9a60162" +dependencies = [ + "dispatch2", + "nix", + "windows-sys 0.61.2", +] + [[package]] name = "derive_arbitrary" version = "1.4.1" @@ -245,6 +307,18 @@ dependencies = [ "syn 2.0.104", ] +[[package]] +name = "dispatch2" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38" +dependencies = [ + "bitflags", + "block2", + "libc", + "objc2", +] + [[package]] name = "either" version = "1.15.0" @@ -392,9 +466,9 @@ checksum = "775bf80d5878ab7c2b1080b5351a48b2f737d9f6f8b383574eebcc22be0dfccb" [[package]] name = "libc" -version = "0.2.174" +version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1171693293099992e19cddea4e8b849964e9846f4acee11b3948bcc337be8776" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" [[package]] name = "liblzma" @@ -428,6 +502,16 @@ version = "0.4.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94" +[[package]] +name = "lz4-sys" +version = "1.11.1+lz4-1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bd8c0d6c6ed0cd30b3652886bb8711dc4bb01d637a68105a3d5158039b418e6" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "memchr" version = "2.7.5" @@ -455,6 +539,18 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" +[[package]] +name = "nix" +version = "0.31.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf20d2fde8ff38632c426f1165ed7436270b44f199fc55284c38276f9db47c3d" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "nom" version = "7.1.3" @@ -493,6 +589,21 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" +[[package]] +name = "objc2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a12a8ed07aefc768292f076dc3ac8c48f3781c8f2d5851dd3d98950e8c5a89f" +dependencies = [ + "objc2-encode", +] + +[[package]] +name = "objc2-encode" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33" + [[package]] name = "once_cell" version = "1.21.3" @@ -510,12 +621,15 @@ name = "otadump" version = "0.1.3" dependencies = [ "anyhow", + "brotli", "bzip2", "clap", "color-print", + "ctrlc", "hex", "indicatif", "liblzma", + "lz4-sys", "memmap2", "nom", "nom-derive", @@ -1007,6 +1121,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + [[package]] name = "windows-sys" version = "0.52.0" @@ -1034,6 +1154,15 @@ dependencies = [ "windows-targets 0.53.2", ] +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-targets" version = "0.52.6" diff --git a/Cargo.toml b/Cargo.toml index 542c789..c2280f7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,12 +26,19 @@ protoc-bin-vendored = "3.2.0" [dependencies] anyhow = "1.0.69" +brotli = "8.0.2" bzip2 = { version = "0.6.0" } clap = { version = "4.2.1", features = ["derive"] } color-print = "0.3.4" +ctrlc = "3.5.0" hex = "0.4.3" indicatif = "0.17.3" liblzma = { version = "0.4.2", features = ["static"] } +# Pinned exactly. The only published 1.11.1 is `1.11.1+lz4-1.10.0` (Cargo +# ignores `+metadata` in requirements, so the lockfile checksum enforces it): +# it bundles liblz4 1.10.0, byte-identical to the AOSP platform/external/lz4 +# sources that Android lz4diff uses. +lz4-sys = "=1.11.1" memmap2 = "0.9.5" nom = "7.1.1" nom-derive = "0.10.1" @@ -49,3 +56,14 @@ codegen-units = 1 debug = 0 lto = true strip = true + +# Static CLI builds (CI + install.sh): panic=abort drops unwinding machinery. +# Kept out of [profile.release] because maturin wheel builds use the `release` +# profile and must preserve unwinding so panics surface as Python exceptions. +[profile.release-cli] +codegen-units = 1 +debug = 0 +inherits = "release" +lto = true +panic = "abort" +strip = true diff --git a/README.md b/README.md index 5da5dd6..5dd320c 100644 --- a/README.md +++ b/README.md @@ -59,9 +59,9 @@ curl -sS https://raw.githubusercontent.com/crazystylus/otadump/mainline/install. Otherwise, using Cargo: ```sh -# Needs LZMA, Protobuf and pkg-config libraries installed. -# - On macOS: brew install protobuf xz pkg-config -# - On Debian / Ubuntu: apt install liblzma-dev protobuf-compiler pkg-config +# Needs LZMA and pkg-config libraries installed. +# - On macOS: brew install xz pkg-config +# - On Debian / Ubuntu: apt install liblzma-dev pkg-config cargo install --locked otadump ``` @@ -80,6 +80,11 @@ otadump ota.zip # Run on payload.bin file. otadump payload.bin + +# Apply a SOURCE_COPY, SOURCE_BSDIFF, BROTLI_BSDIFF, PUFFDIFF (inner BSDIFF or +# ZUCCHINI), standalone ZUCCHINI, or LZ4DIFF (inner BSDIFF or PUFFDIFF) delta +# payload with matching base images. +otadump delta.zip --output-dir output --source-dir source-images ``` ### Python @@ -97,11 +102,47 @@ otadump.extract( Path("output"), partitions=["boot", "system"], overwrite=True, + source_dir=Path("source-images"), ) ``` -The optional keyword arguments are `num_threads`, `overwrite`, `partitions`, -and `verify`. Extraction releases the Python GIL. +The optional keyword arguments are `num_threads`, `overwrite`, `partitions`, `verify`, and `source_dir`. +Set `source_dir` to a directory containing matching base partition images for supported delta operations. +Extraction generates declared dm-verity hash-tree and FEC extents before verifying each partition. +Extraction releases the Python GIL. + +### Local release packages + +Build the Linux x86-64 GNU CLI with `cargo build --profile release-cli --bin otadump --locked` +(`release-cli` adds `panic = "abort"` on top of `release`; the plain `release` +profile keeps unwinding so Python wheels can raise panics as exceptions). +Build the ABI3 Python wheel with `uv build --wheel`. +The wheel supports Python 3.9 and later and should be produced inside a +manylinux2014 (manylinux_2_17) environment. +Release archives and wheels include the notices for vendored Puffin and LZ4 code. + +### Native ZUCCHINI artifact fetch (Linux x86-64 GNU builds) + +Linux x86-64 GNU builds fetch a pinned native archive when `OTADUMP_NATIVE_DIR` +is not set. The helper requires Python 3 and enforces an immutable +`https://raw.githubusercontent.com/AndroidDumps/otadump//...` URL from +`native/zucchini/ARTIFACT_BUNDLE_LOCK.json`. + +- `OTADUMP_NATIVE_DIR`: bypass download/cache and use a fully materialized + directory containing `include/`, `lib/`, `licenses/`, and `provenance.txt`. +- `OTADUMP_NATIVE_CACHE`: cache root for the pinned `tar.gz` bundle. +- `OTADUMP_NATIVE_PRESEED`: local preseed source; either a directory (already + extracted tree) or a `.tar.gz` archive copied into cache before verification. +- `OTADUMP_NATIVE_OFFLINE`: non-empty value disables network fetch. Builds fail + if cache/preseed is missing. + +Behavior summary: + +1. If output already exists, checksums are verified and reused. +2. Concurrent fetches use a cache lock and unique staging directories; one + process atomically wins publication and all contenders verify the winner. +3. Extraction rejects links/path traversal and rejects extra unpinned files or + directories outside the checksum lock. ## Contributors diff --git a/build.rs b/build.rs index 38cec17..edd5a24 100644 --- a/build.rs +++ b/build.rs @@ -1,4 +1,78 @@ fn main() { + println!("cargo:rustc-check-cfg=cfg(otadump_zucchini)"); + println!("cargo:rerun-if-changed=src/protos/chromeos_update_engine/update_metadata.proto"); + println!("cargo:rerun-if-changed=native/zucchini/ARTIFACT_BUNDLE_LOCK.json"); + println!("cargo:rerun-if-changed=native/zucchini/ARTIFACT_LOCK.sha256"); + println!("cargo:rerun-if-changed=scripts/fetch-native-artifact.py"); + + println!("cargo:rerun-if-env-changed=OTADUMP_NATIVE_DIR"); + println!("cargo:rerun-if-env-changed=OTADUMP_NATIVE_OFFLINE"); + println!("cargo:rerun-if-env-changed=OTADUMP_NATIVE_CACHE"); + println!("cargo:rerun-if-env-changed=OTADUMP_NATIVE_PRESEED"); + if std::env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("linux") + && std::env::var("CARGO_CFG_TARGET_ARCH").as_deref() == Ok("x86_64") + && std::env::var("CARGO_CFG_TARGET_ENV").as_deref() == Ok("gnu") + { + let manifest_dir = std::path::PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + let bundle_lock = manifest_dir.join("native/zucchini/ARTIFACT_BUNDLE_LOCK.json"); + let checksum_lock = manifest_dir.join("native/zucchini/ARTIFACT_LOCK.sha256"); + let root = if let Some(path) = std::env::var_os("OTADUMP_NATIVE_DIR") { + std::path::PathBuf::from(path) + } else { + let out = std::path::PathBuf::from(std::env::var_os("OUT_DIR").unwrap()); + let artifact = out.join("otadump-native/linux-x86_64-gnu"); + let cache = std::env::var_os("OTADUMP_NATIVE_CACHE") + .map(std::path::PathBuf::from) + .unwrap_or_else(|| out.join("cache/otadump-native")); + let mut command = std::process::Command::new("python3"); + command + .arg(manifest_dir.join("scripts/fetch-native-artifact.py")) + .arg("--bundle-lock") + .arg(&bundle_lock) + .arg("--checksum-lock") + .arg(&checksum_lock) + .arg("--out") + .arg(&artifact) + .arg("--cache") + .arg(cache); + let status = command.status().expect("unable to run native artifact fetch helper"); + assert!(status.success(), "native artifact fetch helper failed"); + artifact + }; + let include = root.join("include"); + let archive = root.join("lib/libotadump_zucchini.a"); + assert!( + include.join("zucchini_ffi.h").is_file(), + "missing Zucchini artifact header: {}", + include.display() + ); + assert!(archive.is_file(), "missing Zucchini artifact archive: {}", archive.display()); + let status = std::process::Command::new("sha256sum") + .arg("-c") + .arg(&checksum_lock) + .current_dir(&root) + .status() + .expect("unable to verify Zucchini artifact checksum lock"); + assert!( + status.success(), + "Zucchini artifact checksum lock failed: {}", + checksum_lock.display() + ); + println!("cargo:rustc-link-search=native={}", root.join("lib").display()); + println!("cargo:rustc-link-lib=static=otadump_zucchini"); + let cxx_archive = std::process::Command::new("g++") + .args(["-print-file-name=libstdc++.a"]) + .output() + .expect("unable to locate the C++ runtime") + .stdout; + let cxx_archive = std::path::PathBuf::from(String::from_utf8(cxx_archive).unwrap().trim()); + assert!(cxx_archive.is_file(), "static libstdc++ is required: {}", cxx_archive.display()); + println!("cargo:rustc-link-search=native={}", cxx_archive.parent().unwrap().display()); + println!("cargo:rustc-link-lib=static=stdc++"); + println!("cargo:rustc-link-arg=-static-libgcc"); + println!("cargo:rustc-cfg=otadump_zucchini"); + } + let protoc = protoc_bin_vendored::protoc_bin_path().expect("unable to find vendored protoc"); let mut config = prost_build::Config::new(); config.protoc_executable(protoc); diff --git a/docs/zucchini-rust-port.md b/docs/zucchini-rust-port.md new file mode 100644 index 0000000..2b2452c --- /dev/null +++ b/docs/zucchini-rust-port.md @@ -0,0 +1,64 @@ +# Native ZUCCHINI artifact architecture + +Status: active for Linux x86-64 GNU builds. + +## Overview + +`otadump` links a prebuilt `libotadump_zucchini.a` archive for Linux x86-64 GNU +targets and keeps all other targets on the pure-Rust delta stack without the +native bridge. + +The artifact is pinned by: + +- `native/zucchini/ARTIFACT_BUNDLE_LOCK.json` (immutable commit URL + archive digest) +- `native/zucchini/ARTIFACT_LOCK.sha256` (per-file digest lock inside the extracted tree) + +`build.rs` invokes `scripts/fetch-native-artifact.py` unless +`OTADUMP_NATIVE_DIR` is set. + +## Security and integrity guarantees + +The fetch helper enforces: + +1. HTTPS-only URL policy. +2. Exact host policy: `raw.githubusercontent.com`. +3. Exact immutable path policy: + `/AndroidDumps/otadump//native/artifacts/zucchini/zucchini-linux-x86_64-gnu.tar.gz`. +4. Download timeout and max-size bounds. +5. Archive digest verification before extraction. +6. Tar safety checks: no traversal, no symlinks/hardlinks, files/dirs only. +7. Strict extracted-tree verification: no extra files/directories outside the + checksum lock's expected structure. + +## Concurrency model + +Concurrent builds sharing `OTADUMP_NATIVE_CACHE` coordinate through a fetch lock +keyed by bundle digest. + +- Only one process writes the cached archive at a time. +- Each process stages extraction into a unique temp directory. +- Publication uses atomic rename; one contender wins. +- Losers verify the published tree and continue. + +This avoids partial publication and keeps cache/output deterministic under +parallel runners. + +## Offline and preseed behavior + +- `OTADUMP_NATIVE_OFFLINE=1`: network disabled; build requires cached or + preseeded artifact. +- `OTADUMP_NATIVE_PRESEED=/path/to/archive.tar.gz`: copies the archive into the + cache before checksum verification. +- `OTADUMP_NATIVE_PRESEED=/path/to/tree`: verifies and publishes a fully + materialized extracted tree. + +## Test coverage + +`tests/native_artifact_fetch_test.py` covers: + +- parallel fetch/publish contention +- URL policy rejections +- strict extra-file rejection +- checksum tamper rejection +- traversal rejection +- offline-with-empty-cache rejection diff --git a/licenses/LICENSE.lz4 b/licenses/LICENSE.lz4 new file mode 100644 index 0000000..4884916 --- /dev/null +++ b/licenses/LICENSE.lz4 @@ -0,0 +1,24 @@ +LZ4 Library +Copyright (c) 2011-2020, Yann Collet +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, this + list of conditions and the following disclaimer in the documentation and/or + other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/licenses/README.md b/licenses/README.md new file mode 100644 index 0000000..7cc657a --- /dev/null +++ b/licenses/README.md @@ -0,0 +1,10 @@ +# Third-party license notices + +`LICENSE.lz4` is the BSD-2-Clause license of upstream liblz4, statically linked +into release binaries and wheels via the pinned +[`lz4-sys`](https://crates.io/crates/lz4-sys) `=1.11.1+lz4-1.10.0` crate. + +That crate bundles the liblz4 1.10.0 C sources. The `lz4.c`, `lz4.h`, +`lz4hc.c`, and `lz4hc.h` files it compiles are byte-identical (SHA-256 +verified) to the AOSP `platform/external/lz4` sources pinned at commit +`734e07032602e9a72fcc9701028b0aee45147fcd`, which Android `lz4diff` uses. diff --git a/native/zucchini/ARTIFACT_BUNDLE_LOCK.json b/native/zucchini/ARTIFACT_BUNDLE_LOCK.json new file mode 100644 index 0000000..5daebd7 --- /dev/null +++ b/native/zucchini/ARTIFACT_BUNDLE_LOCK.json @@ -0,0 +1,7 @@ +{ + "schema": 1, + "commit": "9534f6e24296b1682d53cdac19da0b81fc97fcbc", + "url": "https://raw.githubusercontent.com/AndroidDumps/otadump/9534f6e24296b1682d53cdac19da0b81fc97fcbc/native/artifacts/zucchini/zucchini-linux-x86_64-gnu.tar.gz", + "sha256": "70655ef27c2327291071590ee726f13cde1abd6f9ed0c3c4dae32103834ef16f", + "extract_root": "linux-x86_64-gnu" +} diff --git a/native/zucchini/ARTIFACT_LOCK.sha256 b/native/zucchini/ARTIFACT_LOCK.sha256 new file mode 100644 index 0000000..b795a88 --- /dev/null +++ b/native/zucchini/ARTIFACT_LOCK.sha256 @@ -0,0 +1,5 @@ +2f7138dd9ff6b0d4cb081536b3a8a16cb5c3ff09ee96d73f6351e8cae0d6a99c lib/libotadump_zucchini.a +53aabfa70f3c3061b63017a8d7c5e82bf62aab22c68b6189fdc1c8f99927c31b SHA256SUMS +11696b3b8d3dc920e2375d98fd95028cd32c6c1ea86b082d6f800e00fb357703 include/zucchini_ffi.h +845022e0c1db1abb41a6ba4cd3c4b674ec290f3359d9d3c78ae558d4c0ed9308 licenses/LICENSE.zucchini +a0a33d4ad9d6a62311eae87175db1bad4686a8d6a0edd47ac1584be7569d67e4 provenance.txt diff --git a/pyproject.toml b/pyproject.toml index 717b81c..9fda8fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,11 @@ description = "Extract partitions from Android OTA files" readme = "README.md" requires-python = ">=3.9" license = "MIT" +license-files = [ + "LICENSE", + "src/puffin/LICENSE", + "licenses/LICENSE.lz4", +] authors = [ { name = "Kartik Sharma", email = "kartik.sharma522@gmail.com" }, { name = "Ajeet D'Souza", email = "98ajeet@gmail.com" }, @@ -19,4 +24,8 @@ classifiers = [ ] [tool.maturin] +# Wheels build with the default `release` profile, which keeps `panic = +# "unwind"` so extraction-thread panics surface as Python exceptions. The +# `release-cli` profile (panic = "abort") is for the static CLI only. features = ["python"] +locked = true diff --git a/scripts/fetch-native-artifact.py b/scripts/fetch-native-artifact.py new file mode 100755 index 0000000..a6aa586 --- /dev/null +++ b/scripts/fetch-native-artifact.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 +import argparse +import hashlib +import json +import os +import pathlib +import shutil +import sys +import tarfile +import tempfile +import time +import urllib.request +import uuid +from urllib.parse import urlsplit + + +DOWNLOAD_TIMEOUT_SECONDS = 30 +DOWNLOAD_MAX_BYTES = 256 * 1024 * 1024 +LOCK_WAIT_SECONDS = 120 +LOCK_POLL_SECONDS = 0.05 + + +class FileLock: + def __init__(self, path: pathlib.Path, timeout_seconds: float) -> None: + self.path = path + self.timeout_seconds = timeout_seconds + self.fd = -1 + + def __enter__(self) -> "FileLock": + deadline = time.monotonic() + self.timeout_seconds + while True: + try: + self.fd = os.open(self.path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + os.write(self.fd, f"{os.getpid()}\n".encode("ascii")) + return self + except FileExistsError: + if time.monotonic() >= deadline: + raise RuntimeError(f"timed out waiting for lock: {self.path}") + time.sleep(LOCK_POLL_SECONDS) + + def __exit__(self, _exc_type, _exc, _tb) -> None: + if self.fd >= 0: + os.close(self.fd) + self.fd = -1 + try: + self.path.unlink() + except FileNotFoundError: + pass + + +def sha256_file(path: pathlib.Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def parse_sha256_lock(path: pathlib.Path) -> dict[str, str]: + checksums = {} + for line in path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line: + continue + digest, relpath = line.split(maxsplit=1) + checksums[relpath.strip()] = digest + return checksums + + +def validate_bundle_url(url: str, commit: str) -> None: + if len(commit) != 40 or any(ch not in "0123456789abcdef" for ch in commit): + raise RuntimeError(f"bundle lock commit must be a lowercase 40-char hex SHA: {commit}") + parsed = urlsplit(url) + if parsed.scheme != "https": + raise RuntimeError(f"native artifact URL must use HTTPS: {url}") + if parsed.hostname != "raw.githubusercontent.com": + raise RuntimeError( + "native artifact URL host must be raw.githubusercontent.com: " + f"{parsed.hostname or ''}" + ) + if parsed.query or parsed.fragment: + raise RuntimeError("native artifact URL must not include query parameters or fragments") + expected_path = ( + f"/AndroidDumps/otadump/{commit}/native/artifacts/zucchini/" + "zucchini-linux-x86_64-gnu.tar.gz" + ) + if parsed.path != expected_path: + raise RuntimeError( + "native artifact URL path must match the immutable commit artifact path: " + f"expected {expected_path}, got {parsed.path}" + ) + + +def verify_tree(root: pathlib.Path, checksums: dict[str, str]) -> None: + expected_files = set(checksums) + allowed_dirs = {"."} + for relpath in checksums: + parent = pathlib.PurePosixPath(relpath).parent + while True: + allowed_dirs.add(parent.as_posix()) + if parent == pathlib.PurePosixPath("."): + break + parent = parent.parent + + discovered_files = set() + for candidate in root.rglob("*"): + relpath = candidate.relative_to(root).as_posix() + if candidate.is_symlink(): + raise RuntimeError(f"links are not allowed in native artifact tree: {relpath}") + if candidate.is_dir(): + if relpath not in allowed_dirs: + raise RuntimeError(f"unexpected native artifact directory: {relpath}") + continue + if not candidate.is_file(): + raise RuntimeError(f"unsupported native artifact entry type: {relpath}") + if relpath not in expected_files: + raise RuntimeError(f"unexpected native artifact file: {relpath}") + discovered_files.add(relpath) + + missing = sorted(expected_files - discovered_files) + if missing: + raise RuntimeError(f"missing native artifact files: {', '.join(missing)}") + + for relpath, expected in checksums.items(): + target = root / relpath + if not target.is_file(): + raise RuntimeError(f"missing native artifact file: {target}") + actual = sha256_file(target) + if actual != expected: + raise RuntimeError( + f"native artifact checksum mismatch for {relpath}: expected {expected}, got {actual}" + ) + + +def safe_extract(archive: pathlib.Path, destination: pathlib.Path) -> None: + with tarfile.open(archive, "r:gz") as tar: + for member in tar.getmembers(): + name = pathlib.PurePosixPath(member.name) + if name.is_absolute() or ".." in name.parts: + raise RuntimeError(f"unsafe entry in native artifact archive: {member.name}") + if member.issym() or member.islnk(): + raise RuntimeError(f"links are not allowed in native artifact archive: {member.name}") + if not (member.isfile() or member.isdir()): + raise RuntimeError(f"unsupported tar entry in native artifact archive: {member.name}") + tar.extractall(destination) + + +def download_archive(url: str, archive_path: pathlib.Path) -> None: + with urllib.request.urlopen(url, timeout=DOWNLOAD_TIMEOUT_SECONDS) as response: + header_length = response.headers.get("Content-Length") + if header_length is not None: + try: + if int(header_length) > DOWNLOAD_MAX_BYTES: + raise RuntimeError( + "native artifact archive is too large: " + f"{header_length} bytes exceeds {DOWNLOAD_MAX_BYTES}" + ) + except ValueError: + pass + with tempfile.NamedTemporaryFile(delete=False, dir=archive_path.parent) as handle: + total = 0 + while True: + chunk = response.read(1024 * 1024) + if not chunk: + break + total += len(chunk) + if total > DOWNLOAD_MAX_BYTES: + raise RuntimeError( + "native artifact archive is too large: " + f"streamed {total} bytes exceeds {DOWNLOAD_MAX_BYTES}" + ) + handle.write(chunk) + temp_name = handle.name + pathlib.Path(temp_name).replace(archive_path) + + +def publish_tree(source: pathlib.Path, out_dir: pathlib.Path, checksums: dict[str, str]) -> int: + out_dir.parent.mkdir(parents=True, exist_ok=True) + staging = out_dir.with_name(f"{out_dir.name}.tmp-{os.getpid()}-{uuid.uuid4().hex}") + if staging.exists(): + shutil.rmtree(staging) + shutil.copytree(source, staging) + try: + os.replace(staging, out_dir) + except OSError: + if staging.exists(): + shutil.rmtree(staging) + if out_dir.is_dir(): + verify_tree(out_dir, checksums) + return 0 + raise + return 0 + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--bundle-lock", required=True) + parser.add_argument("--checksum-lock", required=True) + parser.add_argument("--out", required=True) + parser.add_argument("--cache", required=True) + args = parser.parse_args() + + bundle_lock = json.loads(pathlib.Path(args.bundle_lock).read_text(encoding="utf-8")) + validate_bundle_url(bundle_lock["url"], bundle_lock["commit"]) + checksums = parse_sha256_lock(pathlib.Path(args.checksum_lock)) + out_dir = pathlib.Path(args.out) + + if out_dir.is_dir(): + verify_tree(out_dir, checksums) + return 0 + + cache_root = pathlib.Path(args.cache) + cache_root.mkdir(parents=True, exist_ok=True) + archive_path = cache_root / f"{bundle_lock['sha256']}.tar.gz" + fetch_lock = cache_root / f".{bundle_lock['sha256']}.fetch.lock" + + preseed = os.environ.get("OTADUMP_NATIVE_PRESEED") + offline = os.environ.get("OTADUMP_NATIVE_OFFLINE") + + if preseed: + preseed_path = pathlib.Path(preseed) + if preseed_path.is_dir(): + verify_tree(preseed_path, checksums) + return publish_tree(preseed_path, out_dir, checksums) + if not preseed_path.is_file(): + raise RuntimeError(f"OTADUMP_NATIVE_PRESEED path does not exist: {preseed_path}") + + with FileLock(fetch_lock, LOCK_WAIT_SECONDS): + if preseed and pathlib.Path(preseed).is_file(): + shutil.copy2(pathlib.Path(preseed), archive_path) + if not archive_path.is_file(): + if offline: + raise RuntimeError( + "offline native artifact mode is enabled but cache is empty; " + "set OTADUMP_NATIVE_PRESEED or unset OTADUMP_NATIVE_OFFLINE" + ) + download_archive(bundle_lock["url"], archive_path) + + archive_digest = sha256_file(archive_path) + if archive_digest != bundle_lock["sha256"]: + raise RuntimeError( + f"native artifact archive checksum mismatch: expected {bundle_lock['sha256']}, got {archive_digest}" + ) + + with tempfile.TemporaryDirectory(dir=cache_root) as temp: + temp_root = pathlib.Path(temp) + safe_extract(archive_path, temp_root) + extracted = temp_root / bundle_lock["extract_root"] + if not extracted.is_dir(): + raise RuntimeError(f"missing artifact root after extraction: {bundle_lock['extract_root']}") + verify_tree(extracted, checksums) + + return publish_tree(extracted, out_dir, checksums) + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except Exception as error: + print(f"otadump native artifact fetch failed: {error}", file=sys.stderr) + raise SystemExit(2) diff --git a/src/bsdiff.rs b/src/bsdiff.rs new file mode 100644 index 0000000..6255ee6 --- /dev/null +++ b/src/bsdiff.rs @@ -0,0 +1,121 @@ +//! Local BSDF2/BSDIFF40 patch header and stream reader. +//! +//! Replaces the external `bsdiff-android` crate, of which only +//! `parse_bsdf2_header` was ever linked. The apply loop itself lives in +//! `lib.rs`; this module reproduces the header semantics bit-exactly: +//! +//! - `BSDIFF40`: all three streams are bzip2-compressed. +//! - `BSDF2`: magic bytes 5..8 select a per-stream algorithm +//! (`0` = stored, `1` = bzip2, `2` = brotli). +//! - Lengths use the bspatch sign-magnitude 64-bit encoding (`offtin`). +//! - The extra stream spans the remainder of the patch. + +use std::io::Read; + +use anyhow::{Result, anyhow, bail, ensure}; + +enum Algorithm { + None, + Bz2, + Brotli, +} + +type ParsedPatch = (i64, Vec, Vec, Vec); + +pub fn validate_output_len(patch: &[u8], expected: usize) -> Result<()> { + ensure!(patch.len() >= 32, "Patch data too short"); + ensure!( + &patch[..8] == b"BSDIFF40" || &patch[..5] == b"BSDF2", + "Invalid BSDIFF/BSDF2 magic header" + ); + let encoded = u64::from_le_bytes( + patch[24..32].try_into().map_err(|_| anyhow!("Invalid BSDIFF output length field"))?, + ); + ensure!(encoded & (1 << 63) == 0, "Negative output length in patch header"); + let actual = + usize::try_from(encoded).map_err(|_| anyhow!("Patch output length is too large"))?; + ensure!(actual == expected, "Patch output length mismatch: expected {expected}, got {actual}"); + Ok(()) +} + +fn algorithm_from_byte(value: u8) -> Result { + match value { + 0 => Ok(Algorithm::None), + 1 => Ok(Algorithm::Bz2), + 2 => Ok(Algorithm::Brotli), + _ => bail!("Unknown BSDF2 compression algorithm: {value}"), + } +} + +/// Reads a sign-magnitude 64-bit integer as used by bspatch. This is not plain +/// little-endian: the top bit encodes the sign of the magnitude. +#[inline] +fn offtin(buf: [u8; 8]) -> i64 { + let value = i64::from_le_bytes(buf); + if value & (1 << 63) == 0 { value } else { -(value & !(1 << 63)) } +} + +fn decompress_bz2(data: &[u8]) -> Result> { + let mut output = Vec::new(); + let mut decoder = bzip2::read::BzDecoder::new(data); + decoder.read_to_end(&mut output)?; + Ok(output) +} + +fn decompress_brotli(data: &[u8]) -> Result> { + let mut output = Vec::new(); + let mut decoder = brotli::Decompressor::new(data, 4096); + decoder.read_to_end(&mut output)?; + Ok(output) +} + +fn decompress(algorithm: Algorithm, data: &[u8]) -> Result> { + match algorithm { + Algorithm::None => Ok(data.to_vec()), + Algorithm::Bz2 => decompress_bz2(data), + Algorithm::Brotli => decompress_brotli(data), + } +} + +/// Parses a BSDIFF40/BSDF2 patch header and returns the decompressed +/// `(new_size, control, diff, extra)` streams. +pub fn parse_header(patch: &[u8]) -> Result { + ensure!(patch.len() >= 32, "Patch data too short"); + + let magic = &patch[0..8]; + let (control_algorithm, diff_algorithm, extra_algorithm) = if magic == b"BSDIFF40" { + (Algorithm::Bz2, Algorithm::Bz2, Algorithm::Bz2) + } else if &magic[0..5] == b"BSDF2" { + ( + algorithm_from_byte(magic[5])?, + algorithm_from_byte(magic[6])?, + algorithm_from_byte(magic[7])?, + ) + } else { + bail!("Invalid BSDIFF/BSDF2 magic header"); + }; + + let control_length = offtin(patch[8..16].try_into()?); + let diff_length = offtin(patch[16..24].try_into()?); + let new_size = offtin(patch[24..32].try_into()?); + ensure!( + control_length >= 0 && diff_length >= 0 && new_size >= 0, + "Negative length in patch header" + ); + let (control_length, diff_length) = (control_length as usize, diff_length as usize); + + let control_end = 32usize + .checked_add(control_length) + .ok_or_else(|| anyhow::anyhow!("Control stream length overflows patch bounds"))?; + ensure!(control_end <= patch.len(), "Control stream exceeds patch bounds"); + let diff_end = control_end + .checked_add(diff_length) + .ok_or_else(|| anyhow::anyhow!("Diff stream length overflows patch bounds"))?; + ensure!(diff_end <= patch.len(), "Diff stream exceeds patch bounds"); + + let control = decompress(control_algorithm, &patch[32..control_end])?; + ensure!(control.len() % 24 == 0, "Invalid control data length (not multiple of 24)"); + let diff = decompress(diff_algorithm, &patch[control_end..diff_end])?; + let extra = decompress(extra_algorithm, &patch[diff_end..])?; + Ok((new_size, control, diff, extra)) +} diff --git a/src/lib.rs b/src/lib.rs index 5da31f5..a7b018a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,31 +3,39 @@ mod chromeos_update_engine { include!(concat!(env!("OUT_DIR"), "/chromeos_update_engine.rs")); } +mod bsdiff; +pub mod lz4; +mod lz4diff; mod payload; +mod puffin; #[cfg(feature = "python")] mod python; +mod verity; +pub mod zucchini; use std::cmp::Reverse; use std::collections::HashSet; use std::fs::{self, File, OpenOptions}; use std::io::Read; use std::num::NonZero; -use std::ops::Mul; -use std::path::Path; -use std::sync::OnceLock; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::ops::Range; +use std::path::{Component, Path, PathBuf}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, OnceLock}; use std::{error, result, slice, thread}; use anyhow::{Context as _, Error, Result, bail, ensure}; +use brotli::Decompressor as BrotliDecoder; use bzip2::read::BzDecoder; use chromeos_update_engine::install_operation::Type; use chromeos_update_engine::{DeltaArchiveManifest, InstallOperation, PartitionUpdate}; use liblzma::read::XzDecoder; use memmap2::{Mmap, MmapMut, MmapOptions}; use prost::Message as _; -use rayon::{BroadcastContext, ThreadPoolBuilder}; +use rayon::ThreadPoolBuilder; use ring::digest; use sync_unsafe_cell::SyncUnsafeCell; +use tempfile::TempPath; use zip::ZipArchive; use zip::result::ZipError; use zstd::Decoder as ZstdDecoder; @@ -37,6 +45,55 @@ pub use crate::payload::Payload; // The chunk size is the number of bytes we verify before we count one "tick" in // the progress tracker. const VERIFY_CHUNK_SIZE: usize = 2 * 1024 * 1024; // 2 MiB +const BROTLI_DECODE_CHUNK_SIZE: usize = 32 * 1024; + +#[derive(Clone, Debug, Default)] +pub struct CancellationToken { + cancelled: Arc, +} + +impl CancellationToken { + pub fn new() -> Self { + Self::default() + } + + pub fn cancel(&self) { + self.cancelled.store(true, Ordering::Release); + } + + pub fn is_cancelled(&self) -> bool { + self.cancelled.load(Ordering::Acquire) + } + + fn check(&self) -> Result<()> { + if self.is_cancelled() { + return Err(ExtractionCancelled.into()); + } + Ok(()) + } +} + +#[derive(Debug)] +pub struct ExtractionCancelled; + +impl std::fmt::Display for ExtractionCancelled { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("Extraction cancelled") + } +} + +impl error::Error for ExtractionCancelled {} + +pub fn is_cancellation(error: &(dyn error::Error + 'static)) -> bool { + let mut error = Some(error); + while let Some(current) = error { + if current.downcast_ref::().is_some() { + return true; + } + error = current.source(); + } + false +} pub struct ExtractOptions<'a> { num_threads: Option, @@ -44,6 +101,8 @@ pub struct ExtractOptions<'a> { partitions: Option>, progress_reporter: &'a dyn ProgressReporter, verify: bool, + source_dir: Option, + cancellation_token: CancellationToken, } impl<'a> ExtractOptions<'a> { @@ -55,6 +114,8 @@ impl<'a> ExtractOptions<'a> { partitions: None, progress_reporter: &NoOpProgressReporter, verify: true, + source_dir: None, + cancellation_token: CancellationToken::new(), } } @@ -70,14 +131,18 @@ impl<'a> ExtractOptions<'a> { { let payload_file = payload_file.as_ref(); let output_dir = output_dir.as_ref(); - self.extract_impl(payload_file, output_dir)?; - Ok(()) + match self.extract_impl(payload_file, output_dir) { + Ok(()) => Ok(()), + Err(error) if is_cancellation(error.as_ref()) => Err(Box::new(ExtractionCancelled)), + Err(error) => Err(error.into()), + } } fn extract_impl(&self, payload_file: &Path, output_dir: &Path) -> Result<()> { + self.cancellation_token.check()?; self.progress_reporter.report_progress(0.); - let payload_file = Self::open_payload_file(payload_file)?; + let payload_file = self.open_payload_file(payload_file)?; let payload = Payload::parse(&payload_file)?; let mut manifest = @@ -97,27 +162,25 @@ impl<'a> ExtractOptions<'a> { let extract_ops = manifest.partitions.iter().map(|update| update.operations.len()).sum::(); - let verify_ops = if self.verify { - manifest - .partitions - .iter() - .map(|update| { - let partition_size = - update.new_partition_info.as_ref().and_then(|info| info.size).unwrap_or(0) - as usize; - partition_size.div_ceil(VERIFY_CHUNK_SIZE) - }) - .sum() - } else { - 0 - }; + let verify_ops: usize = manifest + .partitions + .iter() + .map(|update| { + let partition_size = + update.new_partition_info.as_ref().and_then(|info| info.size).unwrap_or(0) + as usize; + partition_size.div_ceil(VERIFY_CHUNK_SIZE) + }) + .sum(); let total_ops = extract_ops + verify_ops; let total_ops_completed = AtomicUsize::new(0); let block_size = manifest.block_size.context("block_size not defined")? as usize; + ensure!(block_size > 0, "block_size must be greater than zero"); // Ensure that all partitions to be extracted are present in the manifest. for partition_name in self.partitions.iter().flatten() { + self.cancellation_token.check()?; ensure!( manifest.partitions.iter().any(|update| &update.partition_name == partition_name), "Partition not found: {partition_name}", @@ -127,6 +190,55 @@ impl<'a> ExtractOptions<'a> { fs::create_dir_all(output_dir) .with_context(|| format!("Could not create output directory: {output_dir:?}"))?; + let validated = manifest + .partitions + .iter() + .map(|update| self.validate_partition(update, &payload, block_size)) + .collect::>>()?; + let (sources, verity_configs): (Vec<_>, Vec<_>) = validated.into_iter().unzip(); + + let mut partition_names = HashSet::with_capacity(manifest.partitions.len()); + let mut destinations = Vec::with_capacity(manifest.partitions.len()); + for update in &manifest.partitions { + self.cancellation_token.check()?; + ensure!( + partition_names.insert(&update.partition_name), + "Duplicate partition mapping: {:?}", + update.partition_name + ); + destinations.push(output_dir.join(format!("{}.img", update.partition_name))); + } + for (update, destination) in manifest.partitions.iter().zip(&destinations) { + self.cancellation_token.check()?; + for source in sources.iter().flatten() { + ensure!( + !paths_alias(&source.path, destination)?, + "Source partition image {:?} and destination partition image {destination:?} refer to the same file", + source.path, + ); + } + if !self.overwrite { + ensure!( + !path_entry_exists(destination)?, + "Destination partition image already exists for {:?}: {destination:?}", + update.partition_name + ); + } + } + + let staging_dir = tempfile::Builder::new() + .prefix(".otadump-") + .tempdir_in(output_dir) + .context("Unable to create staging directory")?; + let mut staging_paths = Vec::with_capacity(manifest.partitions.len()); + let mut partitions = Vec::with_capacity(manifest.partitions.len()); + for update in &manifest.partitions { + self.cancellation_token.check()?; + let (path, partition) = self.create_staged_partition(update, staging_dir.path())?; + staging_paths.push(path); + partitions.push(partition); + } + let num_threads = self .num_threads .unwrap_or_else(|| thread::available_parallelism().map(NonZero::get).unwrap_or(1)) @@ -137,45 +249,98 @@ impl<'a> ExtractOptions<'a> { .context("Unable to start threadpool")?; let mut error = OnceLock::new(); - threadpool.in_place_scope_fifo(|scope| -> Result<()> { - for update in &manifest.partitions { + threadpool.in_place_scope_fifo(|scope| { + for ((update, source), partition) in + manifest.partitions.iter().zip(sources).zip(partitions) + { // Exit early if an error has occurred. if error.get().is_some() { break; } + if let Err(cancelled) = self.cancellation_token.check() { + _ = error.set(cancelled); + break; + } // Create and broadcast a task to the threadpool. - let partition = self.open_partition_file(update, output_dir)?; let task = Task { payload: &payload, block_size, - verify: self.verify, update, + source, op_idx: AtomicUsize::new(0), partition: SyncUnsafeCell::new(partition), total_ops, total_ops_completed: &total_ops_completed, progress_reporter: self.progress_reporter, + cancellation_token: &self.cancellation_token, error: &error, }; - scope.spawn_broadcast(move |_, ctx| { - if let Err(e) = task.run(ctx) { + scope.spawn_broadcast(move |_, _| { + if let Err(e) = task.run() { _ = task.error.set(e); } }); } - Ok(()) - })?; + }); + self.cancellation_token.check()?; if let Some(e) = error.take() { return Err(e); } + for ((update, verity_config), staging_path) in + manifest.partitions.iter().zip(verity_configs).zip(&staging_paths) + { + self.cancellation_token.check()?; + let file = OpenOptions::new() + .read(true) + .write(true) + .open(staging_path) + .context("Unable to reopen staged partition image")?; + let mut partition = unsafe { MmapMut::map_mut(&file) } + .context("Failed to mmap staged partition image")?; + if let Some(config) = verity_config.as_ref() { + verity::generate(&mut partition, config, block_size, &self.cancellation_token) + .map_err(|error| { + if is_cancellation(error.as_ref()) { + error + } else { + error.context(format!( + "Failed to generate verity data for {:?}", + update.partition_name + )) + } + })?; + } + verify_partition( + update, + &partition, + total_ops, + &total_ops_completed, + self.progress_reporter, + &self.cancellation_token, + )?; + self.cancellation_token.check()?; + partition.flush().context("Error while flushing file to disk")?; + } + + for (staging_path, destination) in staging_paths.into_iter().zip(destinations) { + self.cancellation_token.check()?; + let result = if self.overwrite { + staging_path.persist(&destination) + } else { + staging_path.persist_noclobber(&destination) + }; + result.map_err(|error| error.error)?; + } + self.progress_reporter.report_progress(1.); + self.cancellation_token.check()?; Ok(()) } - fn open_payload_file(path: &Path) -> Result { + fn open_payload_file(&self, path: &Path) -> Result { let in_file = File::open(path) .with_context(|| format!("Failed to open file for reading: {path:?}"))?; @@ -206,9 +371,12 @@ impl<'a> ExtractOptions<'a> { let mut out_file = unsafe { MmapMut::map_mut(&out_file) } .context("Failed to mmap temporary file")?; - zip_file - .read_exact(&mut out_file) - .context("Failed to write to temporary file")?; + for chunk in out_file.chunks_mut(VERIFY_CHUNK_SIZE) { + self.cancellation_token.check()?; + zip_file + .read_exact(chunk) + .context("Failed to write to temporary file")?; + } out_file.make_read_only().context("Failed to make temporary file read-only") } } @@ -219,33 +387,183 @@ impl<'a> ExtractOptions<'a> { } } - fn open_partition_file( + fn validate_partition( &self, update: &PartitionUpdate, - partition_dir: impl AsRef, - ) -> Result { + payload: &Payload<'_>, + block_size: usize, + ) -> Result<(Option, Option)> { + validate_partition_name(&update.partition_name)?; + self.cancellation_token.check()?; + let partition_len = update .new_partition_info .as_ref() .and_then(|info| info.size) .context("Unable to determine output file size")?; + let partition_len = + usize::try_from(partition_len).context("Output partition is too large")?; + ensure!( + update.new_partition_info.as_ref().and_then(|info| info.hash.as_ref()).is_some(), + "Unable to determine output partition hash" + ); + + let needs_source = update.operations.iter().try_fold(false, |needed, op| { + self.cancellation_token.check()?; + let op_type = Type::try_from(op.r#type).context("Invalid operation")?; + validate_operation_type(op_type)?; + Ok::<_, Error>(needed || operation_needs_source(op_type)) + })?; + let source = if needs_source { Some(self.open_source_partition(update)?) } else { None }; + + let mut destination_ranges = Vec::new(); + for (op_index, op) in update.operations.iter().enumerate() { + self.cancellation_token.check()?; + let op_type = Type::try_from(op.r#type).context("Invalid operation")?; + let dst_ranges = extent_ranges(&op.dst_extents, block_size, partition_len) + .with_context(|| format!("Invalid destination extents for operation {op_index}"))?; + ensure!(!dst_ranges.is_empty(), "Operation {op_index} has no destination extents"); + let dst_len = ranges_len(&dst_ranges)?; + if let Some(declared_len) = op.dst_length { + ensure!( + usize::try_from(declared_len).ok() == Some(dst_len), + "Destination length mismatch for operation {op_index}" + ); + } + destination_ranges.extend( + dst_ranges + .iter() + .filter(|range| !range.is_empty()) + .map(|range| (range.clone(), op_index)), + ); + + if operation_needs_source(op_type) { + let source = source.as_ref().context("Source partition was not opened")?; + let src_ranges = extent_ranges(&op.src_extents, block_size, source.len) + .with_context(|| format!("Invalid source extents for operation {op_index}"))?; + ensure!(!src_ranges.is_empty(), "Operation {op_index} has no source extents"); + let src_len = ranges_len(&src_ranges)?; + if let Some(declared_len) = op.src_length { + ensure!( + usize::try_from(declared_len).ok() == Some(src_len), + "Source length mismatch for operation {op_index}" + ); + } + if self.verify { + verify_hash_over_ranges( + source.bytes(), + &src_ranges, + op.src_sha256_hash.as_deref(), + "Source extent", + &self.cancellation_token, + ) + .with_context(|| { + format!("Source verification failed for operation {op_index}") + })?; + } + } + + if operation_has_data(op_type) { + let data = extract_operation_data(payload, op) + .with_context(|| format!("Invalid data for operation {op_index}"))?; + if self.verify { + verify_hash( + data, + op.data_sha256_hash.as_deref(), + "Input", + &self.cancellation_token, + ) + .with_context(|| { + format!("Data verification failed for operation {op_index}") + })?; + } + } + } + + destination_ranges.sort_unstable_by_key(|(range, _)| (range.start, range.end)); + for pair in destination_ranges.windows(2) { + let (left, left_op) = &pair[0]; + let (right, right_op) = &pair[1]; + ensure!( + left.end <= right.start, + "Destination extents overlap between operations {left_op} and {right_op}" + ); + } - let filename = Path::new(&update.partition_name).with_extension("img"); - let path = partition_dir.as_ref().join(filename); + let verity = verity::validate(update, block_size, partition_len, &destination_ranges) + .map_err(|error| { + anyhow::anyhow!("Invalid verity metadata for {:?}: {error}", update.partition_name) + })?; + + Ok((source, verity)) + } + + fn open_source_partition(&self, update: &PartitionUpdate) -> Result { + let source_dir = self.source_dir.as_ref().with_context(|| { + format!( + "Delta operations for partition {:?} require a source directory", + update.partition_name + ) + })?; + let path = source_dir.join(format!("{}.img", update.partition_name)); + let file = File::open(&path) + .with_context(|| format!("Unable to open source partition image: {path:?}"))?; + let metadata = file + .metadata() + .with_context(|| format!("Unable to inspect source partition image: {path:?}"))?; + ensure!(metadata.is_file(), "Source partition image is not a file: {path:?}"); + let backing_len = + usize::try_from(metadata.len()).context("Source partition is too large")?; + let info = + update.old_partition_info.as_ref().context("Source partition info is missing")?; + let expected_size = info.size.context("Source partition size is missing")?; + let len = usize::try_from(expected_size).context("Source partition is too large")?; + ensure!( + backing_len >= len, + "Source partition size mismatch for {:?}: expected at least {expected_size}, got {backing_len}", + update.partition_name, + ); + let mmap = if backing_len == 0 { + None + } else { + Some( + unsafe { MmapOptions::new().map(&file) } + .with_context(|| format!("Failed to mmap source partition image: {path:?}"))?, + ) + }; + let source = SourcePartition { path, len, mmap }; + + let expected_hash = info.hash.as_deref().context("Source partition hash is missing")?; + verify_hash( + source.bytes(), + Some(expected_hash), + "Source partition", + &self.cancellation_token, + )?; + Ok(source) + } + + fn create_staged_partition( + &self, + update: &PartitionUpdate, + staging_dir: &Path, + ) -> Result<(TempPath, MmapMut)> { + let partition_len = update + .new_partition_info + .as_ref() + .and_then(|info| info.size) + .context("Unable to determine output file size")?; + ensure!(partition_len > 0, "Output partition size must be greater than zero"); - let file = OpenOptions::new() - .create(true) - .create_new(!self.overwrite) - .read(true) - .write(true) - .truncate(true) - .open(&path) - .with_context(|| format!("Unable to open file for writing: {path:?}"))?; - file.set_len(partition_len)?; + let file = tempfile::Builder::new() + .prefix("partition-") + .tempfile_in(staging_dir) + .context("Unable to create staged partition image")?; + file.as_file().set_len(partition_len)?; - let file = unsafe { MmapMut::map_mut(&file) } - .with_context(|| format!("Failed to mmap file: {path:?}"))?; - Ok(file) + let partition = unsafe { MmapMut::map_mut(file.as_file()) } + .context("Failed to mmap staged partition image")?; + Ok((file.into_temp_path(), partition)) } /// Number of threads to use for extraction. By default, this is set to the @@ -280,11 +598,212 @@ impl<'a> ExtractOptions<'a> { self } - /// Verify the input and output partitions. This is enabled by default. + /// Verify operation data and source extent hashes. This is enabled by default. + /// Whole source images and staged output images are always verified. pub fn verify(&mut self, verify: bool) -> &mut Self { self.verify = verify; self } + + /// Set the directory containing base partition images for delta OTA extraction. + pub fn source_dir>(&mut self, source_dir: P) -> &mut Self { + self.source_dir = Some(source_dir.as_ref().to_path_buf()); + self + } + + /// Stop extraction cooperatively when the caller cancels the token. + pub fn cancellation_token(&mut self, cancellation_token: &CancellationToken) -> &mut Self { + self.cancellation_token = cancellation_token.clone(); + self + } +} + +struct SourcePartition { + path: PathBuf, + len: usize, + mmap: Option, +} + +impl SourcePartition { + fn bytes(&self) -> &[u8] { + &self.mmap.as_deref().unwrap_or_default()[..self.len] + } +} + +fn validate_partition_name(name: &str) -> Result<()> { + let mut components = Path::new(name).components(); + let valid = matches!(components.next(), Some(Component::Normal(component)) if component == name) + && components.next().is_none() + && !name.contains('/') + && !name.contains('\\'); + ensure!(valid, "Unsafe partition name: {name:?}"); + Ok(()) +} + +#[allow(deprecated)] +fn validate_operation_type(op_type: Type) -> Result<()> { + match op_type { + Type::Move | Type::Bsdiff => bail!("Deprecated operation is not supported: {op_type:?}"), + Type::Replace + | Type::ReplaceBz + | Type::SourceCopy + | Type::SourceBsdiff + | Type::Zero + | Type::Discard + | Type::ReplaceXz + | Type::Puffdiff + | Type::BrotliBsdiff + | Type::Zucchini + | Type::Lz4diffBsdiff + | Type::Lz4diffPuffdiff + | Type::ReplaceZstd => Ok(()), + } +} + +fn operation_needs_source(op_type: Type) -> bool { + matches!( + op_type, + Type::SourceCopy + | Type::SourceBsdiff + | Type::Puffdiff + | Type::BrotliBsdiff + | Type::Zucchini + | Type::Lz4diffBsdiff + | Type::Lz4diffPuffdiff + ) +} + +fn operation_has_data(op_type: Type) -> bool { + matches!( + op_type, + Type::Replace + | Type::ReplaceBz + | Type::SourceBsdiff + | Type::ReplaceXz + | Type::Puffdiff + | Type::BrotliBsdiff + | Type::Zucchini + | Type::Lz4diffBsdiff + | Type::Lz4diffPuffdiff + | Type::ReplaceZstd + ) +} + +fn extent_ranges( + extents: &[chromeos_update_engine::Extent], + block_size: usize, + partition_len: usize, +) -> Result>> { + extents + .iter() + .map(|extent| { + let start_block = + usize::try_from(extent.start_block.context("start_block not defined in extent")?) + .context("Extent start block is too large")?; + let num_blocks = + usize::try_from(extent.num_blocks.context("num_blocks not defined in extent")?) + .context("Extent block count is too large")?; + let start = start_block.checked_mul(block_size).context("Extent offset overflow")?; + let len = num_blocks.checked_mul(block_size).context("Extent length overflow")?; + let end = start.checked_add(len).context("Extent end overflow")?; + ensure!(end <= partition_len, "Extent exceeds partition size"); + Ok(start..end) + }) + .collect() +} + +fn ranges_len(ranges: &[Range]) -> Result { + ranges.iter().try_fold(0usize, |len, range| { + len.checked_add(range.len()).context("Combined extent length overflow") + }) +} + +fn extract_operation_data<'a>(payload: &'a Payload<'a>, op: &InstallOperation) -> Result<&'a [u8]> { + let offset = usize::try_from(op.data_offset.context("data_offset not defined")?) + .context("Data offset is too large")?; + let len = usize::try_from(op.data_length.context("data_length not defined")?) + .context("Data length is too large")?; + let end = offset.checked_add(len).context("Data range overflow")?; + payload.data.get(offset..end).context("Data range exceeds payload size") +} + +fn verify_hash_chunks<'a>( + chunks: impl IntoIterator, + expected: Option<&[u8]>, + label: &str, + cancellation_token: &CancellationToken, +) -> Result<()> { + cancellation_token.check()?; + let Some(expected) = expected else { + return Ok(()); + }; + let mut context = digest::Context::new(&digest::SHA256); + for chunk in chunks { + cancellation_token.check()?; + context.update(chunk); + } + let actual = context.finish(); + ensure!( + actual.as_ref() == expected, + "{label} hash mismatch: expected {}, got {}", + hex::encode(expected), + hex::encode(actual.as_ref()) + ); + Ok(()) +} + +fn verify_hash( + data: &[u8], + expected: Option<&[u8]>, + label: &str, + cancellation_token: &CancellationToken, +) -> Result<()> { + verify_hash_chunks(data.chunks(VERIFY_CHUNK_SIZE), expected, label, cancellation_token) +} + +fn verify_hash_over_ranges( + data: &[u8], + ranges: &[Range], + expected: Option<&[u8]>, + label: &str, + cancellation_token: &CancellationToken, +) -> Result<()> { + let chunks = ranges.iter().flat_map(|range| data[range.clone()].chunks(VERIFY_CHUNK_SIZE)); + verify_hash_chunks(chunks, expected, label, cancellation_token) +} + +fn paths_alias(source: &Path, destination: &Path) -> Result { + let destination_metadata = match fs::metadata(destination) { + Ok(metadata) => metadata, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false), + Err(error) => { + return Err(error).with_context(|| { + format!("Unable to inspect destination partition image: {destination:?}") + }); + } + }; + let source_metadata = fs::metadata(source) + .with_context(|| format!("Unable to inspect source partition image: {source:?}"))?; + + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt as _; + Ok(source_metadata.dev() == destination_metadata.dev() + && source_metadata.ino() == destination_metadata.ino()) + } + #[cfg(not(unix))] + { + Ok(fs::canonicalize(source)? == fs::canonicalize(destination)?) + } +} + +fn path_entry_exists(path: &Path) -> Result { + match fs::symlink_metadata(path) { + Ok(_) => Ok(true), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), + Err(error) => Err(error) + .with_context(|| format!("Unable to inspect destination partition image: {path:?}")), + } } impl Default for ExtractOptions<'_> { @@ -296,38 +815,33 @@ impl Default for ExtractOptions<'_> { struct Task<'a> { payload: &'a Payload<'a>, block_size: usize, - verify: bool, update: &'a PartitionUpdate, + source: Option, op_idx: AtomicUsize, partition: SyncUnsafeCell, total_ops: usize, total_ops_completed: &'a AtomicUsize, progress_reporter: &'a dyn ProgressReporter, + cancellation_token: &'a CancellationToken, error: &'a OnceLock, } impl Task<'_> { - fn run(&self, ctx: BroadcastContext<'_>) -> Result<()> { + fn run(&self) -> Result<()> { // If an error has already occurred, stop processing the partition. while self.error.get().is_none() { + self.cancellation_token.check()?; let op_idx = self.op_idx.fetch_add(1, Ordering::AcqRel); match self.update.operations.get(op_idx) { Some(op) => { self.run_op(op)?; + self.cancellation_token.check()?; self.increment_progress(); } None => { - // If this is the last thread to fetch an operation for this - // partition, the partition is fully extracted and can now be - // verified. - if op_idx + 1 == self.update.operations.len() + ctx.num_threads() { - self.verify_partition()?; - unsafe { (*self.partition.get()).flush() } - .context("Error while flushing file to disk")?; - }; break; } } @@ -335,7 +849,9 @@ impl Task<'_> { Ok(()) } + #[allow(deprecated)] fn run_op(&self, op: &InstallOperation) -> Result<()> { + self.cancellation_token.check()?; let mut dst_extents = self.extract_dst_extents(op).context("Error extracting dst_extents")?; match Type::try_from(op.r#type).context("Invalid operation")? { @@ -363,9 +879,188 @@ impl Task<'_> { self.run_op_replace(&mut decoder, &mut dst_extents) .context("Error in REPLACE_ZSTD operation") } - Type::Zero => Ok(()), // This is a no-op since the partition is already zeroed - op => bail!("Unimplemented operation: {op:?}"), + Type::SourceCopy => { + let source = self.extract_source_data(op)?; + self.write_exact(&source, &mut dst_extents) + .context("Error in SOURCE_COPY operation") + } + Type::SourceBsdiff => { + let patch = self.extract_data(op).context("Error extracting SOURCE_BSDIFF data")?; + self.run_op_bsdiff(op, patch, "SOURCE_BSDIFF", &mut dst_extents) + } + Type::BrotliBsdiff => { + let patch = self.extract_data(op).context("Error extracting BROTLI_BSDIFF data")?; + self.run_op_bsdiff(op, patch, "BROTLI_BSDIFF", &mut dst_extents) + } + Type::Zucchini => { + let patch = self.extract_data(op).context("Error extracting ZUCCHINI data")?; + self.run_op_zucchini(op, patch, &mut dst_extents) + } + Type::Puffdiff => { + let patch = self.extract_data(op).context("Error extracting PUFFDIFF data")?; + self.run_op_puffdiff(op, patch, &mut dst_extents) + } + Type::Lz4diffBsdiff | Type::Lz4diffPuffdiff => { + let patch = self.extract_data(op).context("Error extracting LZ4DIFF data")?; + self.run_op_lz4diff(op, patch, &mut dst_extents) + } + Type::Zero | Type::Discard => { + for extent in dst_extents { + for chunk in extent.chunks_mut(VERIFY_CHUNK_SIZE) { + self.cancellation_token.check()?; + chunk.fill(0); + } + } + Ok(()) + } + Type::Move | Type::Bsdiff => { + bail!("Deprecated operation is not supported: {:?}", Type::try_from(op.r#type)?) + } + }?; + self.cancellation_token.check() + } + + fn extract_source_data(&self, op: &InstallOperation) -> Result> { + let source = self.source.as_ref().context("Source partition was not opened")?; + let ranges = extent_ranges(&op.src_extents, self.block_size, source.len)?; + let capacity = ranges_len(&ranges)?; + let mut data = Vec::new(); + data.try_reserve_exact(capacity).context("Unable to allocate source extent buffer")?; + for range in ranges { + for chunk in source.bytes()[range].chunks(VERIFY_CHUNK_SIZE) { + self.cancellation_token.check()?; + data.extend_from_slice(chunk); + } } + Ok(data) + } + + fn run_op_bsdiff( + &self, + op: &InstallOperation, + patch: &[u8], + operation: &str, + dst_extents: &mut [&mut [u8]], + ) -> Result<()> { + let source = self.extract_source_data(op)?; + let expected_output_len = dst_extents.iter().try_fold(0usize, |len, extent| { + len.checked_add(extent.len()).context("Combined destination length overflow") + })?; + let output = apply_bsdiff(&source, patch, expected_output_len, self.cancellation_token) + .map_err(|error| { + if is_cancellation(error.as_ref()) { + error + } else { + anyhow::anyhow!("{operation} patch is invalid: {error}") + } + })?; + self.cancellation_token.check()?; + self.write_exact(&output, dst_extents) + } + + fn run_op_zucchini( + &self, + op: &InstallOperation, + compressed_patch: &[u8], + dst_extents: &mut [&mut [u8]], + ) -> Result<()> { + let source_partition = self.source.as_ref().context("Source partition was not opened")?; + let source_ranges = extent_ranges(&op.src_extents, self.block_size, source_partition.len)?; + let source_len = ranges_len(&source_ranges)?; + ensure!( + source_len < zucchini::OFFSET_BOUND, + "ZUCCHINI source exceeds the maximum size of {} bytes", + zucchini::OFFSET_BOUND - 1 + ); + let source = self.extract_source_data(op)?; + let expected_output_len = dst_extents.iter().try_fold(0usize, |len, extent| { + len.checked_add(extent.len()).context("Combined destination length overflow") + })?; + self.cancellation_token.check()?; + let patch = decode_zucchini_patch(compressed_patch, self.cancellation_token)?; + self.cancellation_token.check()?; + let output = zucchini::apply_with_cancel(&source, &patch, expected_output_len, || { + self.cancellation_token.is_cancelled() + }) + .map_err(map_zucchini_apply_error)?; + self.cancellation_token.check()?; + self.write_exact(&output, dst_extents) + } + + fn run_op_puffdiff( + &self, + op: &InstallOperation, + patch: &[u8], + dst_extents: &mut [&mut [u8]], + ) -> Result<()> { + let source_partition = self.source.as_ref().context("Source partition was not opened")?; + let source_ranges = extent_ranges(&op.src_extents, self.block_size, source_partition.len)?; + let source_size = ranges_len(&source_ranges)?; + let destination_size = dst_extents.iter().try_fold(0usize, |size, extent| { + size.checked_add(extent.len()).context("Combined PUFFDIFF destination length overflow") + })?; + ensure!( + source_size < zucchini::OFFSET_BOUND && destination_size < zucchini::OFFSET_BOUND, + "PUFFDIFF raw stream exceeds the maximum size of {} bytes", + zucchini::OFFSET_BOUND - 1 + ); + let source = self.extract_source_data(op)?; + self.cancellation_token.check()?; + let output = puffin::apply(&source, patch, destination_size, self.cancellation_token) + .map_err(|error| { + let message = error.to_string(); + anyhow::Error::new(error).context(format!("PUFFDIFF apply failed: {message}")) + })?; + self.cancellation_token.check()?; + self.write_exact(&output, dst_extents) + } + + fn run_op_lz4diff( + &self, + op: &InstallOperation, + patch: &[u8], + dst_extents: &mut [&mut [u8]], + ) -> Result<()> { + let source_partition = self.source.as_ref().context("Source partition was not opened")?; + let source_ranges = extent_ranges(&op.src_extents, self.block_size, source_partition.len)?; + let source_size = ranges_len(&source_ranges)?; + let destination_size = dst_extents.iter().try_fold(0usize, |size, extent| { + size.checked_add(extent.len()).context("Combined LZ4DIFF destination length overflow") + })?; + let inner = match Type::try_from(op.r#type).context("Invalid operation")? { + Type::Lz4diffBsdiff => lz4diff::InnerPatch::Bsdiff, + Type::Lz4diffPuffdiff => lz4diff::InnerPatch::Puffdiff, + _ => bail!("Invalid LZ4DIFF operation type"), + }; + let parsed = + lz4diff::parse(patch, source_size, destination_size, inner, self.cancellation_token) + .map_err(lz4diff_error)?; + let source = self.extract_source_data(op)?; + let output = + lz4diff::apply(&source, parsed, self.cancellation_token).map_err(lz4diff_error)?; + self.cancellation_token.check()?; + self.write_exact(&output, dst_extents) + } + + fn write_exact(&self, data: &[u8], dst_extents: &mut [&mut [u8]]) -> Result<()> { + let dst_len = dst_extents.iter().try_fold(0usize, |len, extent| { + len.checked_add(extent.len()).context("Combined destination length overflow") + })?; + ensure!( + data.len() == dst_len, + "Patched output size mismatch: expected {dst_len}, got {}", + data.len() + ); + let mut offset = 0usize; + for extent in dst_extents { + for chunk in extent.chunks_mut(VERIFY_CHUNK_SIZE) { + self.cancellation_token.check()?; + let end = offset.checked_add(chunk.len()).context("Destination offset overflow")?; + chunk.copy_from_slice(&data[offset..end]); + offset = end; + } + } + Ok(()) } fn run_op_replace(&self, reader: &mut impl Read, dst_extents: &mut [&mut [u8]]) -> Result<()> { @@ -375,7 +1070,9 @@ impl Task<'_> { for extent in dst_extents { let mut extent = &mut **extent; loop { - match reader.read(extent).context("Failed to write to buffer")? { + self.cancellation_token.check()?; + let chunk_len = extent.len().min(VERIFY_CHUNK_SIZE); + match reader.read(&mut extent[..chunk_len]).context("Failed to write to buffer")? { 0 => break, n => { bytes_written += n; @@ -384,109 +1081,255 @@ impl Task<'_> { } } } - ensure!(reader.read(&mut [0])? == 0, "Read fewer bytes than expected"); + self.cancellation_token.check()?; + ensure!(reader.read(&mut [0])? == 0, "Decoded data exceeds destination extents"); // Align number of bytes written to block size. The formula for alignment is: // ((operand + alignment - 1) / alignment) * alignment. - let bytes_written_aligned = bytes_written.div_ceil(self.block_size).mul(self.block_size); + let bytes_written_aligned = bytes_written + .checked_add(self.block_size - 1) + .context("Decoded data length overflow")? + / self.block_size + * self.block_size; ensure!(bytes_written_aligned == dst_len, "More dst blocks than data, even with padding"); Ok(()) } - fn extract_dst_extents(&self, op: &InstallOperation) -> Result> { + #[allow(clippy::mut_from_ref)] + fn extract_dst_extents(&self, op: &InstallOperation) -> Result> { let partition = unsafe { (*self.partition.get()).as_mut_ptr() }; let partition_len = unsafe { (&(*self.partition.get())).len() }; + let ranges = extent_ranges(&op.dst_extents, self.block_size, partition_len)?; - op.dst_extents - .iter() - .map(|extent| { - let start_block = - extent.start_block.context("start_block not defined in extent")? as usize; - let num_blocks = - extent.num_blocks.context("num_blocks not defined in extent")? as usize; + // validate_partition rejects overlapping destination ranges before any task starts. + // This guarantees that parallel workers never create aliases into the mutable map. + Ok(ranges + .into_iter() + .map(|range| unsafe { + slice::from_raw_parts_mut(partition.add(range.start), range.len()) + }) + .collect()) + } - let partition_offset = start_block * self.block_size; - let extent_len = num_blocks * self.block_size; + fn extract_data<'a>(&'a self, op: &InstallOperation) -> Result<&'a [u8]> { + extract_operation_data(self.payload, op) + } - ensure!( - partition_offset + extent_len <= partition_len, - "Extent exceeds partition size" - ); - let extent = unsafe { - slice::from_raw_parts_mut(partition.add(partition_offset), extent_len) - }; + fn increment_progress(&self) { + increment_progress(self.total_ops, self.total_ops_completed, self.progress_reporter); + } +} - Ok(extent) - }) - .collect() +fn lz4diff_error(error: Error) -> Error { + if is_cancellation(error.as_ref()) { + error + } else { + let message = error.to_string(); + error.context(format!("Error in LZ4DIFF operation: {message}")) } +} - fn extract_data<'a>(&'a self, op: &InstallOperation) -> Result<&'a [u8]> { - let data_len = op.data_length.context("data_length not defined")? as usize; - let data = { - let offset = op.data_offset.context("data_offset not defined")? as usize; - self.payload - .data - .get(offset..offset + data_len) - .context("Data offset exceeds payload size")? - }; - self.verify_op(op, data)?; - Ok(data) +fn map_zucchini_apply_error(error: zucchini::Error) -> Error { + if error.status() == zucchini::Status::Cancelled { + Error::new(ExtractionCancelled) + } else { + anyhow::anyhow!("ZUCCHINI apply failed: {error}") } +} - fn verify_op(&self, op: &InstallOperation, data: &[u8]) -> Result<()> { - if !self.verify { - return Ok(()); - } - let Some(exp_hash) = &op.data_sha256_hash else { - return Ok(()); - }; +pub(crate) fn decode_zucchini_patch( + compressed_patch: &[u8], + cancellation_token: &CancellationToken, +) -> Result> { + cancellation_token.check()?; + let mut decoder = BrotliDecoder::new(compressed_patch, BROTLI_DECODE_CHUNK_SIZE); + let mut patch = Vec::new(); + let mut chunk = [0u8; BROTLI_DECODE_CHUNK_SIZE]; + let maximum_size = zucchini::OFFSET_BOUND - 1; - let got_hash = digest::digest(&digest::SHA256, data); - let got_hash = got_hash.as_ref(); + loop { + cancellation_token.check()?; + let remaining = maximum_size + .checked_sub(patch.len()) + .context("Decoded ZUCCHINI patch length exceeded its bound")?; + let read_len = chunk.len().min(remaining.saturating_add(1)); + let count = decoder + .read(&mut chunk[..read_len]) + .context("Unable to decode ZUCCHINI Brotli patch")?; + if count == 0 { + cancellation_token.check()?; + // The decoder reports buffered trailing input only on a read after + // EOF. The source check also catches trailing bytes it did not read. + let mut trailing = [0u8; 1]; + let trailing_count = + decoder.read(&mut trailing).context("Unable to decode ZUCCHINI Brotli patch")?; + ensure!( + trailing_count == 0 && decoder.get_ref().is_empty(), + "ZUCCHINI Brotli patch contains trailing data" + ); + break; + } ensure!( - got_hash == exp_hash, - "Input verification failed: hash mismatch: expected {}, got {got_hash:?}", - hex::encode(exp_hash) + count <= remaining, + "Decoded ZUCCHINI patch exceeds the maximum size of {maximum_size} bytes" ); - Ok(()) + let new_len = + patch.len().checked_add(count).context("Decoded ZUCCHINI patch length overflow")?; + patch.try_reserve(count).context("Unable to allocate decoded ZUCCHINI patch buffer")?; + patch.extend_from_slice(&chunk[..count]); + debug_assert_eq!(patch.len(), new_len); } - fn verify_partition(&self) -> Result<()> { - if !self.verify { - return Ok(()); - } + Ok(patch) +} - let Some(exp_hash) = - self.update.new_partition_info.as_ref().and_then(|info| info.hash.as_ref()) - else { - return Ok(()); - }; +fn verify_partition( + update: &PartitionUpdate, + partition: &[u8], + total_ops: usize, + total_ops_completed: &AtomicUsize, + progress_reporter: &dyn ProgressReporter, + cancellation_token: &CancellationToken, +) -> Result<()> { + let exp_hash = update + .new_partition_info + .as_ref() + .and_then(|info| info.hash.as_ref()) + .context("Unable to determine output partition hash")?; - let mut digest = digest::Context::new(&digest::SHA256); - for chunk in unsafe { (*self.partition.get()).chunks(VERIFY_CHUNK_SIZE) } { - digest.update(chunk); - self.increment_progress(); - } + let mut context = digest::Context::new(&digest::SHA256); + for chunk in partition.chunks(VERIFY_CHUNK_SIZE) { + cancellation_token.check()?; + context.update(chunk); + increment_progress(total_ops, total_ops_completed, progress_reporter); + } + + cancellation_token.check()?; + let got_hash = context.finish(); + ensure!( + got_hash.as_ref() == exp_hash, + "Output verification failed: hash mismatch: expected {}, got {}", + hex::encode(exp_hash), + hex::encode(got_hash.as_ref()) + ); + Ok(()) +} + +fn increment_progress( + total_ops: usize, + total_ops_completed: &AtomicUsize, + progress_reporter: &dyn ProgressReporter, +) { + let total_ops_completed = total_ops_completed.fetch_add(1, Ordering::Relaxed) + 1; + if total_ops_completed % 16 == 0 { + let progress = total_ops_completed as f64 / total_ops as f64; + progress_reporter.report_progress(progress); + } +} + +#[inline] +fn offtin(buf: [u8; 8]) -> i64 { + let y = i64::from_le_bytes(buf); + if 0 == y & (1 << 63) { y } else { -(y & !(1 << 63)) } +} + +pub(crate) fn apply_bsdiff( + source: &[u8], + patch: &[u8], + expected_output_len: usize, + cancellation_token: &CancellationToken, +) -> Result> { + bsdiff::validate_output_len(patch, expected_output_len)?; + cancellation_token.check()?; + puffin::validate_bsdiff_resources(patch, expected_output_len, cancellation_token)?; + cancellation_token.check()?; + + let (new_size, control_data, diff_data, extra_data) = bsdiff::parse_header(patch)?; + let new_size = usize::try_from(new_size).context("Patch new_size is negative or overflows")?; + ensure!( + new_size == expected_output_len, + "Patch output length mismatch: expected {expected_output_len}, got {new_size}" + ); + + let mut output = Vec::new(); + output + .try_reserve_exact(expected_output_len) + .context("Unable to allocate patched output buffer")?; + + ensure!(control_data.len() % 24 == 0, "Invalid control data length (not a multiple of 24)"); + let mut oldpos: i64 = 0; + let mut diff_pos: usize = 0; + let mut extra_pos: usize = 0; + + for chunk in control_data.chunks_exact(24) { + cancellation_token.check()?; + let add_len = offtin(chunk[0..8].try_into().unwrap()); + let copy_len = offtin(chunk[8..16].try_into().unwrap()); + let seek_amount = offtin(chunk[16..24].try_into().unwrap()); - let got_hash = digest.finish(); - let got_hash = got_hash.as_ref(); ensure!( - got_hash == exp_hash, - "Output verification failed: hash mismatch: expected {}, got {got_hash:?}", - hex::encode(exp_hash) + add_len >= 0 && copy_len >= 0, + "Negative length in control tuple: add={add_len}, copy={copy_len}" ); - Ok(()) - } + let add_len = usize::try_from(add_len).context("Control tuple add_len overflow")?; + let copy_len = usize::try_from(copy_len).context("Control tuple copy_len overflow")?; - fn increment_progress(&self) { - let total_ops_completed = self.total_ops_completed.fetch_add(1, Ordering::Relaxed) + 1; - if total_ops_completed % 16 == 0 { - let progress = total_ops_completed as f64 / self.total_ops as f64; - self.progress_reporter.report_progress(progress); + let pending_len = add_len.checked_add(copy_len).context("Control tuple length overflow")?; + let target_len = + output.len().checked_add(pending_len).context("Patched output length overflow")?; + ensure!( + target_len <= expected_output_len, + "Control tuple would exceed expected output size" + ); + + if add_len > 0 { + let diff_end = diff_pos.checked_add(add_len).context("Diff stream offset overflow")?; + let diff_chunk = diff_data.get(diff_pos..diff_end).context("Diff stream exhausted")?; + for (i, &diff_byte) in diff_chunk.iter().enumerate() { + let offset = i64::try_from(i).context("Diff offset overflow")?; + let src_pos = oldpos.checked_add(offset).context("Source position overflow")?; + let old_byte = match usize::try_from(src_pos) { + Ok(pos) if pos < source.len() => source[pos], + _ => 0u8, + }; + output.push(old_byte.wrapping_add(diff_byte)); + } + let add_i64 = i64::try_from(add_len).context("Source position overflow")?; + oldpos = oldpos.checked_add(add_i64).context("Source position overflow")?; + diff_pos = diff_end; + } + + if copy_len > 0 { + let extra_end = + extra_pos.checked_add(copy_len).context("Extra stream offset overflow")?; + let extra_chunk = + extra_data.get(extra_pos..extra_end).context("Extra stream exhausted")?; + output.extend_from_slice(extra_chunk); + extra_pos = extra_end; } + + oldpos = oldpos.checked_add(seek_amount).context("Source position overflow")?; } + + cancellation_token.check()?; + ensure!( + output.len() == expected_output_len, + "Patched output size mismatch: expected {expected_output_len}, got {}", + output.len() + ); + ensure!( + diff_pos == diff_data.len(), + "Diff stream not fully consumed: used {diff_pos}/{}", + diff_data.len() + ); + ensure!( + extra_pos == extra_data.len(), + "Extra stream not fully consumed: used {extra_pos}/{}", + extra_data.len() + ); + + Ok(output) } pub trait ProgressReporter: Sync { diff --git a/src/lz4.rs b/src/lz4.rs new file mode 100644 index 0000000..5f60c52 --- /dev/null +++ b/src/lz4.rs @@ -0,0 +1,341 @@ +//! LZ4 block operations for Android `lz4diff`, backed by the pinned `lz4-sys` +//! crate. The crate bundles liblz4 1.10.0 C sources that are byte-identical +//! (SHA-256 verified) to AOSP `platform/external/lz4` commit +//! `734e07032602e9a72fcc9701028b0aee45147fcd`, so compressed output matches +//! the AOSP reference byte for byte. See `licenses/README.md`. + +// Force the pinned `lz4-sys` crate into the crate graph: its rlib carries the +// `cargo:rustc-link-lib=static=lz4` metadata that links the bundled liblz4 +// archive, which the direct extern declarations below resolve against. +use lz4_sys as _; + +use std::error; +use std::ffi::{c_char, c_int, c_void}; +use std::fmt; + +const MAX_INPUT_SIZE: usize = 0x7e00_0000; +pub const HC_LEVEL_MIN: i32 = 2; +pub const HC_LEVEL_MAX: i32 = 12; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Status { + InvalidArgument, + CompressionFailed, + DecompressionFailed, + AllocationFailure, + CompressionDivergence, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct Error { + status: Status, + message: String, +} + +impl Error { + pub fn status(&self) -> Status { + self.status + } +} + +impl fmt::Display for Error { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.message) + } +} + +impl error::Error for Error {} + +pub type Result = std::result::Result; + +/// Decompresses one LZ4 block into a newly allocated, exact-size buffer. +pub fn decompress_safe_partial(input: &[u8], output_size: usize) -> Result> { + validate_sizes(input.len(), output_size)?; + if input.len() >= output_size { + return Err(invalid_argument("stored LZ4 block must be smaller than its raw block")); + } + let mut output = allocate_output(output_size, "decompression")?; + let result = decompress_native(input, &mut output)?; + if result.output_size != output_size { + return Err(Error { + status: Status::DecompressionFailed, + message: format!( + "LZ4 decompressed {} bytes, expected {output_size}", + result.output_size + ), + }); + } + Ok(output) +} + +/// Compresses one complete block with LZ4 acceleration 1. +/// +/// The returned buffer always has `output_size` bytes. Short compressed output +/// receives leading zeros when `zero_padding` is true and trailing zeros otherwise. +pub fn compress_dest_size(input: &[u8], output_size: usize, zero_padding: bool) -> Result> { + compress(input, Some(input.len()), output_size, zero_padding, None) +} + +/// Compresses one complete block with the metadata-selected LZ4HC level. +pub fn compress_hc_dest_size( + input: &[u8], + output_size: usize, + compression_level: i32, + zero_padding: bool, +) -> Result> { + if !(HC_LEVEL_MIN..=HC_LEVEL_MAX).contains(&compression_level) { + return Err(invalid_argument(format!( + "LZ4HC compression level must be in {HC_LEVEL_MIN}..={HC_LEVEL_MAX}" + ))); + } + compress(input, Some(input.len()), output_size, zero_padding, Some(compression_level)) +} + +pub(crate) fn compress_dest_size_partial( + input: &[u8], + raw_block_size: usize, + output_size: usize, + compression_level: Option, + zero_padding: bool, +) -> Result> { + if let Some(level) = compression_level + && !(HC_LEVEL_MIN..=HC_LEVEL_MAX).contains(&level) + { + return Err(invalid_argument(format!( + "LZ4HC compression level must be in {HC_LEVEL_MIN}..={HC_LEVEL_MAX}" + ))); + } + if raw_block_size == 0 || raw_block_size > input.len() { + return Err(invalid_argument("LZ4 source consumption size is invalid")); + } + if output_size >= raw_block_size { + return Err(invalid_argument("stored LZ4 block must be smaller than its raw block")); + } + compress(input, None, output_size, zero_padding, compression_level) +} + +fn compress( + input: &[u8], + expected_source_size: Option, + output_size: usize, + zero_padding: bool, + compression_level: Option, +) -> Result> { + validate_sizes(input.len(), output_size)?; + if let Some(expected) = expected_source_size { + if expected == 0 || expected > input.len() { + return Err(invalid_argument("LZ4 source consumption size is invalid")); + } + if output_size >= expected { + return Err(invalid_argument("stored LZ4 block must be smaller than its raw block")); + } + } else if output_size >= input.len() { + return Err(invalid_argument("stored LZ4 block must be smaller than its raw block")); + } + let mut output = allocate_output(output_size, "compression")?; + let result = compress_native(input, &mut output, compression_level)?; + if result.source_size == 0 || result.source_size > input.len() { + return Err(Error { + status: Status::CompressionFailed, + message: "LZ4 compression consumed invalid source bytes".into(), + }); + } + if let Some(expected) = expected_source_size { + if result.source_size != expected { + return Err(Error { + status: Status::CompressionDivergence, + message: format!( + "LZ4 compression consumed {} of {} source bytes", + result.source_size, expected + ), + }); + } + } + if result.output_size == 0 || result.output_size > output_size { + return Err(Error { + status: Status::CompressionDivergence, + message: format!( + "LZ4 compression produced {} bytes for capacity {output_size}", + result.output_size + ), + }); + } + + let padding = output_size - result.output_size; + if padding > 0 { + if zero_padding { + output.copy_within(0..result.output_size, padding); + output[..padding].fill(0); + } else { + output[result.output_size..].fill(0); + } + } + Ok(output) +} + +fn validate_sizes(input_size: usize, output_size: usize) -> Result<()> { + if input_size == 0 || output_size == 0 { + return Err(invalid_argument("LZ4 block sizes must be greater than zero")); + } + if input_size > MAX_INPUT_SIZE || output_size > MAX_INPUT_SIZE { + return Err(invalid_argument("LZ4 block exceeds LZ4_MAX_INPUT_SIZE")); + } + i32::try_from(input_size).map_err(|_| invalid_argument("LZ4 input size exceeds int"))?; + i32::try_from(output_size).map_err(|_| invalid_argument("LZ4 output size exceeds int"))?; + Ok(()) +} + +fn allocate_output(size: usize, operation: &str) -> Result> { + let mut output = Vec::new(); + output.try_reserve_exact(size).map_err(|error| Error { + status: Status::AllocationFailure, + message: format!("unable to allocate LZ4 {operation} output: {error}"), + })?; + output.resize(size, 0); + Ok(output) +} + +fn invalid_argument(message: impl Into) -> Error { + Error { status: Status::InvalidArgument, message: message.into() } +} + +#[derive(Clone, Copy)] +struct NativeResult { + source_size: usize, + output_size: usize, +} + +fn decompress_native(input: &[u8], output: &mut [u8]) -> Result { + // SAFETY: The immutable input and mutable output slices cannot alias. They + // remain alive for the call, and their checked lengths match the pointers. + let written = unsafe { + LZ4_decompress_safe_partial( + input.as_ptr().cast::(), + output.as_mut_ptr().cast::(), + to_c_int(input.len())?, + to_c_int(output.len())?, + to_c_int(output.len())?, + ) + }; + if written < 0 { + return Err(Error { + status: Status::DecompressionFailed, + message: "LZ4 decompression failed".into(), + }); + } + Ok(NativeResult { source_size: input.len(), output_size: written as usize }) +} + +fn compress_native( + input: &[u8], + output: &mut [u8], + compression_level: Option, +) -> Result { + let mut consumed = to_c_int(input.len())?; + let capacity = to_c_int(output.len())?; + // SAFETY: The immutable input and mutable output slices cannot alias. They + // remain alive for the call, and their checked lengths match the pointers. + // `consumed` is a valid in-out parameter initialized to the input length. + let written = unsafe { + if let Some(level) = compression_level { + let stream = LZ4_createStreamHC(); + if stream.is_null() { + return Err(Error { + status: Status::AllocationFailure, + message: "unable to allocate the LZ4HC stream state".into(), + }); + } + let written = LZ4_compress_HC_destSize( + stream, + input.as_ptr().cast::(), + output.as_mut_ptr().cast::(), + &mut consumed, + capacity, + level, + ); + LZ4_freeStreamHC(stream); + written + } else { + LZ4_compress_destSize( + input.as_ptr().cast::(), + output.as_mut_ptr().cast::(), + &mut consumed, + capacity, + ) + } + }; + if written <= 0 { + return Err(Error { + status: Status::CompressionFailed, + message: "LZ4 compression failed".into(), + }); + } + Ok(NativeResult { source_size: consumed as usize, output_size: written as usize }) +} + +fn to_c_int(size: usize) -> Result { + c_int::try_from(size).map_err(|_| invalid_argument("LZ4 size exceeds int")) +} + +// Direct declarations for the liblz4 block API: the pinned `lz4-sys` crate +// statically links the bundled liblz4 archive but its bindings omit these +// symbols. Signatures match `lz4.h`/`lz4hc.h` 1.10.0. +unsafe extern "C" { + fn LZ4_decompress_safe_partial( + source: *const c_char, + dest: *mut c_char, + compressed_size: c_int, + target_output_size: c_int, + max_decompressed_size: c_int, + ) -> c_int; + fn LZ4_compress_destSize( + source: *const c_char, + dest: *mut c_char, + source_size: *mut c_int, + target_dest_size: c_int, + ) -> c_int; + fn LZ4_createStreamHC() -> *mut c_void; + fn LZ4_freeStreamHC(stream: *mut c_void) -> c_int; + fn LZ4_compress_HC_destSize( + stream: *mut c_void, + source: *const c_char, + dest: *mut c_char, + source_size: *mut c_int, + target_dest_size: c_int, + compression_level: c_int, + ) -> c_int; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn compress_and_decompress_round_trip() { + let raw = (0..2048).map(|i| (i % 37) as u8).collect::>(); + let compressed = compress_dest_size(&raw, 512, false).unwrap(); + let decompressed = decompress_safe_partial(&compressed, raw.len()).unwrap(); + assert_eq!(decompressed, raw); + } + + #[test] + fn zero_padding_round_trip() { + let raw = vec![0x5a; 1024]; + let compressed = compress_dest_size(&raw, 64, true).unwrap(); + let start = compressed.iter().position(|byte| *byte != 0).unwrap(); + let decompressed = decompress_safe_partial(&compressed[start..], raw.len()).unwrap(); + assert_eq!(decompressed, raw); + } + + #[test] + fn invalid_sizes_return_invalid_argument() { + assert_eq!( + compress_dest_size(&[], 32, false).unwrap_err().status(), + Status::InvalidArgument + ); + assert_eq!( + decompress_safe_partial(&[1, 2, 3], 0).unwrap_err().status(), + Status::InvalidArgument + ); + } +} diff --git a/src/lz4diff.rs b/src/lz4diff.rs new file mode 100644 index 0000000..754ce3f --- /dev/null +++ b/src/lz4diff.rs @@ -0,0 +1,669 @@ +// Copyright 2026 otadump contributors +// SPDX-License-Identifier: MIT + +use anyhow::{Context as _, Result, bail, ensure}; +use prost::Message as _; +use ring::digest; + +use crate::{CancellationToken, bsdiff, is_cancellation, lz4, puffin, zucchini}; + +const MAGIC: &[u8; 7] = b"LZ4DIFF"; +const FRAMING_SIZE: usize = 16; +const VERSION: u32 = 1; +const EROFS_BLOCK_SIZE: usize = 4096; +const PROTOBUF_RECURSION_LIMIT: usize = 100; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum InnerPatch { + Bsdiff, + Puffdiff, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Algorithm { + Uncompressed, + Lz4, + Lz4hc, +} + +struct Block<'a> { + uncompressed_offset: usize, + uncompressed_length: usize, + compressed_length: usize, + sha256_hash: &'a [u8], + postfix_bspatch: &'a [u8], +} + +struct CompressionInfo<'a> { + algorithm: Algorithm, + level: i32, + blocks: Vec>, + zero_padding_enabled: bool, + uncompressed_size: usize, + physical_size: usize, +} + +pub(crate) struct Patch<'a> { + source: CompressionInfo<'a>, + destination: CompressionInfo<'a>, + inner_patch: &'a [u8], + inner_type: InnerPatch, +} + +#[derive(Clone, Copy, PartialEq, Eq, prost::Message)] +struct CompressionAlgorithmProto { + #[prost(enumeration = "AlgorithmProto", tag = "1")] + r#type: i32, + #[prost(int32, tag = "2")] + level: i32, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, prost::Enumeration)] +#[repr(i32)] +enum AlgorithmProto { + Uncompressed = 0, + Lz4 = 1, + Lz4hc = 2, +} + +#[derive(Clone, PartialEq, Eq, prost::Message)] +struct BlockProto { + #[prost(uint64, tag = "1")] + uncompressed_offset: u64, + #[prost(uint64, tag = "2")] + uncompressed_length: u64, + #[prost(uint64, tag = "3")] + compressed_length: u64, +} + +#[derive(Clone, PartialEq, prost::Message)] +struct CompressionInfoProto { + #[prost(message, optional, tag = "1")] + algo: Option, + #[prost(message, repeated, tag = "2")] + block_info: Vec, + #[prost(bool, tag = "3")] + zero_padding_enabled: bool, +} + +#[derive(Clone, PartialEq, prost::Message)] +struct HeaderProto { + #[prost(message, optional, tag = "1")] + src_info: Option, + #[prost(message, optional, tag = "2")] + dst_info: Option, + #[prost(enumeration = "InnerPatchProto", tag = "3")] + inner_type: i32, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, prost::Enumeration)] +#[repr(i32)] +enum InnerPatchProto { + Bsdiff = 0, + Puffdiff = 1, +} + +struct HeaderLayout<'a> { + source_seen: bool, + destination_seen: bool, + source_blocks: Vec<&'a [u8]>, + destination_blocks: Vec<&'a [u8]>, +} + +struct ProtoCursor<'a> { + bytes: &'a [u8], + position: usize, +} + +impl<'a> ProtoCursor<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, position: 0 } + } + + fn at_end(&self) -> bool { + self.position == self.bytes.len() + } + + fn varint(&mut self) -> Result { + let mut value = 0u64; + for index in 0..10 { + let byte = + *self.bytes.get(self.position).context("truncated LZ4DIFF protobuf varint")?; + self.position += 1; + if index == 9 && byte > 1 { + bail!("LZ4DIFF protobuf varint overflows u64"); + } + value |= u64::from(byte & 0x7f) << (index * 7); + if byte & 0x80 == 0 { + return Ok(value); + } + } + bail!("LZ4DIFF protobuf varint is too long") + } + + fn key(&mut self) -> Result<(u64, u8)> { + let key = self.varint()?; + ensure!(key >> 3 != 0, "LZ4DIFF protobuf field number is zero"); + Ok((key >> 3, (key & 7) as u8)) + } + + fn bytes(&mut self) -> Result<&'a [u8]> { + let length = usize::try_from(self.varint()?) + .context("LZ4DIFF protobuf field length is too large")?; + let end = + self.position.checked_add(length).context("LZ4DIFF protobuf field range overflow")?; + let bytes = + self.bytes.get(self.position..end).context("LZ4DIFF protobuf field overruns header")?; + self.position = end; + Ok(bytes) + } + + fn advance(&mut self, length: usize) -> Result<()> { + let end = + self.position.checked_add(length).context("LZ4DIFF protobuf field range overflow")?; + ensure!(end <= self.bytes.len(), "LZ4DIFF protobuf field overruns header"); + self.position = end; + Ok(()) + } + + fn skip(&mut self, field: u64, wire_type: u8) -> Result<()> { + self.skip_with_depth(field, wire_type, 0) + } + + fn skip_with_depth(&mut self, field: u64, wire_type: u8, depth: usize) -> Result<()> { + match wire_type { + 0 => { + self.varint()?; + } + 1 => self.advance(8)?, + 2 => { + self.bytes()?; + } + 3 => { + ensure!( + depth < PROTOBUF_RECURSION_LIMIT, + "LZ4DIFF protobuf recursion limit exceeded" + ); + loop { + let (nested_field, nested_wire_type) = self.key()?; + if nested_wire_type == 4 { + ensure!(nested_field == field, "mismatched LZ4DIFF protobuf end group"); + break; + } + self.skip_with_depth(nested_field, nested_wire_type, depth + 1)?; + } + } + 4 => bail!("unexpected LZ4DIFF protobuf end group"), + 5 => self.advance(4)?, + _ => bail!("unsupported LZ4DIFF protobuf wire type {wire_type}"), + } + Ok(()) + } +} + +fn scan_header(bytes: &[u8]) -> Result> { + let mut layout = HeaderLayout { + source_seen: false, + destination_seen: false, + source_blocks: Vec::new(), + destination_blocks: Vec::new(), + }; + let mut cursor = ProtoCursor::new(bytes); + while !cursor.at_end() { + let (field, wire_type) = cursor.key()?; + match (field, wire_type) { + (1, 2) => { + layout.source_seen = true; + scan_compression_info(cursor.bytes()?, &mut layout.source_blocks)?; + } + (2, 2) => { + layout.destination_seen = true; + scan_compression_info(cursor.bytes()?, &mut layout.destination_blocks)?; + } + _ => cursor.skip(field, wire_type)?, + } + } + Ok(layout) +} + +fn scan_compression_info<'a>(bytes: &'a [u8], blocks: &mut Vec<&'a [u8]>) -> Result<()> { + let mut cursor = ProtoCursor::new(bytes); + while !cursor.at_end() { + let (field, wire_type) = cursor.key()?; + if (field, wire_type) == (2, 2) { + let block = cursor.bytes()?; + blocks.try_reserve(1).context("unable to allocate LZ4DIFF block preflight")?; + blocks.push(block); + } else { + cursor.skip(field, wire_type)?; + } + } + Ok(()) +} + +fn scan_block_blobs(bytes: &[u8]) -> Result<(&[u8], &[u8])> { + let mut sha256_hash = &[][..]; + let mut postfix_bspatch = &[][..]; + let mut cursor = ProtoCursor::new(bytes); + while !cursor.at_end() { + let (field, wire_type) = cursor.key()?; + match (field, wire_type) { + (4, 2) => sha256_hash = cursor.bytes()?, + (5, 2) => postfix_bspatch = cursor.bytes()?, + _ => cursor.skip(field, wire_type)?, + } + } + Ok((sha256_hash, postfix_bspatch)) +} + +fn preallocated_info(block_count: usize) -> Result { + let mut block_info = Vec::new(); + block_info + .try_reserve_exact(block_count) + .context("unable to allocate LZ4DIFF protobuf block metadata")?; + Ok(CompressionInfoProto { algo: None, block_info, zero_padding_enabled: false }) +} + +pub(crate) fn apply( + source: &[u8], + patch: Patch<'_>, + cancellation_token: &CancellationToken, +) -> Result> { + cancellation_token.check()?; + ensure!( + source.len() == patch.source.physical_size, + "LZ4DIFF source size changed after validation" + ); + let uncompressed_source = decompress_source(source, &patch.source, cancellation_token)?; + cancellation_token.check()?; + let uncompressed_destination = match patch.inner_type { + InnerPatch::Bsdiff => apply_bsdiff( + &uncompressed_source, + patch.inner_patch, + patch.destination.uncompressed_size, + "inner BSDIFF", + cancellation_token, + )?, + InnerPatch::Puffdiff => puffin::apply( + &uncompressed_source, + patch.inner_patch, + patch.destination.uncompressed_size, + cancellation_token, + ) + .map_err(|error| preserve_puffin_error(error, "inner PUFFDIFF patch is invalid"))?, + }; + ensure!( + uncompressed_destination.len() == patch.destination.uncompressed_size, + "LZ4DIFF inner output size mismatch: expected {}, got {}", + patch.destination.uncompressed_size, + uncompressed_destination.len() + ); + cancellation_token.check()?; + compress_destination(&uncompressed_destination, &patch.destination, cancellation_token) +} + +pub(crate) fn parse<'a>( + patch_data: &'a [u8], + source_size: usize, + destination_size: usize, + expected_inner: InnerPatch, + cancellation_token: &CancellationToken, +) -> Result> { + ensure_below_bound(patch_data.len(), "patch")?; + let framing = patch_data.get(..FRAMING_SIZE).context("LZ4DIFF framing is truncated")?; + ensure!(&framing[..MAGIC.len()] == MAGIC, "invalid LZ4DIFF magic"); + let version = + u32::from_be_bytes(framing[7..11].try_into().context("invalid LZ4DIFF version field")?); + ensure!(version == VERSION, "unsupported LZ4DIFF version {version}"); + let header_size = usize::try_from(u32::from_be_bytes( + framing[11..15].try_into().context("invalid LZ4DIFF header size field")?, + )) + .context("LZ4DIFF protobuf header is too large")?; + ensure!(framing[15] == 0, "LZ4DIFF framing padding byte is not zero"); + ensure_below_bound(header_size, "protobuf header")?; + let header_end = + FRAMING_SIZE.checked_add(header_size).context("LZ4DIFF protobuf header range overflow")?; + let header_bytes = patch_data + .get(FRAMING_SIZE..header_end) + .context("LZ4DIFF protobuf header overruns patch")?; + let layout = scan_header(header_bytes)?; + ensure!(layout.source_seen, "LZ4DIFF source compression info is missing"); + ensure!(layout.destination_seen, "LZ4DIFF destination compression info is missing"); + let mut header = HeaderProto { + src_info: Some(preallocated_info(layout.source_blocks.len())?), + dst_info: Some(preallocated_info(layout.destination_blocks.len())?), + inner_type: 0, + }; + header.merge(header_bytes).context("invalid LZ4DIFF protobuf header")?; + let inner = InnerPatchProto::try_from(header.inner_type) + .map_err(|_| anyhow::anyhow!("unknown LZ4DIFF inner patch type {}", header.inner_type))?; + let actual_inner = match inner { + InnerPatchProto::Bsdiff => InnerPatch::Bsdiff, + InnerPatchProto::Puffdiff => InnerPatch::Puffdiff, + }; + ensure!( + actual_inner == expected_inner, + "LZ4DIFF operation type does not match inner patch type" + ); + let source = validate_info( + header.src_info.context("LZ4DIFF source compression info is missing")?, + &layout.source_blocks, + "source", + true, + cancellation_token, + )?; + let destination = validate_info( + header.dst_info.context("LZ4DIFF destination compression info is missing")?, + &layout.destination_blocks, + "destination", + false, + cancellation_token, + )?; + ensure!( + source.physical_size == source_size, + "LZ4DIFF source physical size mismatch: metadata describes {}, operation provides {source_size}", + source.physical_size + ); + ensure!( + destination.physical_size == destination_size, + "LZ4DIFF destination physical size mismatch: metadata describes {}, operation provides {destination_size}", + destination.physical_size + ); + let inner_patch = patch_data.get(header_end..).context("invalid LZ4DIFF inner patch offset")?; + ensure!(!inner_patch.is_empty(), "LZ4DIFF inner patch is empty"); + ensure_below_bound(inner_patch.len(), "inner patch")?; + Ok(Patch { source, destination, inner_patch, inner_type: actual_inner }) +} + +fn validate_info<'a>( + info: CompressionInfoProto, + block_messages: &[&'a [u8]], + label: &str, + is_source: bool, + cancellation_token: &CancellationToken, +) -> Result> { + cancellation_token.check()?; + let algorithm = info.algo.context(format!("LZ4DIFF {label} algorithm is missing"))?; + let algorithm_type = AlgorithmProto::try_from(algorithm.r#type).map_err(|_| { + anyhow::anyhow!("unknown LZ4DIFF {label} compression algorithm {}", algorithm.r#type) + })?; + let algorithm_type = match algorithm_type { + AlgorithmProto::Uncompressed => { + ensure!(algorithm.level == 0, "LZ4DIFF {label} uncompressed level must be zero"); + Algorithm::Uncompressed + } + AlgorithmProto::Lz4 => { + ensure!(algorithm.level == 0, "LZ4DIFF {label} LZ4 level must be zero"); + Algorithm::Lz4 + } + AlgorithmProto::Lz4hc => { + ensure!( + (lz4::HC_LEVEL_MIN..=lz4::HC_LEVEL_MAX).contains(&algorithm.level), + "LZ4DIFF {label} LZ4HC level must be in {}..={}", + lz4::HC_LEVEL_MIN, + lz4::HC_LEVEL_MAX + ); + Algorithm::Lz4hc + } + }; + ensure!(!info.block_info.is_empty(), "LZ4DIFF {label} block list is empty"); + if info.zero_padding_enabled { + ensure!( + algorithm_type != Algorithm::Uncompressed, + "LZ4DIFF {label} zero padding requires LZ4 compression" + ); + } + + let mut blocks = Vec::new(); + blocks + .try_reserve_exact(info.block_info.len()) + .context("unable to allocate LZ4DIFF block metadata")?; + let mut uncompressed_size = 0usize; + let mut physical_size = 0usize; + let mut has_compressed_block = false; + ensure!( + info.block_info.len() == block_messages.len(), + "LZ4DIFF {label} block preflight count mismatch" + ); + for (index, (block, block_message)) in + info.block_info.into_iter().zip(block_messages).enumerate() + { + cancellation_token.check()?; + let (sha256_hash, postfix_bspatch) = scan_block_blobs(block_message)?; + let uncompressed_offset = to_size(block.uncompressed_offset, label, index, "offset")?; + let uncompressed_length = to_size(block.uncompressed_length, label, index, "length")?; + let compressed_length = to_size(block.compressed_length, label, index, "physical length")?; + ensure!( + uncompressed_offset == uncompressed_size, + "LZ4DIFF {label} blocks are not contiguous at block {index}" + ); + ensure!( + uncompressed_length > 0 && compressed_length > 0, + "LZ4DIFF {label} block {index} has an empty length" + ); + ensure!( + compressed_length <= uncompressed_length, + "LZ4DIFF {label} block {index} physical length exceeds its raw length" + ); + has_compressed_block |= compressed_length < uncompressed_length; + uncompressed_size = uncompressed_size + .checked_add(uncompressed_length) + .context(format!("LZ4DIFF {label} uncompressed size overflow"))?; + physical_size = physical_size + .checked_add(compressed_length) + .context(format!("LZ4DIFF {label} physical size overflow"))?; + ensure_below_bound(uncompressed_size, &format!("{label} uncompressed stream"))?; + ensure_below_bound(physical_size, &format!("{label} physical stream"))?; + + if is_source { + ensure!( + sha256_hash.is_empty() && postfix_bspatch.is_empty(), + "LZ4DIFF source block {index} contains destination-only metadata" + ); + } else { + ensure!( + sha256_hash.is_empty() || sha256_hash.len() == digest::SHA256_OUTPUT_LEN, + "LZ4DIFF destination block {index} SHA-256 hash must be empty or 32 bytes" + ); + if !postfix_bspatch.is_empty() { + ensure!( + compressed_length < uncompressed_length, + "LZ4DIFF destination raw block {index} cannot have a postfix patch" + ); + ensure!( + sha256_hash.len() == digest::SHA256_OUTPUT_LEN, + "LZ4DIFF destination block {index} postfix patch requires a SHA-256 hash" + ); + ensure_below_bound(postfix_bspatch.len(), "postfix patch")?; + bsdiff::validate_output_len(postfix_bspatch, compressed_length) + .with_context(|| format!("LZ4DIFF destination block {index} postfix patch"))?; + puffin::validate_bsdiff_resources( + postfix_bspatch, + compressed_length, + cancellation_token, + ) + .map_err(|error| { + preserve_puffin_error( + error, + &format!("LZ4DIFF destination block {index} postfix patch is invalid"), + ) + })?; + } + } + blocks.push(Block { + uncompressed_offset, + uncompressed_length, + compressed_length, + sha256_hash, + postfix_bspatch, + }); + } + ensure!( + algorithm_type != Algorithm::Uncompressed || !has_compressed_block, + "LZ4DIFF {label} uses compressed blocks with the uncompressed algorithm" + ); + Ok(CompressionInfo { + algorithm: algorithm_type, + level: algorithm.level, + blocks, + zero_padding_enabled: info.zero_padding_enabled, + uncompressed_size, + physical_size, + }) +} + +fn to_size(value: u64, label: &str, index: usize, field: &str) -> Result { + usize::try_from(value) + .with_context(|| format!("LZ4DIFF {label} block {index} {field} is too large")) +} + +fn ensure_below_bound(size: usize, label: &str) -> Result<()> { + ensure!( + size < zucchini::OFFSET_BOUND, + "LZ4DIFF {label} exceeds the maximum size of {} bytes", + zucchini::OFFSET_BOUND - 1 + ); + Ok(()) +} + +fn decompress_source( + source: &[u8], + info: &CompressionInfo, + cancellation_token: &CancellationToken, +) -> Result> { + let mut output = Vec::new(); + output + .try_reserve_exact(info.uncompressed_size) + .context("unable to allocate LZ4DIFF uncompressed source")?; + let mut physical_offset = 0usize; + for (index, block) in info.blocks.iter().enumerate() { + cancellation_token.check()?; + let end = physical_offset + .checked_add(block.compressed_length) + .context("LZ4DIFF source block range overflow")?; + let physical = source + .get(physical_offset..end) + .with_context(|| format!("LZ4DIFF source block {index} exceeds source data"))?; + if block.compressed_length == block.uncompressed_length { + output.extend_from_slice(physical); + } else { + let margin = if info.zero_padding_enabled { + physical + .iter() + .take(EROFS_BLOCK_SIZE.min(physical.len())) + .take_while(|byte| **byte == 0) + .count() + } else { + 0 + }; + let compressed = physical + .get(margin..) + .with_context(|| format!("invalid LZ4DIFF source block {index} input margin"))?; + let decompressed = lz4::decompress_safe_partial(compressed, block.uncompressed_length) + .with_context(|| format!("unable to decompress LZ4DIFF source block {index}"))?; + output.extend_from_slice(&decompressed); + } + physical_offset = end; + } + ensure!( + physical_offset == source.len() && output.len() == info.uncompressed_size, + "LZ4DIFF source block layout size mismatch" + ); + Ok(output) +} + +fn apply_bsdiff( + source: &[u8], + patch: &[u8], + output_size: usize, + label: &str, + cancellation_token: &CancellationToken, +) -> Result> { + crate::apply_bsdiff(source, patch, output_size, cancellation_token).map_err(|error| { + if is_cancellation(error.as_ref()) { + error + } else { + anyhow::anyhow!("LZ4DIFF {label} patch is invalid: {error}") + } + }) +} + +fn preserve_puffin_error(error: puffin::Error, context: &str) -> anyhow::Error { + let message = error.to_string(); + anyhow::Error::new(error).context(format!("{context}: {message}")) +} + +fn compress_destination( + uncompressed: &[u8], + info: &CompressionInfo, + cancellation_token: &CancellationToken, +) -> Result> { + ensure!( + uncompressed.len() == info.uncompressed_size, + "LZ4DIFF destination uncompressed size mismatch" + ); + let mut output = Vec::new(); + output + .try_reserve_exact(info.physical_size) + .context("unable to allocate LZ4DIFF destination")?; + for (index, block) in info.blocks.iter().enumerate() { + cancellation_token.check()?; + let raw_end = block + .uncompressed_offset + .checked_add(block.uncompressed_length) + .context("LZ4DIFF destination block range overflow")?; + let raw = uncompressed + .get(block.uncompressed_offset..raw_end) + .with_context(|| format!("LZ4DIFF destination block {index} exceeds inner output"))?; + if block.compressed_length == block.uncompressed_length { + output.extend_from_slice(raw); + continue; + } + let mut physical = { + let remaining = uncompressed + .get(block.uncompressed_offset..) + .with_context(|| format!("invalid LZ4DIFF destination block {index} offset"))?; + let level = match info.algorithm { + Algorithm::Lz4 => None, + Algorithm::Lz4hc => Some(info.level), + Algorithm::Uncompressed => { + bail!("LZ4DIFF destination compressed block uses uncompressed algorithm") + } + }; + lz4::compress_dest_size_partial( + remaining, + block.uncompressed_length, + block.compressed_length, + level, + info.zero_padding_enabled, + ) + .with_context(|| format!("unable to compress LZ4DIFF destination block {index}"))? + }; + ensure!( + physical.len() == block.compressed_length, + "LZ4DIFF destination block {index} physical size mismatch" + ); + if !block.postfix_bspatch.is_empty() { + cancellation_token.check()?; + let actual = digest::digest(&digest::SHA256, &physical); + ensure!( + actual.as_ref() == block.sha256_hash, + "LZ4DIFF destination block {index} reference recompression hash mismatch" + ); + physical = apply_bsdiff( + &physical, + block.postfix_bspatch, + block.compressed_length, + "postfix BSDIFF", + cancellation_token, + )?; + } + output.extend_from_slice(&physical); + } + ensure!( + output.len() == info.physical_size, + "LZ4DIFF destination physical size mismatch after compression" + ); + cancellation_token.check()?; + Ok(output) +} diff --git a/src/main.rs b/src/main.rs index 5c2e977..70e3e6d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,18 @@ -use anyhow::Result; +use std::process::ExitCode; + use clap::Parser; use indicatif::{ProgressBar, ProgressStyle}; -use otadump::{ExtractOptions, ProgressReporter}; +use otadump::{CancellationToken, ExtractOptions, ProgressReporter, is_cancellation}; -fn main() -> Result<()> { - extract(); - Ok(()) +fn main() -> Result { + let cancellation_token = CancellationToken::new(); + let signal_token = cancellation_token.clone(); + ctrlc::set_handler(move || signal_token.cancel())?; + let exit_code = extract(&cancellation_token); + if cancellation_token.is_cancelled() { + return Ok(ExitCode::from(130)); + } + Ok(exit_code) } const HELP_TEMPLATE: &str = color_print::cstr!( @@ -39,26 +46,37 @@ pub struct Args { /// Path to the output directory #[clap(long)] pub output_dir: String, + + /// Directory containing source partition images for delta payloads + #[clap(long)] + pub source_dir: Option, } -pub fn extract() { +pub fn extract(cancellation_token: &CancellationToken) -> ExitCode { let args = Args::parse(); let reporter = Box::new(CliProgressReporter::new()); let reporter = reporter.as_ref(); - let result = ExtractOptions::new() - .progress_reporter(reporter) - .extract(&args.payload_file, &args.output_dir); + let mut options = ExtractOptions::new(); + options.progress_reporter(reporter).cancellation_token(cancellation_token); + if let Some(source_dir) = &args.source_dir { + options.source_dir(source_dir); + } + let result = options.extract(&args.payload_file, &args.output_dir); reporter.progress_bar.finish_and_clear(); match result { Ok(()) => { let message = format!("Extraction complete: {}", args.output_dir); reporter.progress_bar.println(message); + ExitCode::SUCCESS } Err(e) => { - let message = format!("Error: {e:?}"); - reporter.progress_bar.println(message); + if is_cancellation(e.as_ref()) { + return ExitCode::from(130); + } + eprintln!("Error: {e:?}"); + ExitCode::FAILURE } } } diff --git a/src/puffin/LICENSE b/src/puffin/LICENSE new file mode 100644 index 0000000..f6d41c6 --- /dev/null +++ b/src/puffin/LICENSE @@ -0,0 +1,27 @@ +// Copyright 2017 The ChromiumOS Authors +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are +// met: +// +// * Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// * Redistributions in binary form must reproduce the above +// copyright notice, this list of conditions and the following disclaimer +// in the documentation and/or other materials provided with the +// distribution. +// * Neither the name of Google Inc. nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/src/puffin/README.md b/src/puffin/README.md new file mode 100644 index 0000000..9cc6e37 --- /dev/null +++ b/src/puffin/README.md @@ -0,0 +1,10 @@ +# Puffin-derived Rust code + +This module adapts the pure-Rust Puffin implementation from `puffdiff` 0.1.0. +That crate ports ChromiumOS Puffin's parser, bit I/O, PuffData I/O, Huffman table, puffer, huffer, and whole-buffer stream code. + +The reference Puffin revision is `343e23db1b4d81045e91a10244244893f5acd73b`. +The Rust code is restricted to PUF1 parsing and inner BSDIFF application. +It adds checked arithmetic, fallible allocations, operation size bounds, and cooperative cancellation for otadump. + +The code is available under the BSD-3-Clause terms in `LICENSE`. diff --git a/src/puffin/bit_io.rs b/src/puffin/bit_io.rs new file mode 100644 index 0000000..45ab098 --- /dev/null +++ b/src/puffin/bit_io.rs @@ -0,0 +1,179 @@ +// Copyright 2017 The ChromiumOS Authors +// Use of this source code is governed by a BSD-style license that can be +// found in the LICENSE file. +// +// Adapted from puffdiff 0.1.0, a pure-Rust port of ChromiumOS Puffin. + +//! LSB-first bit reader/writer, ported from puffin `bit_reader.cc` / +//! `bit_writer.cc`. Deflate packs bits least-significant-first within each byte. + +use crate::puffin::{Error, Result}; + +/// Reads bits from a byte buffer, LSB-first, exactly like `BufferBitReader`. +pub struct BitReader<'a> { + buf: &'a [u8], + index: usize, // next byte to pull into the cache + cache: u32, // pending bits, LSB = next bit to read + cache_bits: usize, // number of valid bits in `cache` +} + +impl<'a> BitReader<'a> { + pub fn new(buf: &'a [u8]) -> Self { + BitReader { buf, index: 0, cache: 0, cache_bits: 0 } + } + + /// Ensure at least `nbits` bits are cached. Max 32. + pub fn cache_bits(&mut self, nbits: usize) -> bool { + if (self.buf.len() - self.index) * 8 + self.cache_bits < nbits { + return false; + } + if nbits > 32 { + return false; + } + while self.cache_bits < nbits { + self.cache |= (self.buf[self.index] as u32) << self.cache_bits; + self.index += 1; + self.cache_bits += 8; + } + true + } + + /// Peek the low `nbits` cached bits (call `cache_bits` first). + pub fn read_bits(&self, nbits: usize) -> u32 { + if nbits == 0 { + return 0; + } + if nbits >= 32 { + return self.cache; + } + self.cache & ((1u32 << nbits) - 1) + } + + pub fn drop_bits(&mut self, nbits: usize) { + self.cache >>= nbits; + self.cache_bits -= nbits; + } + + pub fn cached_bits(&self) -> usize { + self.cache_bits + } + + pub fn offset_in_bits(&self) -> Result { + self.index + .checked_mul(8) + .and_then(|offset| offset.checked_sub(self.cache_bits)) + .ok_or_else(|| Error::Corrupt("bit reader position is invalid".into())) + } + + pub fn read_boundary_bits(&self) -> u8 { + (self.cache & ((1u32 << (self.cache_bits & 7)) - 1)) as u8 + } + + pub fn skip_boundary_bits(&mut self) -> usize { + let nbits = self.cache_bits & 7; + self.cache >>= nbits; + self.cache_bits -= nbits; + nbits + } + + pub fn offset(&self) -> usize { + self.index - self.cache_bits / 8 + } + + pub fn bits_remaining(&self) -> u64 { + ((self.buf.len() - self.index) as u64) * 8 + self.cache_bits as u64 + } + + /// Rewind the cache to the current byte boundary and read `count` raw bytes. + /// Mirrors `GetByteReaderFn`: subsequent bit reads start on a byte boundary. + pub fn get_bytes(&mut self, count: usize) -> Result<&'a [u8]> { + self.index -= self.cache_bits.div_ceil(8); + self.cache = 0; + self.cache_bits = 0; + if count > self.buf.len() - self.index { + return Err(Error::Corrupt("get_bytes out of range".into())); + } + let slice = &self.buf[self.index..self.index + count]; + self.index += count; + Ok(slice) + } +} + +/// Writes bits into a fixed-size byte buffer, LSB-first, like `BufferBitWriter`. +pub struct BitWriter<'a> { + out: &'a mut [u8], + index: usize, + holder: u32, + holder_bits: usize, +} + +impl<'a> BitWriter<'a> { + pub fn new(out: &'a mut [u8]) -> Self { + BitWriter { out, index: 0, holder: 0, holder_bits: 0 } + } + + pub fn write_bits(&mut self, mut nbits: usize, mut bits: u32) -> Result<()> { + if ((self.out.len() - self.index) * 8).saturating_sub(self.holder_bits) < nbits { + return Err(Error::Corrupt("bit writer overflow".into())); + } + if nbits > 32 { + return Err(Error::Corrupt("write_bits > 32".into())); + } + while nbits > 0 { + while self.holder_bits >= 8 { + self.out[self.index] = (self.holder & 0xFF) as u8; + self.index += 1; + self.holder >>= 8; + self.holder_bits -= 8; + } + while self.holder_bits < 24 && nbits > 0 { + self.holder |= (bits & 0xFF) << self.holder_bits; + let min = nbits.min(8); + self.holder_bits += min; + bits >>= min; + nbits -= min; + } + } + Ok(()) + } + + /// Write `data` verbatim; requires the writer to be on a byte boundary. + pub fn write_bytes(&mut self, data: &[u8]) -> Result<()> { + if ((self.out.len() - self.index) * 8).saturating_sub(self.holder_bits) < data.len() * 8 { + return Err(Error::Corrupt("bit writer byte overflow".into())); + } + if self.holder_bits % 8 != 0 { + return Err(Error::Corrupt("write_bytes not byte aligned".into())); + } + self.flush()?; + self.out[self.index..self.index + data.len()].copy_from_slice(data); + self.index += data.len(); + Ok(()) + } + + pub fn write_boundary_bits(&mut self, bits: u8) -> Result<()> { + self.write_bits((8 - (self.holder_bits & 7)) & 7, bits as u32) + } + + pub fn flush(&mut self) -> Result<()> { + self.write_boundary_bits(0)?; + while self.holder_bits > 0 { + self.out[self.index] = (self.holder & 0xFF) as u8; + self.index += 1; + self.holder >>= 8; + self.holder_bits -= 8; + } + Ok(()) + } + + pub fn size(&self) -> usize { + self.index + } + + pub fn offset_in_bits(&self) -> Result { + self.index + .checked_mul(8) + .and_then(|offset| offset.checked_add(self.holder_bits)) + .ok_or_else(|| Error::Corrupt("bit writer position overflows".into())) + } +} diff --git a/src/puffin/huffer.rs b/src/puffin/huffer.rs new file mode 100644 index 0000000..f2712e9 --- /dev/null +++ b/src/puffin/huffer.rs @@ -0,0 +1,146 @@ +// Copyright 2017 The ChromiumOS Authors +// Use of this source code is governed by a BSD-style license that can be +// found in the LICENSE file. +// +// Adapted from puffdiff 0.1.0, a pure-Rust port of ChromiumOS Puffin. + +//! Puff -> deflate transform, ported from puffin `huffer.cc`. + +use crate::CancellationToken; +use crate::puffin::bit_io::BitWriter; +use crate::puffin::huffman::{ + BlockType, DISTANCE_BASES, DISTANCE_EXTRA_BITS, HuffmanTable, LENGTH_BASES, LENGTH_EXTRA_BITS, +}; +use crate::puffin::puff_io::{PuffData, PuffReader}; +use crate::puffin::{Error, Result, check_cancelled}; + +/// Re-encode puff tokens from `pr` back into a bit-exact deflate stream in `bw`. +pub fn huff_deflate( + pr: &mut PuffReader, + bw: &mut BitWriter, + cancellation_token: &CancellationToken, +) -> Result<()> { + let mut fixed_ht = HuffmanTable::new()?; + let mut dyn_ht = HuffmanTable::new()?; + let mut fixed_built = false; + + while pr.bytes_left() != 0 { + check_cancelled(cancellation_token)?; + let pd = pr.get_next()?; + let meta = match pd { + PuffData::BlockMetadata(m) => m, + _ => return Err(Error::Corrupt("expected block metadata".into())), + }; + let header = meta[0]; + let final_bit = (header & 0x80) >> 7; + let btype = (header & 0x60) >> 5; + let skipped_bits = header & 0x1F; + + bw.write_bits(1, final_bit as u32)?; + bw.write_bits(2, btype as u32)?; + + let use_dynamic = match BlockType::from_bits(btype)? { + BlockType::Uncompressed => { + bw.write_boundary_bits(skipped_bits)?; + let next = pr.get_next()?; + match next { + PuffData::Literals(raw) => { + bw.write_bits(16, raw.len() as u32)?; + bw.write_bits(16, !(raw.len() as u32) & 0xFFFF)?; + bw.write_bytes(raw)?; + match pr.get_next()? { + PuffData::EndOfBlock => {} + _ => return Err(Error::Corrupt("uncompressed block not ended".into())), + } + } + PuffData::EndOfBlock => { + bw.write_bits(16, 0)?; + bw.write_bits(16, 0xFFFF)?; + } + PuffData::Literal(_) => { + return Err(Error::Corrupt("unexpected single literal".into())); + } + _ => return Err(Error::Corrupt("uncompressed block malformed".into())), + } + continue; + } + BlockType::Fixed => { + if !fixed_built { + fixed_ht.build_fixed()?; + fixed_built = true; + } + false + } + BlockType::Dynamic => { + dyn_ht.build_dynamic_from_puff(&meta[1..], bw)?; + true + } + }; + + let mut block_ended = false; + while !block_ended { + check_cancelled(cancellation_token)?; + let pd = pr.get_next()?; + let cur_ht: &HuffmanTable = if use_dynamic { &dyn_ht } else { &fixed_ht }; + match pd { + PuffData::Literal(b) => { + let (code, nbits) = cur_ht.lit_len_huffman(b as u16)?; + bw.write_bits(nbits, code as u32)?; + } + PuffData::Literals(bytes) => { + for &b in bytes { + check_cancelled(cancellation_token)?; + let (code, nbits) = cur_ht.lit_len_huffman(b as u16)?; + bw.write_bits(nbits, code as u32)?; + } + } + PuffData::LenDist { length, distance } => { + if !(3..=258).contains(&length) { + return Err(Error::Corrupt("length out of range".into())); + } + // Linear search for the length base (guard stops us short). + let mut index = 0usize; + while length > LENGTH_BASES[index] as usize { + index += 1; + } + if length < LENGTH_BASES[index] as usize { + index -= 1; + } + let extra_len = LENGTH_EXTRA_BITS[index] as usize; + let (code, nbits) = cur_ht.lit_len_huffman((index + 257) as u16)?; + bw.write_bits(nbits, code as u32)?; + if extra_len > 0 { + bw.write_bits(extra_len, (length - LENGTH_BASES[index] as usize) as u32)?; + } + + let mut index = 0usize; + while distance > DISTANCE_BASES[index] as usize { + index += 1; + } + if distance < DISTANCE_BASES[index] as usize { + index -= 1; + } + let extra_len = DISTANCE_EXTRA_BITS[index] as usize; + let (code, nbits) = cur_ht.distance_huffman(index as u16)?; + bw.write_bits(nbits, code as u32)?; + if extra_len > 0 { + bw.write_bits( + extra_len, + (distance - DISTANCE_BASES[index] as usize) as u32, + )?; + } + } + PuffData::EndOfBlock => { + let (code, nbits) = cur_ht.lit_len_huffman(256)?; + bw.write_bits(nbits, code as u32)?; + block_ended = true; + } + PuffData::BlockMetadata(_) => { + return Err(Error::Corrupt("unexpected block metadata".into())); + } + } + } + } + check_cancelled(cancellation_token)?; + Ok(()) +} diff --git a/src/puffin/huffman.rs b/src/puffin/huffman.rs new file mode 100644 index 0000000..f5eeb1f --- /dev/null +++ b/src/puffin/huffman.rs @@ -0,0 +1,766 @@ +// Copyright 2017 The ChromiumOS Authors +// Use of this source code is governed by a BSD-style license that can be +// found in the LICENSE file. +// +// Adapted from puffdiff 0.1.0, a pure-Rust port of ChromiumOS Puffin. + +//! RFC1951 Huffman handling, ported from puffin `huffman_table.cc`. +//! +//! Builds canonical decode/encode tables and (de)serializes a dynamic Huffman +//! block header to/from the puff `block_metadata` blob. Byte-for-byte fidelity +//! with puffin is required. + +use crate::puffin::bit_io::{BitReader, BitWriter}; +use crate::puffin::{Error, Result}; + +pub const MAX_HUFFMAN_BITS: usize = 15; + +const PERMUTATIONS: [u8; 19] = [16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15]; + +pub const LENGTH_BASES: [u16; 30] = [ + 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 67, 83, 99, 115, 131, + 163, 195, 227, 258, 0xFFFF, +]; + +pub const LENGTH_EXTRA_BITS: [u8; 29] = + [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0]; + +pub const DISTANCE_BASES: [u16; 31] = [ + 1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1025, 1537, + 2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577, 0xFFFF, +]; + +pub const DISTANCE_EXTRA_BITS: [u8; 30] = [ + 0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, + 13, +]; + +fn replace_bytes(target: &mut Vec, bytes: &[u8], label: &str) -> Result<()> { + target.clear(); + target + .try_reserve_exact(bytes.len()) + .map_err(|error| Error::Allocation(format!("{label}: {error}")))?; + target.extend_from_slice(bytes); + Ok(()) +} + +fn fill_bytes(target: &mut Vec, length: usize, value: u8, label: &str) -> Result<()> { + target.clear(); + target + .try_reserve_exact(length) + .map_err(|error| Error::Allocation(format!("{label}: {error}")))?; + target.resize(length, value); + Ok(()) +} + +#[derive(Clone, Copy)] +pub enum BlockType { + Uncompressed, + Fixed, + Dynamic, +} + +impl BlockType { + pub fn from_bits(v: u8) -> Result { + match v { + 0 => Ok(BlockType::Uncompressed), + 1 => Ok(BlockType::Fixed), + 2 => Ok(BlockType::Dynamic), + _ => Err(Error::Corrupt("invalid block type".into())), + } + } +} + +/// One (canonical code, alphabet index) pair used while building tables. +#[derive(Clone, Copy)] +struct CodeIndexPair { + code: u16, + index: u16, +} + +/// Holds the current block's literal/length, distance and code-length tables. +pub struct HuffmanTable { + lit_len_lens: Vec, + lit_len_hcodes: Vec, + lit_len_rcodes: Vec, + lit_len_max_bits: usize, + + distance_lens: Vec, + distance_hcodes: Vec, + distance_rcodes: Vec, + distance_max_bits: usize, + + code_lens: Vec, + code_hcodes: Vec, + code_rcodes: Vec, + code_max_bits: usize, + + tmp_lens: Vec, + pairs: Vec, +} + +impl HuffmanTable { + pub fn new() -> Result { + let mut code_lens = Vec::new(); + code_lens + .try_reserve_exact(19) + .map_err(|error| Error::Allocation(format!("Huffman code lengths: {error}")))?; + code_lens.resize(19, 0); + let mut pairs = Vec::new(); + pairs + .try_reserve_exact(288) + .map_err(|error| Error::Allocation(format!("Huffman code pairs: {error}")))?; + Ok(HuffmanTable { + lit_len_lens: Vec::new(), + lit_len_hcodes: Vec::new(), + lit_len_rcodes: Vec::new(), + lit_len_max_bits: 0, + distance_lens: Vec::new(), + distance_hcodes: Vec::new(), + distance_rcodes: Vec::new(), + distance_max_bits: 0, + code_lens, + code_hcodes: Vec::new(), + code_rcodes: Vec::new(), + code_max_bits: 0, + tmp_lens: Vec::new(), + pairs, + }) + } + + pub fn lit_len_max_bits(&self) -> usize { + self.lit_len_max_bits + } + pub fn distance_max_bits(&self) -> usize { + self.distance_max_bits + } + + /// Decode lookup: map `bits` (max_bits wide) to a lit/len alphabet + code + /// length. Returns `(alphabet, nbits)`. + pub fn lit_len_alphabet(&self, bits: u32) -> Result<(u16, usize)> { + let hc = *self + .lit_len_hcodes + .get(bits as usize) + .ok_or_else(|| Error::Corrupt("lit/len lookup is out of range".into()))?; + if hc & 0x8000 == 0 { + return Err(Error::Corrupt("invalid lit/len code".into())); + } + let alphabet = hc & 0x7FFF; + let length = *self + .lit_len_lens + .get(alphabet as usize) + .ok_or_else(|| Error::Corrupt("lit/len alphabet is out of range".into()))?; + Ok((alphabet, length as usize)) + } + + pub fn distance_alphabet(&self, bits: u32) -> Result<(u16, usize)> { + let hc = *self + .distance_hcodes + .get(bits as usize) + .ok_or_else(|| Error::Corrupt("distance lookup is out of range".into()))?; + if hc & 0x8000 == 0 { + return Err(Error::Corrupt("invalid distance code".into())); + } + let alphabet = hc & 0x7FFF; + let length = *self + .distance_lens + .get(alphabet as usize) + .ok_or_else(|| Error::Corrupt("distance alphabet is out of range".into()))?; + Ok((alphabet, length as usize)) + } + + fn code_alphabet(&self, bits: u32) -> Result<(u16, usize)> { + let hc = *self + .code_hcodes + .get(bits as usize) + .ok_or_else(|| Error::Corrupt("code-length lookup is out of range".into()))?; + if hc & 0x8000 == 0 { + return Err(Error::Corrupt("invalid code-length code".into())); + } + let alphabet = hc & 0x7FFF; + let length = *self + .code_lens + .get(alphabet as usize) + .ok_or_else(|| Error::Corrupt("code-length alphabet is out of range".into()))?; + Ok((alphabet, length as usize)) + } + + pub fn end_of_block_bit_length(&self) -> Result { + if 256 >= self.lit_len_lens.len() { + return Err(Error::Corrupt("no end-of-block symbol".into())); + } + Ok(self.lit_len_lens[256] as usize) + } + + /// Encode lookup: alphabet -> (huffman code, nbits). + pub fn lit_len_huffman(&self, alphabet: u16) -> Result<(u16, usize)> { + if (alphabet as usize) >= self.lit_len_lens.len() { + return Err(Error::Corrupt("lit/len alphabet out of range".into())); + } + Ok((self.lit_len_rcodes[alphabet as usize], self.lit_len_lens[alphabet as usize] as usize)) + } + + pub fn distance_huffman(&self, alphabet: u16) -> Result<(u16, usize)> { + if (alphabet as usize) >= self.distance_lens.len() { + return Err(Error::Corrupt("distance alphabet out of range".into())); + } + Ok(( + self.distance_rcodes[alphabet as usize], + self.distance_lens[alphabet as usize] as usize, + )) + } + + fn code_huffman(&self, alphabet: u16) -> Result<(u16, usize)> { + if (alphabet as usize) >= self.code_lens.len() { + return Err(Error::Corrupt("code alphabet out of range".into())); + } + Ok((self.code_rcodes[alphabet as usize], self.code_lens[alphabet as usize] as usize)) + } + + /// Compute canonical codes from an array of code lengths, filling `pairs`. + /// Returns the max code length seen (0 if all lengths are zero). + fn init_huffman_codes(lens: &[u8], pairs: &mut Vec) -> Result { + let mut len_count = [0u16; MAX_HUFFMAN_BITS + 1]; + for &l in lens { + let count = len_count + .get_mut(l as usize) + .ok_or_else(|| Error::Corrupt("Huffman code length exceeds 15 bits".into()))?; + *count = count + .checked_add(1) + .ok_or_else(|| Error::Corrupt("Huffman code count overflows".into()))?; + } + + let mut max_bits = 0usize; + for bits in (1..=MAX_HUFFMAN_BITS).rev() { + if len_count[bits] != 0 { + max_bits = bits; + break; + } + } + + let mut available = 1i32; + for count in len_count.iter().take(MAX_HUFFMAN_BITS + 1).skip(1) { + available = available + .checked_mul(2) + .and_then(|slots| slots.checked_sub(i32::from(*count))) + .ok_or_else(|| Error::Corrupt("Huffman code-space accounting overflows".into()))?; + if available < 0 { + return Err(Error::Corrupt("oversubscribed code lengths".into())); + } + } + + let mut code = 0u16; + len_count[0] = 0; + let mut next_code = [0u16; MAX_HUFFMAN_BITS + 1]; + for bits in 1..=MAX_HUFFMAN_BITS { + code = (code + len_count[bits - 1]) << 1; + next_code[bits] = code; + } + + pairs.clear(); + for (idx, &len_u8) in lens.iter().enumerate() { + let len = len_u8 as usize; + if len == 0 { + continue; + } + // Bit-reverse the canonical code (deflate reads codes MSB-first). + let mut rev = 0u16; + let mut tmp = next_code[len]; + for _ in 0..len { + rev <<= 1; + rev |= tmp & 1; + tmp >>= 1; + } + pairs.push(CodeIndexPair { code: rev, index: idx as u16 }); + next_code[len] += 1; + } + Ok(max_bits) + } + + fn build_huffman_codes( + lens: &[u8], + pairs: &mut [CodeIndexPair], + hcodes: &mut Vec, + max_bits: usize, + ) -> Result<()> { + // Sort descending by code length so shorter codes fill the don't-care + // slots without overwriting exact matches. + pairs.sort_by(|a, b| lens[b.index as usize].cmp(&lens[a.index as usize])); + hcodes.clear(); + let size = 1usize + .checked_shl(max_bits as u32) + .ok_or_else(|| Error::Corrupt("Huffman lookup size overflows".into()))?; + hcodes + .try_reserve_exact(size) + .map_err(|error| Error::Allocation(format!("Huffman lookup table: {error}")))?; + hcodes.resize(size, 0); + for cip in pairs.iter() { + let slot = hcodes + .get_mut(cip.code as usize) + .ok_or_else(|| Error::Corrupt("Huffman code exceeds lookup table".into()))?; + *slot = cip.index | 0x8000; + let code_len = lens[cip.index as usize] as usize; + let fill_bits = max_bits - code_len; + for idx in 1..(1usize << fill_bits) { + let location = (idx << code_len) | cip.code as usize; + let slot = hcodes + .get_mut(location) + .ok_or_else(|| Error::Corrupt("Huffman fill exceeds lookup table".into()))?; + if *slot & 0x8000 == 0 { + *slot = cip.index | 0x8000; + } + } + } + Ok(()) + } + + fn build_huffman_reverse_codes( + lens: &[u8], + pairs: &mut [CodeIndexPair], + rcodes: &mut Vec, + ) -> Result<()> { + pairs.sort_by_key(|pair| pair.index); + rcodes.clear(); + rcodes + .try_reserve_exact(lens.len()) + .map_err(|error| Error::Allocation(format!("Huffman reverse table: {error}")))?; + rcodes.resize(lens.len(), 0); + let mut i = 0usize; + for (idx, slot) in rcodes.iter_mut().enumerate() { + if i < pairs.len() && idx == pairs[i].index as usize { + *slot = pairs[i].code; + i += 1; + } else { + *slot = 0; + } + } + Ok(()) + } + + pub fn build_fixed(&mut self) -> Result<()> { + fill_bytes(&mut self.lit_len_lens, 288, 0, "fixed literal lengths")?; + let mut i = 0; + while i < 144 { + self.lit_len_lens[i] = 8; + i += 1; + } + while i < 256 { + self.lit_len_lens[i] = 9; + i += 1; + } + while i < 280 { + self.lit_len_lens[i] = 7; + i += 1; + } + while i < 288 { + self.lit_len_lens[i] = 8; + i += 1; + } + fill_bytes(&mut self.distance_lens, 30, 5, "fixed distance lengths")?; + + let mut pairs; + { + self.lit_len_max_bits = Self::init_huffman_codes(&self.lit_len_lens, &mut self.pairs)?; + pairs = std::mem::take(&mut self.pairs); + Self::build_huffman_codes( + &self.lit_len_lens, + &mut pairs, + &mut self.lit_len_hcodes, + self.lit_len_max_bits, + )?; + self.pairs = pairs; + } + { + self.distance_max_bits = + Self::init_huffman_codes(&self.distance_lens, &mut self.pairs)?; + pairs = std::mem::take(&mut self.pairs); + Self::build_huffman_codes( + &self.distance_lens, + &mut pairs, + &mut self.distance_hcodes, + self.distance_max_bits, + )?; + self.pairs = pairs; + } + // reverse codes for the encode direction + { + let _ = Self::init_huffman_codes(&self.lit_len_lens, &mut self.pairs)?; + pairs = std::mem::take(&mut self.pairs); + Self::build_huffman_reverse_codes( + &self.lit_len_lens, + &mut pairs, + &mut self.lit_len_rcodes, + )?; + self.pairs = pairs; + } + { + let _ = Self::init_huffman_codes(&self.distance_lens, &mut self.pairs)?; + pairs = std::mem::take(&mut self.pairs); + Self::build_huffman_reverse_codes( + &self.distance_lens, + &mut pairs, + &mut self.distance_rcodes, + )?; + self.pairs = pairs; + } + Ok(()) + } + + fn check_array_lengths( + num_lit_len: usize, + num_distance: usize, + num_codes: usize, + ) -> Result<()> { + if num_lit_len > 286 || num_distance > 30 || num_codes > 19 { + return Err(Error::Corrupt("dynamic huffman array lengths invalid".into())); + } + Ok(()) + } + + /// Read a dynamic Huffman block header from the deflate stream `br`, writing + /// its puff encoding into `buffer`. Returns bytes written to `buffer`. + pub fn build_dynamic_from_deflate( + &mut self, + br: &mut BitReader, + buffer: &mut [u8], + ) -> Result { + if buffer.len() < 3 { + return Err(Error::Corrupt("metadata buffer too small".into())); + } + let mut index = 0usize; + if !br.cache_bits(14) { + return Err(Error::Corrupt("eof reading dynamic header".into())); + } + buffer[index] = br.read_bits(5) as u8; + let num_lit_len = br.read_bits(5) as usize + 257; + br.drop_bits(5); + index += 1; + + buffer[index] = br.read_bits(5) as u8; + let num_distance = br.read_bits(5) as usize + 1; + br.drop_bits(5); + index += 1; + + buffer[index] = br.read_bits(4) as u8; + let num_codes = br.read_bits(4) as usize + 4; + br.drop_bits(4); + index += 1; + + Self::check_array_lengths(num_lit_len, num_distance, num_codes)?; + + let mut checked = false; + if buffer.len() - index < num_codes.div_ceil(2) { + return Err(Error::Corrupt("metadata buffer too small for codes".into())); + } + let mut idx = 0usize; + while idx < num_codes { + if !br.cache_bits(3) { + return Err(Error::Corrupt("eof reading code lengths".into())); + } + self.code_lens[PERMUTATIONS[idx] as usize] = br.read_bits(3) as u8; + if checked { + buffer[index] |= br.read_bits(3) as u8; + index += 1; + } else { + buffer[index] = (br.read_bits(3) as u8) << 4; + } + checked = !checked; + br.drop_bits(3); + idx += 1; + } + if checked { + index += 1; + } + while idx < 19 { + self.code_lens[PERMUTATIONS[idx] as usize] = 0; + idx += 1; + } + + self.code_max_bits = Self::init_huffman_codes(&self.code_lens, &mut self.pairs)?; + let mut pairs = std::mem::take(&mut self.pairs); + Self::build_huffman_codes( + &self.code_lens, + &mut pairs, + &mut self.code_hcodes, + self.code_max_bits, + )?; + self.pairs = pairs; + + // Read the lit/len + distance code-length arrays (RLE-coded). + let written = self.read_code_lengths_from_deflate( + br, + &mut buffer[index..], + self.code_max_bits, + num_lit_len + num_distance, + )?; + index += written; + + let tmp = std::mem::take(&mut self.tmp_lens); + replace_bytes(&mut self.lit_len_lens, &tmp[..num_lit_len], "dynamic literal lengths")?; + replace_bytes(&mut self.distance_lens, &tmp[num_lit_len..], "dynamic distance lengths")?; + self.tmp_lens = tmp; + + self.lit_len_max_bits = Self::init_huffman_codes(&self.lit_len_lens, &mut self.pairs)?; + let mut pairs = std::mem::take(&mut self.pairs); + Self::build_huffman_codes( + &self.lit_len_lens, + &mut pairs, + &mut self.lit_len_hcodes, + self.lit_len_max_bits, + )?; + self.pairs = pairs; + + self.distance_max_bits = Self::init_huffman_codes(&self.distance_lens, &mut self.pairs)?; + let mut pairs = std::mem::take(&mut self.pairs); + Self::build_huffman_codes( + &self.distance_lens, + &mut pairs, + &mut self.distance_hcodes, + self.distance_max_bits, + )?; + self.pairs = pairs; + + Ok(index) + } + + fn read_code_lengths_from_deflate( + &mut self, + br: &mut BitReader, + buffer: &mut [u8], + max_bits: usize, + num_codes: usize, + ) -> Result { + let mut index = 0usize; + let mut lens = Vec::new(); + lens.try_reserve_exact(num_codes) + .map_err(|error| Error::Allocation(format!("dynamic code lengths: {error}")))?; + let mut idx = 0usize; + while idx < num_codes { + if !br.cache_bits(max_bits) { + return Err(Error::Corrupt("eof reading code length codes".into())); + } + let bits = br.read_bits(max_bits); + let (code, nbits) = self.code_alphabet(bits)?; + if index >= buffer.len() { + return Err(Error::Corrupt("code length buffer overflow".into())); + } + br.drop_bits(nbits); + if code < 16 { + buffer[index] = code as u8; + index += 1; + lens.push(code as u8); + idx += 1; + } else { + if code >= 19 { + return Err(Error::Corrupt("invalid code-length code".into())); + } + let (copy_num, copy_val) = match code { + 16 => { + if idx == 0 { + return Err(Error::Corrupt("repeat with no previous length".into())); + } + if !br.cache_bits(2) { + return Err(Error::Corrupt("eof reading repeat".into())); + } + let n = 3 + br.read_bits(2) as usize; + buffer[index] = 16 + br.read_bits(2) as u8; + index += 1; + let v = lens[idx - 1]; + br.drop_bits(2); + (n, v) + } + 17 => { + if !br.cache_bits(3) { + return Err(Error::Corrupt("eof reading repeat".into())); + } + let n = 3 + br.read_bits(3) as usize; + buffer[index] = 20 + br.read_bits(3) as u8; + index += 1; + br.drop_bits(3); + (n, 0) + } + 18 => { + if !br.cache_bits(7) { + return Err(Error::Corrupt("eof reading repeat".into())); + } + let n = 11 + br.read_bits(7) as usize; + buffer[index] = 28 + br.read_bits(7) as u8; + index += 1; + br.drop_bits(7); + (n, 0) + } + _ => return Err(Error::Corrupt("invalid code-length code".into())), + }; + if copy_num > num_codes - idx { + return Err(Error::Corrupt("code-length repeat exceeds array".into())); + } + idx += copy_num; + for _ in 0..copy_num { + lens.push(copy_val); + } + } + } + if lens.len() != num_codes { + return Err(Error::Corrupt("code length count mismatch".into())); + } + self.tmp_lens = lens; + Ok(index) + } + + /// Write a dynamic Huffman block header into the deflate stream `bw` from + /// its puff encoding in `buffer`. + pub fn build_dynamic_from_puff(&mut self, buffer: &[u8], bw: &mut BitWriter) -> Result<()> { + if buffer.len() < 3 { + return Err(Error::Corrupt("metadata buffer too small".into())); + } + let mut index = 0usize; + let num_lit_len = buffer[index] as usize + 257; + bw.write_bits(5, buffer[index] as u32)?; + index += 1; + + let num_distance = buffer[index] as usize + 1; + bw.write_bits(5, buffer[index] as u32)?; + index += 1; + + let num_codes = buffer[index] as usize + 4; + bw.write_bits(4, buffer[index] as u32)?; + index += 1; + + Self::check_array_lengths(num_lit_len, num_distance, num_codes)?; + + if buffer.len() - index < num_codes.div_ceil(2) { + return Err(Error::Corrupt("metadata truncated".into())); + } + let mut checked = false; + let mut idx = 0usize; + while idx < num_codes { + let len = if checked { + let v = buffer[index] & 0x0F; + index += 1; + v + } else { + buffer[index] >> 4 + }; + checked = !checked; + self.code_lens[PERMUTATIONS[idx] as usize] = len; + bw.write_bits(3, len as u32)?; + idx += 1; + } + if checked { + index += 1; + } + while idx < 19 { + self.code_lens[PERMUTATIONS[idx] as usize] = 0; + idx += 1; + } + + let _ = Self::init_huffman_codes(&self.code_lens, &mut self.pairs)?; + let mut pairs = std::mem::take(&mut self.pairs); + Self::build_huffman_reverse_codes(&self.code_lens, &mut pairs, &mut self.code_rcodes)?; + self.pairs = pairs; + + let read = + self.write_code_lengths_to_deflate(&buffer[index..], bw, num_lit_len + num_distance)?; + index += read; + + let tmp = std::mem::take(&mut self.tmp_lens); + replace_bytes(&mut self.lit_len_lens, &tmp[..num_lit_len], "dynamic literal lengths")?; + replace_bytes(&mut self.distance_lens, &tmp[num_lit_len..], "dynamic distance lengths")?; + self.tmp_lens = tmp; + + let _ = Self::init_huffman_codes(&self.lit_len_lens, &mut self.pairs)?; + let mut pairs = std::mem::take(&mut self.pairs); + Self::build_huffman_reverse_codes( + &self.lit_len_lens, + &mut pairs, + &mut self.lit_len_rcodes, + )?; + self.pairs = pairs; + // lit_len_max_bits is unused when encoding, but keep it consistent. + self.lit_len_max_bits = Self::init_huffman_codes(&self.lit_len_lens, &mut self.pairs)?; + self.pairs.clear(); + + let _ = Self::init_huffman_codes(&self.distance_lens, &mut self.pairs)?; + let mut pairs = std::mem::take(&mut self.pairs); + Self::build_huffman_reverse_codes( + &self.distance_lens, + &mut pairs, + &mut self.distance_rcodes, + )?; + self.pairs = pairs; + + if index != buffer.len() { + return Err(Error::Corrupt("metadata length mismatch".into())); + } + Ok(()) + } + + fn write_code_lengths_to_deflate( + &mut self, + buffer: &[u8], + bw: &mut BitWriter, + num_codes: usize, + ) -> Result { + let mut lens = Vec::new(); + lens.try_reserve_exact(num_codes) + .map_err(|error| Error::Allocation(format!("dynamic code lengths: {error}")))?; + let mut index = 0usize; + let mut idx = 0usize; + while idx < num_codes { + if index >= buffer.len() { + return Err(Error::Corrupt("code length buffer underflow".into())); + } + let pcode = buffer[index]; + index += 1; + if pcode > 155 { + return Err(Error::Corrupt("invalid puffed code length".into())); + } + let code: u16 = if pcode < 16 { + pcode as u16 + } else if pcode < 20 { + 16 + } else if pcode < 28 { + 17 + } else { + 18 + }; + let (hcode, nbits) = self.code_huffman(code)?; + bw.write_bits(nbits, hcode as u32)?; + if code < 16 { + lens.push(code as u8); + idx += 1; + } else { + let (copy_num, copy_val) = match code { + 16 => { + if idx == 0 { + return Err(Error::Corrupt("repeat with no previous length".into())); + } + bw.write_bits(2, (pcode - 16) as u32)?; + (3 + (pcode - 16) as usize, lens[idx - 1]) + } + 17 => { + bw.write_bits(3, (pcode - 20) as u32)?; + (3 + (pcode - 20) as usize, 0) + } + 18 => { + bw.write_bits(7, (pcode - 28) as u32)?; + (11 + (pcode - 28) as usize, 0) + } + _ => return Err(Error::Corrupt("invalid code-length code".into())), + }; + if copy_num > num_codes - idx { + return Err(Error::Corrupt("code-length repeat exceeds array".into())); + } + idx += copy_num; + for _ in 0..copy_num { + lens.push(copy_val); + } + } + } + if lens.len() != num_codes { + return Err(Error::Corrupt("code length count mismatch".into())); + } + self.tmp_lens = lens; + Ok(index) + } +} diff --git a/src/puffin/mod.rs b/src/puffin/mod.rs new file mode 100644 index 0000000..5d1f25d --- /dev/null +++ b/src/puffin/mod.rs @@ -0,0 +1,108 @@ +// Copyright 2017 The ChromiumOS Authors +// Use of this source code is governed by a BSD-style license that can be +// found in the LICENSE file. +// +// Adapted from puffdiff 0.1.0, a pure-Rust port of ChromiumOS Puffin. + +mod bit_io; +mod huffer; +mod huffman; +mod patch; +mod puff_io; +mod puffer; +mod stream; + +use std::error; +use std::fmt; + +use crate::ExtractionCancelled; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct BitExtent { + offset: usize, + length: usize, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct ByteExtent { + offset: usize, + length: usize, +} + +#[derive(Debug)] +pub(crate) enum Error { + BadPatchHeader(String), + BadProto(String), + UnsupportedPatchType(i32), + InvalidMetadata(String), + Corrupt(String), + Bsdiff(String), + Zucchini(String), + SizeMismatch { expected: usize, actual: usize }, + Allocation(String), + Cancelled(ExtractionCancelled), +} + +impl fmt::Display for Error { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::BadPatchHeader(message) => { + write!(formatter, "invalid PUFFDIFF header: {message}") + } + Self::BadProto(message) => { + write!(formatter, "invalid PUFFDIFF protobuf: {message}") + } + Self::UnsupportedPatchType(patch_type) => { + write!(formatter, "unsupported PUFFDIFF patch type {patch_type}") + } + Self::InvalidMetadata(message) => { + write!(formatter, "invalid PUFFDIFF metadata: {message}") + } + Self::Corrupt(message) => write!(formatter, "corrupt Puffin stream: {message}"), + Self::Bsdiff(message) => write!(formatter, "inner BSDIFF patch is invalid: {message}"), + Self::Zucchini(message) => { + write!(formatter, "inner ZUCCHINI patch is invalid: {message}") + } + Self::SizeMismatch { expected, actual } => { + write!(formatter, "PUFFDIFF size mismatch: expected {expected}, got {actual}") + } + Self::Allocation(message) => write!(formatter, "unable to allocate {message}"), + Self::Cancelled(error) => error.fmt(formatter), + } + } +} + +impl error::Error for Error { + fn source(&self) -> Option<&(dyn error::Error + 'static)> { + match self { + Self::Cancelled(error) => Some(error), + _ => None, + } + } +} + +type Result = std::result::Result; + +fn check_cancelled(token: &crate::CancellationToken) -> Result<()> { + if token.is_cancelled() { + return Err(Error::Cancelled(ExtractionCancelled)); + } + Ok(()) +} + +pub(crate) fn apply( + source: &[u8], + patch: &[u8], + destination_size: usize, + cancellation_token: &crate::CancellationToken, +) -> Result> { + patch::apply(source, patch, destination_size, cancellation_token) +} + +pub(crate) fn validate_bsdiff_resources( + patch: &[u8], + output_size: usize, + cancellation_token: &crate::CancellationToken, +) -> Result<()> { + patch::validate_bsdiff_resources(patch, output_size, cancellation_token) +} diff --git a/src/puffin/patch.rs b/src/puffin/patch.rs new file mode 100644 index 0000000..6dcb311 --- /dev/null +++ b/src/puffin/patch.rs @@ -0,0 +1,551 @@ +// Copyright 2017 The ChromiumOS Authors +// Use of this source code is governed by a BSD-style license that can be +// found in the LICENSE file. +// +// Adapted from puffdiff 0.1.0 and ChromiumOS Puffin's puffpatch.cc. + +use std::io::Read; + +use brotli::Decompressor as BrotliDecoder; +use bzip2::read::BzDecoder; + +use crate::puffin::{BitExtent, ByteExtent, Error, Result, check_cancelled, stream}; +use crate::{ + CancellationToken, ExtractionCancelled, decode_zucchini_patch, is_cancellation, zucchini, +}; + +const MAGIC: &[u8; 4] = b"PUF1"; +const PATCH_TYPE_BSDIFF: i32 = 0; +const PATCH_TYPE_ZUCCHINI: i32 = 1; +const DECODE_CHUNK_SIZE: usize = 32 * 1024; + +#[derive(Clone, Copy)] +enum Compression { + None, + Bzip2, + Brotli, +} + +struct StreamInfo { + deflates: Vec, + puffs: Vec, + puff_length: usize, +} + +impl StreamInfo { + fn empty() -> Self { + Self { deflates: Vec::new(), puffs: Vec::new(), puff_length: 0 } + } +} + +struct Cursor<'a> { + bytes: &'a [u8], + position: usize, +} + +impl<'a> Cursor<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, position: 0 } + } + + fn at_end(&self) -> bool { + self.position == self.bytes.len() + } + + fn varint(&mut self) -> Result { + let mut value = 0u64; + for index in 0..10 { + let byte = *self + .bytes + .get(self.position) + .ok_or_else(|| Error::BadProto("truncated protobuf varint".into()))?; + self.position += 1; + if index == 9 && byte > 1 { + return Err(Error::BadProto("protobuf varint overflows u64".into())); + } + value |= u64::from(byte & 0x7f) << (index * 7); + if byte & 0x80 == 0 { + return Ok(value); + } + } + Err(Error::BadProto("protobuf varint is too long".into())) + } + + fn length_delimited(&mut self) -> Result<&'a [u8]> { + let length = usize::try_from(self.varint()?) + .map_err(|_| Error::BadProto("protobuf field length is too large".into()))?; + let end = self + .position + .checked_add(length) + .ok_or_else(|| Error::BadProto("protobuf field range overflows".into()))?; + let field = self.bytes.get(self.position..end).ok_or_else(|| { + Error::BadProto("length-delimited protobuf field overruns header".into()) + })?; + self.position = end; + Ok(field) + } + + fn skip(&mut self, wire_type: u8) -> Result<()> { + match wire_type { + 0 => { + self.varint()?; + } + 1 => self.advance(8)?, + 2 => { + self.length_delimited()?; + } + 5 => self.advance(4)?, + _ => { + return Err(Error::BadProto(format!("unsupported protobuf wire type {wire_type}"))); + } + } + Ok(()) + } + + fn advance(&mut self, count: usize) -> Result<()> { + let end = self + .position + .checked_add(count) + .ok_or_else(|| Error::BadProto("protobuf field range overflows".into()))?; + if end > self.bytes.len() { + return Err(Error::BadProto("protobuf field overruns header".into())); + } + self.position = end; + Ok(()) + } + + fn key(&mut self) -> Result<(u64, u8)> { + let key = self.varint()?; + let field = key >> 3; + if field == 0 { + return Err(Error::BadProto("protobuf field number is zero".into())); + } + Ok((field, (key & 7) as u8)) + } +} + +fn parse_bit_extent(bytes: &[u8]) -> Result<(u64, u64)> { + let mut cursor = Cursor::new(bytes); + let mut offset = 0u64; + let mut length = 0u64; + while !cursor.at_end() { + let (field, wire_type) = cursor.key()?; + match (field, wire_type) { + (1, 0) => offset = cursor.varint()?, + (2, 0) => length = cursor.varint()?, + _ => cursor.skip(wire_type)?, + } + } + Ok((offset, length)) +} + +fn push_extent(extents: &mut Vec, extent: T) -> Result<()> { + extents + .try_reserve(1) + .map_err(|error| Error::Allocation(format!("PUFFDIFF extent metadata: {error}")))?; + extents.push(extent); + Ok(()) +} + +fn parse_stream_info(bytes: &[u8]) -> Result { + let mut info = StreamInfo::empty(); + let mut cursor = Cursor::new(bytes); + while !cursor.at_end() { + let (field, wire_type) = cursor.key()?; + match (field, wire_type) { + (1, 2) => { + let (offset, length) = parse_bit_extent(cursor.length_delimited()?)?; + push_extent( + &mut info.deflates, + BitExtent { + offset: usize::try_from(offset).map_err(|_| { + Error::InvalidMetadata("deflate offset is too large".into()) + })?, + length: usize::try_from(length).map_err(|_| { + Error::InvalidMetadata("deflate length is too large".into()) + })?, + }, + )?; + } + (2, 2) => { + let (offset, length) = parse_bit_extent(cursor.length_delimited()?)?; + if offset % 8 != 0 || length % 8 != 0 { + return Err(Error::InvalidMetadata("puff extent is not byte-aligned".into())); + } + push_extent( + &mut info.puffs, + ByteExtent { + offset: usize::try_from(offset / 8).map_err(|_| { + Error::InvalidMetadata("puff offset is too large".into()) + })?, + length: usize::try_from(length / 8).map_err(|_| { + Error::InvalidMetadata("puff length is too large".into()) + })?, + }, + )?; + } + (3, 0) => { + info.puff_length = usize::try_from(cursor.varint()?).map_err(|_| { + Error::InvalidMetadata("puff stream length is too large".into()) + })?; + } + _ => cursor.skip(wire_type)?, + } + } + Ok(info) +} + +struct PatchHeader { + source: StreamInfo, + destination: StreamInfo, + patch_type: i32, +} + +fn parse_header(bytes: &[u8]) -> Result { + let mut source = StreamInfo::empty(); + let mut destination = StreamInfo::empty(); + let mut patch_type = PATCH_TYPE_BSDIFF; + let mut cursor = Cursor::new(bytes); + while !cursor.at_end() { + let (field, wire_type) = cursor.key()?; + match (field, wire_type) { + (1, 0) => { + cursor.varint()?; + } + (2, 2) => source = parse_stream_info(cursor.length_delimited()?)?, + (3, 2) => destination = parse_stream_info(cursor.length_delimited()?)?, + (4, 0) => patch_type = cursor.varint()? as u32 as i32, + _ => cursor.skip(wire_type)?, + } + } + Ok(PatchHeader { source, destination, patch_type }) +} + +fn validate_extents( + extents: &[T], + label: &str, + limit: usize, + parts: impl Fn(&T) -> (usize, usize), +) -> Result<()> { + let mut previous_end = 0usize; + for (index, extent) in extents.iter().enumerate() { + let (offset, length) = parts(extent); + if length == 0 { + return Err(Error::InvalidMetadata(format!("{label} extent {index} is empty"))); + } + let end = offset + .checked_add(length) + .ok_or_else(|| Error::InvalidMetadata(format!("{label} extent {index} overflows")))?; + if index != 0 && offset < previous_end { + return Err(Error::InvalidMetadata(format!("{label} extents overlap or are unsorted"))); + } + if end > limit { + return Err(Error::InvalidMetadata(format!( + "{label} extent {index} exceeds its stream" + ))); + } + previous_end = end; + } + Ok(()) +} + +fn validate_stream(info: &StreamInfo, raw_size: usize, label: &str) -> Result<()> { + if info.puff_length >= zucchini::OFFSET_BOUND { + return Err(Error::InvalidMetadata(format!( + "{label} puff length exceeds the maximum of {} bytes", + zucchini::OFFSET_BOUND - 1 + ))); + } + if info.deflates.len() != info.puffs.len() { + return Err(Error::InvalidMetadata(format!( + "{label} deflate and puff extent counts differ" + ))); + } + validate_extents(&info.deflates, "deflate", usize::MAX, |extent| { + (extent.offset, extent.length) + })?; + validate_extents(&info.puffs, "puff", info.puff_length, |extent| { + (extent.offset, extent.length) + })?; + + let computed_raw_size = + match (info.deflates.last(), info.puffs.last()) { + (Some(deflate), Some(puff)) => { + let deflate_end = deflate + .offset + .checked_add(deflate.length) + .ok_or_else(|| Error::InvalidMetadata("deflate end overflows".into()))?; + let puff_end = puff + .offset + .checked_add(puff.length) + .ok_or_else(|| Error::InvalidMetadata("puff end overflows".into()))?; + (deflate_end / 8) + .checked_add(info.puff_length.checked_sub(puff_end).ok_or_else(|| { + Error::InvalidMetadata("puff extent exceeds stream".into()) + })?) + .ok_or_else(|| Error::InvalidMetadata(format!("{label} raw size overflows")))? + } + (None, None) => info.puff_length, + _ => { + return Err(Error::InvalidMetadata(format!( + "{label} deflate and puff extent counts differ" + ))); + } + }; + if computed_raw_size != raw_size { + return Err(Error::InvalidMetadata(format!( + "{label} raw size mismatch: metadata describes {computed_raw_size} bytes, operation provides {raw_size}" + ))); + } + let raw_bits = raw_size + .checked_mul(8) + .ok_or_else(|| Error::InvalidMetadata(format!("{label} raw bit length overflows")))?; + for (index, extent) in info.deflates.iter().enumerate() { + let end = extent + .offset + .checked_add(extent.length) + .ok_or_else(|| Error::InvalidMetadata("deflate extent end overflows".into()))?; + if end > raw_bits { + return Err(Error::InvalidMetadata(format!( + "deflate extent {index} exceeds its stream" + ))); + } + } + Ok(()) +} + +fn sign_magnitude_usize(bytes: &[u8], label: &str) -> Result { + let encoded = u64::from_le_bytes( + bytes + .get(..8) + .ok_or_else(|| Error::Bsdiff(format!("truncated {label}")))? + .try_into() + .map_err(|_| Error::Bsdiff(format!("invalid {label}")))?, + ); + if encoded & (1 << 63) != 0 { + return Err(Error::Bsdiff(format!("negative {label}"))); + } + usize::try_from(encoded).map_err(|_| Error::Bsdiff(format!("{label} is too large"))) +} + +fn compression(value: u8, label: &str) -> Result { + match value { + 0 => Ok(Compression::None), + 1 => Ok(Compression::Bzip2), + 2 => Ok(Compression::Brotli), + _ => Err(Error::Bsdiff(format!("unsupported {label} compression algorithm {value}"))), + } +} + +fn check_decoded_bound( + mut reader: impl Read, + maximum: usize, + label: &str, + cancellation_token: &CancellationToken, +) -> Result<()> { + let mut decoded = 0usize; + let mut buffer = [0u8; DECODE_CHUNK_SIZE]; + loop { + check_cancelled(cancellation_token)?; + let remaining = maximum + .checked_sub(decoded) + .ok_or_else(|| Error::Bsdiff(format!("decoded {label} stream exceeds its bound")))?; + let read_size = buffer.len().min(remaining.saturating_add(1)); + let count = reader + .read(&mut buffer[..read_size]) + .map_err(|error| Error::Bsdiff(format!("unable to decode {label} stream: {error}")))?; + if count == 0 { + return Ok(()); + } + if count > remaining { + return Err(Error::Bsdiff(format!( + "decoded {label} stream exceeds its bound of {maximum} bytes" + ))); + } + decoded = decoded + .checked_add(count) + .ok_or_else(|| Error::Bsdiff(format!("decoded {label} length overflows")))?; + } +} + +fn validate_compressed_stream( + algorithm: Compression, + bytes: &[u8], + maximum: usize, + label: &str, + cancellation_token: &CancellationToken, +) -> Result<()> { + match algorithm { + Compression::None => { + if bytes.len() > maximum { + return Err(Error::Bsdiff(format!( + "uncompressed {label} stream exceeds its bound of {maximum} bytes" + ))); + } + Ok(()) + } + Compression::Bzip2 => { + check_decoded_bound(BzDecoder::new(bytes), maximum, label, cancellation_token) + } + Compression::Brotli => check_decoded_bound( + BrotliDecoder::new(bytes, DECODE_CHUNK_SIZE), + maximum, + label, + cancellation_token, + ), + } +} + +fn validate_inner_bsdiff_resources( + patch: &[u8], + output_size: usize, + cancellation_token: &CancellationToken, +) -> Result<()> { + crate::bsdiff::validate_output_len(patch, output_size) + .map_err(|error| Error::Bsdiff(error.to_string()))?; + let header = patch + .get(..32) + .ok_or_else(|| Error::Bsdiff("patch data is shorter than 32 bytes".into()))?; + let algorithms = if &header[..8] == b"BSDIFF40" { + [Compression::Bzip2; 3] + } else if &header[..5] == b"BSDF2" { + [ + compression(header[5], "control")?, + compression(header[6], "diff")?, + compression(header[7], "extra")?, + ] + } else { + return Err(Error::Bsdiff("invalid BSDIFF/BSDF2 magic header".into())); + }; + let control_length = sign_magnitude_usize(&header[8..16], "control stream length")?; + let diff_length = sign_magnitude_usize(&header[16..24], "diff stream length")?; + let control_end = 32usize + .checked_add(control_length) + .ok_or_else(|| Error::Bsdiff("control stream range overflows".into()))?; + let diff_end = control_end + .checked_add(diff_length) + .ok_or_else(|| Error::Bsdiff("diff stream range overflows".into()))?; + let control = patch + .get(32..control_end) + .ok_or_else(|| Error::Bsdiff("control stream exceeds patch".into()))?; + let diff = patch + .get(control_end..diff_end) + .ok_or_else(|| Error::Bsdiff("diff stream exceeds patch".into()))?; + let extra = patch + .get(diff_end..) + .ok_or_else(|| Error::Bsdiff("extra stream offset exceeds patch".into()))?; + let control_bound = + output_size.saturating_mul(48).saturating_add(4096).min(zucchini::OFFSET_BOUND - 1); + validate_compressed_stream( + algorithms[0], + control, + control_bound, + "control", + cancellation_token, + )?; + validate_compressed_stream(algorithms[1], diff, output_size, "diff", cancellation_token)?; + validate_compressed_stream(algorithms[2], extra, output_size, "extra", cancellation_token)?; + Ok(()) +} + +fn map_cancelled_to_zucchini(error: anyhow::Error) -> Error { + if is_cancellation(error.as_ref()) { + Error::Cancelled(ExtractionCancelled) + } else { + Error::Zucchini(error.to_string()) + } +} + +pub(crate) fn validate_bsdiff_resources( + patch: &[u8], + output_size: usize, + cancellation_token: &CancellationToken, +) -> Result<()> { + validate_inner_bsdiff_resources(patch, output_size, cancellation_token) +} + +pub fn apply( + source: &[u8], + patch: &[u8], + destination_size: usize, + cancellation_token: &CancellationToken, +) -> Result> { + check_cancelled(cancellation_token)?; + let prefix = patch + .get(..8) + .ok_or_else(|| Error::BadPatchHeader("patch is shorter than eight bytes".into()))?; + if &prefix[..4] != MAGIC { + return Err(Error::BadPatchHeader("missing PUF1 magic".into())); + } + let header_size = + usize::try_from(u32::from_be_bytes([prefix[4], prefix[5], prefix[6], prefix[7]])) + .map_err(|_| Error::BadPatchHeader("protobuf header is too large".into()))?; + let header_end = 8usize + .checked_add(header_size) + .ok_or_else(|| Error::BadPatchHeader("protobuf header range overflows".into()))?; + let header_bytes = patch + .get(8..header_end) + .ok_or_else(|| Error::BadPatchHeader("protobuf header overruns patch".into()))?; + let header = parse_header(header_bytes)?; + if !matches!(header.patch_type, PATCH_TYPE_BSDIFF | PATCH_TYPE_ZUCCHINI) { + return Err(Error::UnsupportedPatchType(header.patch_type)); + } + validate_stream(&header.source, source.len(), "source")?; + validate_stream(&header.destination, destination_size, "destination")?; + let raw_patch = patch + .get(header_end..) + .ok_or_else(|| Error::BadPatchHeader("inner patch offset is invalid".into()))?; + + check_cancelled(cancellation_token)?; + let puffed_source = stream::puff( + source, + &header.source.deflates, + &header.source.puffs, + header.source.puff_length, + cancellation_token, + )?; + + check_cancelled(cancellation_token)?; + let puffed_destination = if header.patch_type == PATCH_TYPE_BSDIFF { + crate::apply_bsdiff( + &puffed_source, + raw_patch, + header.destination.puff_length, + cancellation_token, + ) + .map_err(|error| { + if is_cancellation(error.as_ref()) { + Error::Cancelled(ExtractionCancelled) + } else { + Error::Bsdiff(error.to_string()) + } + })? + } else { + let zucchini_patch = decode_zucchini_patch(raw_patch, cancellation_token) + .map_err(map_cancelled_to_zucchini)?; + check_cancelled(cancellation_token)?; + let destination = + zucchini::apply(&puffed_source, &zucchini_patch, header.destination.puff_length) + .map_err(|error| Error::Zucchini(format!("apply failed: {error}")))?; + check_cancelled(cancellation_token)?; + destination + }; + if puffed_destination.len() != header.destination.puff_length { + return Err(Error::SizeMismatch { + expected: header.destination.puff_length, + actual: puffed_destination.len(), + }); + } + + check_cancelled(cancellation_token)?; + let destination = stream::huff( + &puffed_destination, + &header.destination.deflates, + &header.destination.puffs, + header.destination.puff_length, + destination_size, + cancellation_token, + )?; + check_cancelled(cancellation_token)?; + Ok(destination) +} diff --git a/src/puffin/puff_io.rs b/src/puffin/puff_io.rs new file mode 100644 index 0000000..f1138c4 --- /dev/null +++ b/src/puffin/puff_io.rs @@ -0,0 +1,337 @@ +// Copyright 2017 The ChromiumOS Authors +// Use of this source code is governed by a BSD-style license that can be +// found in the LICENSE file. +// +// Adapted from puffdiff 0.1.0, a pure-Rust port of ChromiumOS Puffin. + +//! The "puff" intermediate serialization, ported from puffin `puff_data.h`, +//! `puff_writer.cc` and `puff_reader.cc`. +//! +//! The byte layout produced here must match puffin exactly, because the bsdiff +//! patch inside a PUFFDIFF op is computed against puffin's puffed bytes. + +use crate::puffin::{Error, Result}; + +/// Maximum bytes a block-metadata blob can occupy: 1 header + 3 lengths + +/// 286 lit/len code lengths + 30 distance + 19 code-length codes. +pub const BLOCK_METADATA_MAX: usize = 1 + 3 + 286 + 30 + 19; + +const LITERALS_HEADER: u8 = 0x00; +const LEN_DIST_HEADER: u8 = 0x80; +const LITERALS_MAX_LENGTH: usize = (1 << 16) + 127; // 65663 + +/// One decoded token exchanged between puffer/huffer and the puff buffer. +pub enum PuffData<'a> { + Literal(u8), + Literals(&'a [u8]), + LenDist { length: usize, distance: usize }, + BlockMetadata(&'a [u8]), + EndOfBlock, +} + +#[inline] +fn write_u16_be(buf: &mut [u8], value: u16) -> Result<()> { + let target = + buf.get_mut(..2).ok_or_else(|| Error::Corrupt("puff writer u16 overflow".into()))?; + target.copy_from_slice(&value.to_be_bytes()); + Ok(()) +} + +#[inline] +fn read_u16_be(buf: &[u8]) -> Result { + let bytes: [u8; 2] = buf + .get(..2) + .ok_or_else(|| Error::Corrupt("puff reader u16 overflow".into()))? + .try_into() + .map_err(|_| Error::Corrupt("puff reader u16 conversion failed".into()))?; + Ok(u16::from_be_bytes(bytes)) +} + +#[derive(PartialEq)] +enum WState { + None, + Small, + Large, +} + +/// Serializes `PuffData` tokens into the puff byte format. +pub struct PuffWriter<'a> { + out: &'a mut [u8], + index: usize, + state: WState, + len_index: usize, + cur_literals_length: usize, +} + +impl<'a> PuffWriter<'a> { + pub fn new(out: &'a mut [u8]) -> Self { + PuffWriter { out, index: 0, state: WState::None, len_index: 0, cur_literals_length: 0 } + } + + fn insert_literal_bytes(&mut self, is_single: bool, byte: u8, bytes: &[u8]) -> Result<()> { + let length = if is_single { 1 } else { bytes.len() }; + if !is_single && length == 0 { + return Ok(()); + } + if self.state == WState::None { + self.need(1)?; + self.len_index = self.index; + self.index += 1; + self.state = WState::Small; + } + let combined_length = self + .cur_literals_length + .checked_add(length) + .ok_or_else(|| Error::Corrupt("puff literal length overflows".into()))?; + if combined_length > LITERALS_MAX_LENGTH { + return Err(Error::Corrupt("puff literal run is too long".into())); + } + if self.state == WState::Small && combined_length > 127 { + // Open two bytes of space for a large-literal length prefix by + // shifting the already-written literals forward. + self.need(2)?; + let literal_start = self + .len_index + .checked_add(1) + .ok_or_else(|| Error::Corrupt("puff literal offset overflows".into()))?; + let literal_end = literal_start + .checked_add(self.cur_literals_length) + .ok_or_else(|| Error::Corrupt("puff literal range overflows".into()))?; + let destination = self + .len_index + .checked_add(3) + .ok_or_else(|| Error::Corrupt("puff literal shift overflows".into()))?; + self.out.copy_within(literal_start..literal_end, destination); + self.index += 2; + self.state = WState::Large; + } + + self.need(length)?; + let end = self + .index + .checked_add(length) + .ok_or_else(|| Error::Corrupt("puff writer literal range overflows".into()))?; + if is_single { + self.out[self.index] = byte; + } else { + self.out[self.index..end].copy_from_slice(bytes); + } + self.index = end; + self.cur_literals_length = combined_length; + + if self.cur_literals_length == LITERALS_MAX_LENGTH { + self.flush_literals()?; + } + Ok(()) + } + + pub fn insert(&mut self, pd: &PuffData) -> Result<()> { + match pd { + PuffData::Literal(b) => self.insert_literal_bytes(true, *b, &[]), + PuffData::Literals(bytes) => self.insert_literal_bytes(false, 0, bytes), + PuffData::LenDist { length, distance } => { + self.flush_literals()?; + let (length, distance) = (*length, *distance); + if !(3..=258).contains(&length) || !(1..=32768).contains(&distance) { + return Err(Error::Corrupt("len/dist out of range".into())); + } + if length < 130 { + self.need(3)?; + self.out[self.index] = LEN_DIST_HEADER | (length - 3) as u8; + self.index += 1; + } else { + self.need(4)?; + self.out[self.index] = LEN_DIST_HEADER | 127; + self.index += 1; + self.out[self.index] = (length - 3 - 127) as u8; + self.index += 1; + } + write_u16_be(&mut self.out[self.index..], (distance - 1) as u16)?; + self.index += 2; + self.len_index = self.index; + self.state = WState::None; + Ok(()) + } + PuffData::BlockMetadata(meta) => { + self.flush_literals()?; + let length = meta.len(); + if length == 0 || length > BLOCK_METADATA_MAX { + return Err(Error::Corrupt("block metadata length invalid".into())); + } + self.need(length.checked_add(2).ok_or_else(|| { + Error::Corrupt("block metadata output length overflows".into()) + })?)?; + write_u16_be(&mut self.out[self.index..], (length - 1) as u16)?; + self.index += 2; + self.out[self.index..self.index + length].copy_from_slice(meta); + self.index += length; + self.len_index = self.index; + self.state = WState::None; + Ok(()) + } + PuffData::EndOfBlock => { + self.flush_literals()?; + self.need(2)?; + self.out[self.index] = LEN_DIST_HEADER | 127; + self.index += 1; + self.out[self.index] = (259 - 3 - 127) as u8; + self.index += 1; + self.len_index = self.index; + self.state = WState::None; + Ok(()) + } + } + } + + fn need(&self, count: usize) -> Result<()> { + let end = self + .index + .checked_add(count) + .ok_or_else(|| Error::Corrupt("puff writer range overflows".into()))?; + if end > self.out.len() { + return Err(Error::Corrupt("puff writer output overflow".into())); + } + Ok(()) + } + + fn flush_literals(&mut self) -> Result<()> { + if self.cur_literals_length == 0 { + return Ok(()); + } + match self.state { + WState::Small => { + if self.cur_literals_length != self.index - self.len_index - 1 { + return Err(Error::Corrupt("small literal length mismatch".into())); + } + self.out[self.len_index] = LITERALS_HEADER | (self.cur_literals_length - 1) as u8; + self.len_index = self.index; + self.state = WState::None; + } + WState::Large => { + if self.cur_literals_length != self.index - self.len_index - 3 { + return Err(Error::Corrupt("large literal length mismatch".into())); + } + self.out[self.len_index] = LITERALS_HEADER | 127; + write_u16_be( + &mut self.out[self.len_index + 1..], + (self.cur_literals_length - 127 - 1) as u16, + )?; + self.len_index = self.index; + self.state = WState::None; + } + WState::None => {} + } + self.cur_literals_length = 0; + Ok(()) + } + + pub fn flush(&mut self) -> Result<()> { + self.flush_literals() + } + + pub fn size(&self) -> usize { + self.index + } +} + +enum RState { + ReadingLenDist, + ReadingBlockMetadata, +} + +/// Deserializes puff bytes back into `PuffData` tokens. +pub struct PuffReader<'a> { + buf: &'a [u8], + index: usize, + state: RState, +} + +impl<'a> PuffReader<'a> { + pub fn new(buf: &'a [u8]) -> Self { + PuffReader { buf, index: 0, state: RState::ReadingBlockMetadata } + } + + pub fn bytes_left(&self) -> usize { + self.buf.len() - self.index + } + + fn need(&self, n: usize) -> Result<()> { + let end = self + .index + .checked_add(n) + .ok_or_else(|| Error::Corrupt("puff reader range overflows".into()))?; + if end > self.buf.len() { + return Err(Error::Corrupt("puff reader out of range".into())); + } + Ok(()) + } + + pub fn get_next(&mut self) -> Result> { + match self.state { + RState::ReadingLenDist => { + self.need(1)?; + let header = self.buf[self.index]; + if header & 0x80 != 0 { + // length/distance (or end of block) + let mut length; + if (header & 0x7F) < 127 { + length = (header & 0x7F) as usize; + } else { + self.index += 1; + self.need(1)?; + length = self.buf[self.index] as usize + 127; + } + length += 3; + if length > 259 { + return Err(Error::Corrupt("len/dist length too large".into())); + } + self.index += 1; + + if length == 259 { + self.state = RState::ReadingBlockMetadata; + return Ok(PuffData::EndOfBlock); + } + + self.need(2)?; + let distance = read_u16_be(&self.buf[self.index..])?; + if distance >= (1 << 15) { + return Err(Error::Corrupt("distance too large".into())); + } + let distance = distance as usize + 1; + self.index += 2; + Ok(PuffData::LenDist { length, distance }) + } else { + // literals + let mut length; + if (header & 0x7F) < 127 { + length = (header & 0x7F) as usize; + self.index += 1; + } else { + self.index += 1; + self.need(2)?; + length = usize::from(read_u16_be(&self.buf[self.index..])?) + 127; + self.index += 2; + } + length += 1; + self.need(length)?; + let slice = &self.buf[self.index..self.index + length]; + self.index += length; + Ok(PuffData::Literals(slice)) + } + } + RState::ReadingBlockMetadata => { + self.need(2)?; + let length = usize::from(read_u16_be(&self.buf[self.index..])?) + 1; + self.index += 2; + self.need(length)?; + if length > BLOCK_METADATA_MAX { + return Err(Error::Corrupt("block metadata too large".into())); + } + let slice = &self.buf[self.index..self.index + length]; + self.index += length; + self.state = RState::ReadingLenDist; + Ok(PuffData::BlockMetadata(slice)) + } + } + } +} diff --git a/src/puffin/puffer.rs b/src/puffin/puffer.rs new file mode 100644 index 0000000..aba3c9e --- /dev/null +++ b/src/puffin/puffer.rs @@ -0,0 +1,159 @@ +// Copyright 2017 The ChromiumOS Authors +// Use of this source code is governed by a BSD-style license that can be +// found in the LICENSE file. +// +// Adapted from puffdiff 0.1.0, a pure-Rust port of ChromiumOS Puffin. + +//! Deflate -> puff transform, ported from puffin `puffer.cc`. + +use crate::CancellationToken; +use crate::puffin::bit_io::BitReader; +use crate::puffin::huffman::{ + BlockType, DISTANCE_BASES, DISTANCE_EXTRA_BITS, HuffmanTable, LENGTH_BASES, LENGTH_EXTRA_BITS, +}; +use crate::puffin::puff_io::{BLOCK_METADATA_MAX, PuffData, PuffWriter}; +use crate::puffin::{Error, Result, check_cancelled}; + +/// Puff every deflate block in `br` into `pw`, until the input is exhausted. +pub fn puff_deflate( + br: &mut BitReader, + pw: &mut PuffWriter, + cancellation_token: &CancellationToken, +) -> Result<()> { + puff_deflate_impl(br, pw, cancellation_token) +} + +fn puff_deflate_impl( + br: &mut BitReader, + pw: &mut PuffWriter, + cancellation_token: &CancellationToken, +) -> Result<()> { + let mut fixed_ht = HuffmanTable::new()?; + let mut dyn_ht = HuffmanTable::new()?; + let mut fixed_built = false; + // Minimum deflate block is 8 bits (3-bit header + a 5-bit fixed symbol). + while br.cache_bits(8) { + check_cancelled(cancellation_token)?; + if !br.cache_bits(3) { + return Err(Error::Corrupt("eof reading block header".into())); + } + let final_bit = br.read_bits(1) as u8; + br.drop_bits(1); + let btype = br.read_bits(2) as u8; + br.drop_bits(2); + + let block_header = (final_bit << 7) | (btype << 5); + + let use_dynamic = match BlockType::from_bits(btype)? { + BlockType::Uncompressed => { + let skipped_bits = br.read_boundary_bits(); + br.skip_boundary_bits(); + if !br.cache_bits(32) { + return Err(Error::Corrupt("eof reading uncompressed lengths".into())); + } + let len = br.read_bits(16); + br.drop_bits(16); + let nlen = br.read_bits(16); + br.drop_bits(16); + if (len ^ nlen) != 0xFFFF { + return Err(Error::Corrupt("bad uncompressed LEN/NLEN".into())); + } + + let header = block_header | skipped_bits; + pw.insert(&PuffData::BlockMetadata(&[header]))?; + + let raw = br.get_bytes(len as usize)?; + pw.insert(&PuffData::Literals(raw))?; + pw.insert(&PuffData::EndOfBlock)?; + continue; + } + BlockType::Fixed => { + if !fixed_built { + fixed_ht.build_fixed()?; + fixed_built = true; + } + pw.insert(&PuffData::BlockMetadata(&[block_header]))?; + false + } + BlockType::Dynamic => { + let mut meta = [0u8; BLOCK_METADATA_MAX]; + meta[0] = block_header; + let written = dyn_ht.build_dynamic_from_deflate(br, &mut meta[1..])?; + pw.insert(&PuffData::BlockMetadata(&meta[..written + 1]))?; + true + } + }; + + loop { + check_cancelled(cancellation_token)?; + let cur_ht: &HuffmanTable = if use_dynamic { &dyn_ht } else { &fixed_ht }; + let mut max_bits = cur_ht.lit_len_max_bits(); + if !br.cache_bits(max_bits) { + max_bits = cur_ht.end_of_block_bit_length()?; + } + if !br.cache_bits(max_bits) { + return Err(Error::Corrupt("eof reading lit/len".into())); + } + let bits = br.read_bits(max_bits); + let (lit_len_alphabet, nbits) = cur_ht.lit_len_alphabet(bits)?; + if nbits > br.cached_bits() { + return Err(Error::Corrupt("truncated lit/len Huffman code".into())); + } + br.drop_bits(nbits); + + if lit_len_alphabet < 256 { + pw.insert(&PuffData::Literal(lit_len_alphabet as u8))?; + } else if lit_len_alphabet == 256 { + pw.insert(&PuffData::EndOfBlock)?; + break; + } else { + if lit_len_alphabet > 285 { + return Err(Error::Corrupt("lit/len alphabet > 285".into())); + } + let len_code_start = (lit_len_alphabet - 257) as usize; + let extra_len = LENGTH_EXTRA_BITS[len_code_start] as usize; + let mut extra_val = 0u32; + if extra_len != 0 { + if !br.cache_bits(extra_len) { + return Err(Error::Corrupt("eof reading length extra".into())); + } + extra_val = br.read_bits(extra_len); + br.drop_bits(extra_len); + } + let length = LENGTH_BASES[len_code_start] as usize + extra_val as usize; + + let mut bits_to_cache = cur_ht.distance_max_bits(); + if !br.cache_bits(bits_to_cache) { + // Rare legacy corner case (crbug.com/915559): not enough + // bits for a full-width distance code near the stream end. + bits_to_cache = br.bits_remaining() as usize; + if !br.cache_bits(bits_to_cache) { + return Err(Error::Corrupt("eof reading distance".into())); + } + } + let dbits = br.read_bits(bits_to_cache); + let (distance_alphabet, nbits) = cur_ht.distance_alphabet(dbits)?; + if nbits > br.cached_bits() { + return Err(Error::Corrupt("truncated distance Huffman code".into())); + } + br.drop_bits(nbits); + + let extra_len = DISTANCE_EXTRA_BITS[distance_alphabet as usize] as usize; + let mut extra_val = 0u32; + if extra_len != 0 { + if !br.cache_bits(extra_len) { + return Err(Error::Corrupt("eof reading distance extra".into())); + } + extra_val = br.read_bits(extra_len); + br.drop_bits(extra_len); + } + let distance = + DISTANCE_BASES[distance_alphabet as usize] as usize + extra_val as usize; + pw.insert(&PuffData::LenDist { length, distance })?; + } + } + } + pw.flush()?; + check_cancelled(cancellation_token)?; + Ok(()) +} diff --git a/src/puffin/stream.rs b/src/puffin/stream.rs new file mode 100644 index 0000000..4626b8c --- /dev/null +++ b/src/puffin/stream.rs @@ -0,0 +1,481 @@ +// Copyright 2017 The ChromiumOS Authors +// Use of this source code is governed by a BSD-style license that can be +// found in the LICENSE file. +// +// Adapted from puffdiff 0.1.0 and ChromiumOS Puffin's puffin_stream.cc. + +use crate::CancellationToken; +use crate::puffin::bit_io::{BitReader, BitWriter}; +use crate::puffin::huffer::huff_deflate; +use crate::puffin::puff_io::{PuffReader, PuffWriter}; +use crate::puffin::puffer::puff_deflate; +use crate::puffin::{BitExtent, ByteExtent, Error, Result, check_cancelled}; + +#[derive(Clone, Copy)] +struct Extent { + offset: usize, + length: usize, +} + +struct Driver<'a> { + source: Option<&'a [u8]>, + output: Vec, + output_position: usize, + deflates: Vec, + puffs: Vec, + upper_bounds: Vec, + current: usize, + puff_position: usize, + skip_bytes: usize, + deflate_bit_position: usize, + last_byte: u32, + extra_byte: usize, + max_puff_length: usize, +} + +fn allocate_zeroed(length: usize, label: &str) -> Result> { + let mut bytes = Vec::new(); + bytes + .try_reserve_exact(length) + .map_err(|error| Error::Allocation(format!("{label}: {error}")))?; + bytes.resize(length, 0); + Ok(bytes) +} + +fn checked_end(extent: Extent, label: &str) -> Result { + extent + .offset + .checked_add(extent.length) + .ok_or_else(|| Error::Corrupt(format!("{label} extent end overflows"))) +} + +impl<'a> Driver<'a> { + fn new( + source: Option<&'a [u8]>, + raw_size: usize, + puff_size: usize, + deflates: &[BitExtent], + puffs: &[ByteExtent], + ) -> Result { + if deflates.len() != puffs.len() { + return Err(Error::Corrupt("deflate and puff extent counts differ".into())); + } + + let mut driver_deflates = Vec::new(); + driver_deflates + .try_reserve_exact(deflates.len().saturating_add(1)) + .map_err(|error| Error::Allocation(format!("deflate stream metadata: {error}")))?; + driver_deflates.extend( + deflates.iter().map(|extent| Extent { offset: extent.offset, length: extent.length }), + ); + driver_deflates.push(Extent { + offset: raw_size + .checked_mul(8) + .ok_or_else(|| Error::Corrupt("raw stream bit length overflows".into()))?, + length: 0, + }); + + let mut driver_puffs = Vec::new(); + driver_puffs + .try_reserve_exact(puffs.len().saturating_add(1)) + .map_err(|error| Error::Allocation(format!("puff stream metadata: {error}")))?; + driver_puffs.extend( + puffs.iter().map(|extent| Extent { offset: extent.offset, length: extent.length }), + ); + driver_puffs.push(Extent { offset: puff_size, length: 0 }); + + let mut upper_bounds = Vec::new(); + upper_bounds + .try_reserve_exact(driver_puffs.len()) + .map_err(|error| Error::Allocation(format!("puff stream index: {error}")))?; + for puff in &driver_puffs { + upper_bounds.push(checked_end(*puff, "puff")?); + } + + let output = if source.is_some() { + Vec::new() + } else { + allocate_zeroed(raw_size, "raw Puffin output")? + }; + Ok(Self { + source, + output, + output_position: 0, + deflates: driver_deflates, + puffs: driver_puffs, + upper_bounds, + current: 0, + puff_position: 0, + skip_bytes: 0, + deflate_bit_position: 0, + last_byte: 0, + extra_byte: 0, + max_puff_length: puffs.iter().map(|extent| extent.length).max().unwrap_or(0), + }) + } + + fn seek_to_zero(&mut self, huffing: bool) -> Result<()> { + self.current = self + .upper_bounds + .iter() + .position(|&bound| bound > 0) + .ok_or_else(|| Error::Corrupt("empty Puffin stream index".into()))?; + let puff = self.puffs[self.current]; + let deflate = self.deflates[self.current]; + if puff.offset > 0 { + self.puff_position = 0; + let rounded_deflate = deflate + .offset + .checked_add(7) + .ok_or_else(|| Error::Corrupt("deflate offset rounding overflows".into()))? + / 8; + self.deflate_bit_position = rounded_deflate + .checked_sub(puff.offset) + .and_then(|offset| offset.checked_mul(8)) + .ok_or_else(|| Error::Corrupt("initial raw region is inconsistent".into()))?; + } else { + self.puff_position = puff.offset; + self.deflate_bit_position = deflate.offset; + } + if huffing { + self.set_extra_byte()?; + } + Ok(()) + } + + fn set_extra_byte(&mut self) -> Result<()> { + let deflate = *self + .deflates + .get(self.current) + .ok_or_else(|| Error::Corrupt("deflate index is out of range".into()))?; + let Some(next) = self.deflates.get(self.current + 1).copied() else { + self.extra_byte = 0; + return Ok(()); + }; + let end_bit = checked_end(deflate, "deflate")?; + let rounded_end = end_bit + .checked_add(7) + .ok_or_else(|| Error::Corrupt("deflate end rounding overflows".into()))? + & !7usize; + self.extra_byte = usize::from(end_bit & 7 != 0 && rounded_end <= next.offset); + Ok(()) + } + + fn source_slice(&self, range: std::ops::Range) -> Result<&'a [u8]> { + self.source + .and_then(|source| source.get(range)) + .ok_or_else(|| Error::Corrupt("read past end of raw stream".into())) + } + + fn write(&mut self, bytes: &[u8]) -> Result<()> { + let end = self + .output_position + .checked_add(bytes.len()) + .ok_or_else(|| Error::Corrupt("raw output position overflows".into()))?; + let destination = self + .output + .get_mut(self.output_position..end) + .ok_or_else(|| Error::Corrupt("write past end of raw stream".into()))?; + destination.copy_from_slice(bytes); + self.output_position = end; + Ok(()) + } + + fn run_puff( + &mut self, + output: &mut [u8], + cancellation_token: &CancellationToken, + ) -> Result<()> { + let mut bytes_read = 0usize; + while bytes_read < output.len() { + check_cancelled(cancellation_token)?; + let deflate = self.deflates[self.current]; + let puff = self.puffs[self.current]; + if self.puff_position < puff.offset { + let start_byte = self.deflate_bit_position / 8; + let end_byte = + deflate.offset.checked_add(7).ok_or_else(|| { + Error::Corrupt("deflate offset rounding overflows".into()) + })? / 8; + let available = end_byte + .checked_sub(start_byte) + .ok_or_else(|| Error::Corrupt("raw region runs backward".into()))?; + let count = (output.len() - bytes_read).min(available); + if count == 0 { + return Err(Error::Corrupt("raw region makes no progress".into())); + } + let source_end = start_byte + .checked_add(count) + .ok_or_else(|| Error::Corrupt("raw source range overflows".into()))?; + let destination_end = bytes_read + .checked_add(count) + .ok_or_else(|| Error::Corrupt("puff output range overflows".into()))?; + output[bytes_read..destination_end] + .copy_from_slice(self.source_slice(start_byte..source_end)?); + if source_end + .checked_mul(8) + .ok_or_else(|| Error::Corrupt("raw bit position overflows".into()))? + > deflate.offset + { + let mask = ((1u16 << (deflate.offset & 7)) - 1) as u8; + output[destination_end - 1] &= mask; + } + if start_byte * 8 < self.deflate_bit_position { + output[bytes_read] >>= self.deflate_bit_position & 7; + } + self.deflate_bit_position = (self.deflate_bit_position & !7usize) + .checked_add(count * 8) + .ok_or_else(|| Error::Corrupt("raw bit position overflows".into()))? + .min(deflate.offset); + bytes_read = destination_end; + self.puff_position = self + .puff_position + .checked_add(count) + .ok_or_else(|| Error::Corrupt("puff position overflows".into()))?; + if self.puff_position > puff.offset { + return Err(Error::Corrupt("raw region overruns puff extent".into())); + } + continue; + } + + let deflate_end = checked_end(deflate, "deflate")?; + let start_byte = deflate.offset / 8; + let end_byte = deflate_end + .checked_add(7) + .ok_or_else(|| Error::Corrupt("deflate end rounding overflows".into()))? + / 8; + let deflate_bytes = self.source_slice(start_byte..end_byte)?; + let direct = self.skip_bytes == 0 && output.len() - bytes_read >= puff.length; + let mut temporary = Vec::new(); + let target = if direct { + let end = bytes_read + .checked_add(puff.length) + .ok_or_else(|| Error::Corrupt("puff output range overflows".into()))?; + &mut output[bytes_read..end] + } else { + temporary = allocate_zeroed(puff.length, "temporary puff buffer")?; + &mut temporary + }; + let mut bit_reader = BitReader::new(deflate_bytes); + let skipped_bits = deflate.offset & 7; + if !bit_reader.cache_bits(skipped_bits) { + return Err(Error::Corrupt("deflate extent starts past its bytes".into())); + } + bit_reader.drop_bits(skipped_bits); + let mut puff_writer = PuffWriter::new(target); + puff_deflate(&mut bit_reader, &mut puff_writer, cancellation_token)?; + let expected_bits = skipped_bits + .checked_add(deflate.length) + .ok_or_else(|| Error::Corrupt("deflate bit length overflows".into()))?; + if bit_reader.offset_in_bits()? != expected_bits + || bit_reader.offset() != deflate_bytes.len() + || puff_writer.size() != puff.length + { + return Err(Error::Corrupt("puffed extent size does not match metadata".into())); + } + let copy_count = (output.len() - bytes_read).min( + puff.length + .checked_sub(self.skip_bytes) + .ok_or_else(|| Error::Corrupt("puff skip exceeds extent".into()))?, + ); + if !direct { + let copy_end = self + .skip_bytes + .checked_add(copy_count) + .ok_or_else(|| Error::Corrupt("puff copy range overflows".into()))?; + output[bytes_read..bytes_read + copy_count] + .copy_from_slice(&temporary[self.skip_bytes..copy_end]); + } + self.skip_bytes += copy_count; + bytes_read += copy_count; + if self + .puff_position + .checked_add(self.skip_bytes) + .ok_or_else(|| Error::Corrupt("puff position overflows".into()))? + == checked_end(puff, "puff")? + { + self.puff_position += self.skip_bytes; + self.skip_bytes = 0; + self.deflate_bit_position = deflate_end; + self.current += 1; + if self.current >= self.puffs.len() { + break; + } + } else if copy_count == 0 { + return Err(Error::Corrupt("puff extent makes no progress".into())); + } + } + if bytes_read != output.len() { + return Err(Error::Corrupt("puffer did not fill its output".into())); + } + Ok(()) + } + + fn run_huff(&mut self, input: &[u8], cancellation_token: &CancellationToken) -> Result<()> { + let puff_buffer_size = self + .max_puff_length + .checked_add(1) + .ok_or_else(|| Error::Corrupt("puff buffer size overflows".into()))?; + let mut puff_buffer = allocate_zeroed(puff_buffer_size, "huffer puff buffer")?; + let mut bytes_written = 0usize; + while bytes_written < input.len() { + check_cancelled(cancellation_token)?; + let deflate = self.deflates[self.current]; + let puff = self.puffs[self.current]; + if self.deflate_bit_position < (deflate.offset & !7usize) { + if self.deflate_bit_position & 7 != 0 { + return Err(Error::Corrupt("raw output is not byte-aligned".into())); + } + let raw_count = (deflate.offset / 8) + .checked_sub(self.deflate_bit_position / 8) + .ok_or_else(|| Error::Corrupt("raw output region runs backward".into()))?; + let copy_count = raw_count.min(input.len() - bytes_written); + if copy_count == 0 { + return Err(Error::Corrupt("raw output makes no progress".into())); + } + self.write(&input[bytes_written..bytes_written + copy_count])?; + bytes_written += copy_count; + self.puff_position += copy_count; + self.deflate_bit_position += copy_count * 8; + continue; + } + + if self.deflate_bit_position < deflate.offset { + let byte = *input + .get(bytes_written) + .ok_or_else(|| Error::Corrupt("missing shared boundary byte".into()))?; + self.last_byte |= u32::from(byte) << (self.deflate_bit_position & 7); + bytes_written += 1; + self.skip_bytes = 0; + self.deflate_bit_position = deflate.offset; + self.puff_position += 1; + if self.puff_position != puff.offset { + return Err(Error::Corrupt("puff boundary position does not match".into())); + } + } + + let required = puff + .length + .checked_add(self.extra_byte) + .and_then(|length| length.checked_sub(self.skip_bytes)) + .ok_or_else(|| Error::Corrupt("puff buffer accounting is invalid".into()))?; + let copy_count = (input.len() - bytes_written).min(required); + let copy_end = self + .skip_bytes + .checked_add(copy_count) + .ok_or_else(|| Error::Corrupt("puff buffer range overflows".into()))?; + let destination = puff_buffer + .get_mut(self.skip_bytes..copy_end) + .ok_or_else(|| Error::Corrupt("puff buffer is too small".into()))?; + destination.copy_from_slice(&input[bytes_written..bytes_written + copy_count]); + self.skip_bytes = copy_end; + bytes_written += copy_count; + + if self.skip_bytes == puff.length + self.extra_byte { + let deflate_end = checked_end(deflate, "deflate")?; + let start_byte = deflate.offset / 8; + let end_byte = deflate_end + .checked_add(7) + .ok_or_else(|| Error::Corrupt("deflate end rounding overflows".into()))? + / 8; + let mut deflate_buffer = allocate_zeroed( + end_byte + .checked_sub(start_byte) + .ok_or_else(|| Error::Corrupt("deflate byte range runs backward".into()))?, + "huffed deflate buffer", + )?; + let mut count = deflate_buffer.len(); + { + let mut bit_writer = BitWriter::new(&mut deflate_buffer); + bit_writer.write_bits(deflate.offset & 7, self.last_byte)?; + self.last_byte = 0; + let puff_bytes = puff_buffer + .get(..puff.length) + .ok_or_else(|| Error::Corrupt("puff extent exceeds buffer".into()))?; + let mut puff_reader = PuffReader::new(puff_bytes); + huff_deflate(&mut puff_reader, &mut bit_writer, cancellation_token)?; + let expected_bits = (deflate.offset & 7) + .checked_add(deflate.length) + .ok_or_else(|| Error::Corrupt("deflate bit length overflows".into()))?; + if bit_writer.offset_in_bits()? != expected_bits + || puff_reader.bytes_left() != 0 + { + return Err(Error::Corrupt( + "huffed deflate bit length does not match metadata".into(), + )); + } + bit_writer.flush()?; + if bit_writer.size() != count { + return Err(Error::Corrupt( + "huffed deflate size does not match metadata".into(), + )); + } + } + self.deflate_bit_position = deflate_end; + if self.extra_byte == 1 { + let last = deflate_buffer + .last_mut() + .ok_or_else(|| Error::Corrupt("empty deflate buffer".into()))?; + *last |= puff_buffer[puff.length] << (self.deflate_bit_position & 7); + self.deflate_bit_position = self + .deflate_bit_position + .checked_add(7) + .ok_or_else(|| Error::Corrupt("deflate position overflows".into()))? + & !7usize; + } else if self.deflate_bit_position & 7 != 0 { + self.last_byte = u32::from( + *deflate_buffer + .last() + .ok_or_else(|| Error::Corrupt("empty deflate buffer".into()))?, + ); + count -= 1; + } + self.write(&deflate_buffer[..count])?; + self.puff_position += self.skip_bytes; + self.skip_bytes = 0; + self.current += 1; + if self.current >= self.puffs.len() { + break; + } + self.set_extra_byte()?; + } else if copy_count == 0 { + return Err(Error::Corrupt("huffer makes no progress".into())); + } + } + if bytes_written != input.len() || self.output_position != self.output.len() { + return Err(Error::Corrupt("huffer output size does not match destination".into())); + } + Ok(()) + } +} + +pub fn puff( + source: &[u8], + deflates: &[BitExtent], + puffs: &[ByteExtent], + puff_size: usize, + cancellation_token: &CancellationToken, +) -> Result> { + let mut driver = Driver::new(Some(source), source.len(), puff_size, deflates, puffs)?; + driver.seek_to_zero(false)?; + let mut output = allocate_zeroed(puff_size, "puffed source")?; + driver.run_puff(&mut output, cancellation_token)?; + Ok(output) +} + +pub fn huff( + puffed: &[u8], + deflates: &[BitExtent], + puffs: &[ByteExtent], + puff_size: usize, + raw_size: usize, + cancellation_token: &CancellationToken, +) -> Result> { + if puffed.len() != puff_size { + return Err(Error::SizeMismatch { expected: puff_size, actual: puffed.len() }); + } + let mut driver = Driver::new(None, raw_size, puff_size, deflates, puffs)?; + driver.seek_to_zero(true)?; + driver.run_huff(puffed, cancellation_token)?; + Ok(driver.output) +} diff --git a/src/python.rs b/src/python.rs index 3173fb8..5ee7f16 100644 --- a/src/python.rs +++ b/src/python.rs @@ -1,15 +1,38 @@ use std::path::PathBuf; +use std::time::Duration; use pyo3::create_exception; -use pyo3::exceptions::PyException; +use pyo3::exceptions::{PyException, PyKeyboardInterrupt}; use pyo3::prelude::*; -use crate::ExtractOptions; +use crate::{CancellationToken, ExtractOptions, is_cancellation}; create_exception!(otadump, OtaDumpError, PyException); +#[pyclass(name = "CancellationToken")] +struct PyCancellationToken { + token: CancellationToken, +} + +#[pymethods] +impl PyCancellationToken { + #[new] + fn new() -> Self { + Self { token: CancellationToken::new() } + } + + fn cancel(&self) { + self.token.cancel(); + } + + fn is_cancelled(&self) -> bool { + self.token.is_cancelled() + } +} + /// Extract partitions from an Android OTA payload or OTA ZIP. #[pyfunction] +#[allow(clippy::too_many_arguments)] #[pyo3(signature = ( payload_file, output_dir, @@ -18,6 +41,8 @@ create_exception!(otadump, OtaDumpError, PyException); overwrite = false, partitions = None, verify = true, + source_dir = None, + cancellation_token = None, ))] fn extract( py: Python<'_>, @@ -27,9 +52,13 @@ fn extract( overwrite: bool, partitions: Option>, verify: bool, + source_dir: Option, + cancellation_token: Option>, ) -> PyResult<()> { + let cancellation_token = + cancellation_token.map(|token| token.token.clone()).unwrap_or_default(); let mut options = ExtractOptions::new(); - options.overwrite(overwrite).verify(verify); + options.overwrite(overwrite).verify(verify).cancellation_token(&cancellation_token); if let Some(num_threads) = num_threads { options.num_threads(num_threads); @@ -37,17 +66,42 @@ fn extract( if let Some(partitions) = partitions { options.partitions(partitions); } + if let Some(source_dir) = source_dir { + options.source_dir(source_dir); + } + + let extraction = std::thread::spawn(move || match options.extract(payload_file, output_dir) { + Ok(()) => Ok(()), + Err(error) if is_cancellation(error.as_ref()) => Err(None), + Err(error) => Err(Some(error.to_string())), + }); - py.detach(move || { - options - .extract(payload_file, output_dir) - .map_err(|error| OtaDumpError::new_err(error.to_string())) - }) + let mut signal_error = None; + while !extraction.is_finished() { + if signal_error.is_none() { + if let Err(error) = py.check_signals() { + cancellation_token.cancel(); + signal_error = Some(error); + } + } + py.detach(|| std::thread::sleep(Duration::from_millis(20))); + } + let result = + extraction.join().map_err(|_| OtaDumpError::new_err("Extraction worker panicked"))?; + if let Some(error) = signal_error { + return Err(error); + } + match result { + Ok(()) => Ok(()), + Err(None) => Err(PyKeyboardInterrupt::new_err("Extraction cancelled")), + Err(Some(error)) => Err(OtaDumpError::new_err(error)), + } } #[pymodule] fn otadump(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add("OtaDumpError", module.py().get_type::())?; + module.add_class::()?; module.add_function(wrap_pyfunction!(extract, module)?)?; Ok(()) } diff --git a/src/verity.rs b/src/verity.rs new file mode 100644 index 0000000..6708f5e --- /dev/null +++ b/src/verity.rs @@ -0,0 +1,436 @@ +use std::ops::Range; + +use anyhow::{Context as _, Result, ensure}; +use ring::digest; + +use crate::CancellationToken; +use crate::chromeos_update_engine::{Extent, PartitionUpdate}; + +const SHA256_SIZE: usize = 32; +const FEC_BLOCK_SIZE: usize = 4096; +const FEC_SYMBOLS: usize = 255; + +pub(crate) struct VerityConfig { + hash_tree: Option, + fec: Option, +} + +struct HashTreeConfig { + data: Range, + tree: Range, + salt: Vec, +} + +struct FecConfig { + data: Range, + storage: Range, + roots: usize, + rounds: usize, +} + +pub(crate) fn validate( + update: &PartitionUpdate, + block_size: usize, + partition_len: usize, + operation_destinations: &[(Range, usize)], +) -> Result> { + let hash_tree = validate_hash_tree(update, block_size, partition_len, operation_destinations)?; + let fec = validate_fec( + update, + block_size, + partition_len, + operation_destinations, + hash_tree.as_ref(), + )?; + if hash_tree.is_none() && fec.is_none() { + return Ok(None); + } + Ok(Some(VerityConfig { hash_tree, fec })) +} + +fn validate_hash_tree( + update: &PartitionUpdate, + block_size: usize, + partition_len: usize, + operation_destinations: &[(Range, usize)], +) -> Result> { + let data_blocks = + extent_num_blocks(update.hash_tree_data_extent.as_ref(), "Hash tree data")?.unwrap_or(0); + let tree_blocks = + extent_num_blocks(update.hash_tree_extent.as_ref(), "Hash tree")?.unwrap_or(0); + ensure!( + (data_blocks == 0) == (tree_blocks == 0), + "Hash tree data and storage extents must both be nonempty" + ); + if tree_blocks == 0 { + return Ok(None); + } + + let algorithm = update.hash_tree_algorithm.as_deref().unwrap_or_default(); + let is_sha256 = + algorithm.eq_ignore_ascii_case("sha256") || algorithm.eq_ignore_ascii_case("sha-256"); + ensure!(is_sha256, "Unsupported verity hash algorithm: {algorithm:?}"); + ensure!( + block_size > SHA256_SIZE * 2 && block_size % SHA256_SIZE == 0, + "Block size {block_size} is not supported for SHA-256 verity trees" + ); + + let data_extent = + update.hash_tree_data_extent.as_ref().context("Hash tree data extent is missing")?; + let tree_extent = + update.hash_tree_extent.as_ref().context("Hash tree storage extent is missing")?; + let data = extent_range(data_extent, block_size, partition_len, "Hash tree data")?; + let tree = extent_range(tree_extent, block_size, partition_len, "Hash tree")?; + ensure!(!ranges_overlap(&data, &tree), "Hash tree data and storage extents overlap"); + ensure!( + operation_destinations.iter().all(|(destination, _)| !ranges_overlap(&tree, destination)), + "Hash tree extent overlaps destination extents" + ); + + let hashes_per_block = block_size / SHA256_SIZE; + let mut level_blocks = div_ceil(data_blocks, hashes_per_block)?; + let mut required_blocks = 0usize; + loop { + required_blocks = + required_blocks.checked_add(level_blocks).context("Hash tree size overflow")?; + if level_blocks == 1 { + break; + } + level_blocks = div_ceil(level_blocks, hashes_per_block)?; + } + let required_size = + required_blocks.checked_mul(block_size).context("Hash tree size overflow")?; + ensure!( + tree.len() == required_size, + "Hash tree size mismatch: expected {required_size} bytes, got {}", + tree.len() + ); + + Ok(Some(HashTreeConfig { data, tree, salt: update.hash_tree_salt.clone().unwrap_or_default() })) +} + +pub(crate) fn generate( + partition: &mut [u8], + config: &VerityConfig, + block_size: usize, + cancellation_token: &CancellationToken, +) -> Result<()> { + if let Some(hash_tree) = config.hash_tree.as_ref() { + generate_hash_tree(partition, hash_tree, block_size, cancellation_token)?; + } + if let Some(fec) = config.fec.as_ref() { + generate_fec(partition, fec, block_size, cancellation_token)?; + } + Ok(()) +} + +fn generate_hash_tree( + partition: &mut [u8], + config: &HashTreeConfig, + block_size: usize, + cancellation_token: &CancellationToken, +) -> Result<()> { + cancellation_token.check()?; + let hashes_per_block = block_size / SHA256_SIZE; + let data_blocks = config.data.len() / block_size; + let base_blocks = div_ceil(data_blocks, hashes_per_block)?; + let mut level = zeroed_level(base_blocks, block_size)?; + for (block, hash_slot) in partition[config.data.clone()] + .chunks_exact(block_size) + .zip(level.chunks_exact_mut(SHA256_SIZE)) + { + cancellation_token.check()?; + hash_slot.copy_from_slice(salted_hash(&config.salt, block).as_ref()); + } + + let mut levels = vec![level]; + while levels.last().expect("base level exists").len() > block_size { + let current = levels.last().expect("base level exists"); + let next_blocks = div_ceil(current.len() / block_size, hashes_per_block)?; + let mut next = zeroed_level(next_blocks, block_size)?; + for (block, hash_slot) in + current.chunks_exact(block_size).zip(next.chunks_exact_mut(SHA256_SIZE)) + { + cancellation_token.check()?; + hash_slot.copy_from_slice(salted_hash(&config.salt, block).as_ref()); + } + levels.push(next); + } + + let tree = &mut partition[config.tree.clone()]; + let mut offset = 0usize; + for level in levels.iter().rev() { + for block in level.chunks_exact(block_size) { + cancellation_token.check()?; + let end = offset.checked_add(block_size).context("Hash tree write overflow")?; + tree[offset..end].copy_from_slice(block); + offset = end; + } + } + ensure!(offset == tree.len(), "Generated hash tree size mismatch"); + cancellation_token.check() +} + +fn validate_fec( + update: &PartitionUpdate, + block_size: usize, + partition_len: usize, + operation_destinations: &[(Range, usize)], + hash_tree: Option<&HashTreeConfig>, +) -> Result> { + let data_blocks = extent_num_blocks(update.fec_data_extent.as_ref(), "FEC data")?.unwrap_or(0); + let storage_blocks = extent_num_blocks(update.fec_extent.as_ref(), "FEC")?.unwrap_or(0); + ensure!( + (data_blocks == 0) == (storage_blocks == 0), + "FEC data and storage extents must both be nonempty" + ); + if storage_blocks == 0 { + return Ok(None); + } + + ensure!(block_size == FEC_BLOCK_SIZE, "FEC requires 4096-byte blocks"); + let roots = + usize::try_from(update.fec_roots.unwrap_or(2)).context("FEC roots are too large")?; + ensure!((1..FEC_SYMBOLS).contains(&roots), "FEC roots must be between 1 and 254"); + let data = extent_range( + update.fec_data_extent.as_ref().context("FEC data extent is missing")?, + block_size, + partition_len, + "FEC data", + )?; + let storage = extent_range( + update.fec_extent.as_ref().context("FEC storage extent is missing")?, + block_size, + partition_len, + "FEC storage", + )?; + ensure!(storage.start >= data.end, "FEC storage extent must follow the protected data extent"); + ensure!( + operation_destinations + .iter() + .all(|(destination, _)| !ranges_overlap(&storage, destination)), + "FEC storage extent overlaps destination extents" + ); + if let Some(hash_tree) = hash_tree { + ensure!( + !ranges_overlap(&storage, &hash_tree.tree), + "FEC storage extent overlaps hash tree extent" + ); + ensure!( + !ranges_overlap(&storage, &hash_tree.data), + "FEC storage extent overlaps hash tree data extent" + ); + } + + let rs_n = FEC_SYMBOLS - roots; + let rounds = div_ceil(data_blocks, rs_n)?; + let required_blocks = rounds.checked_mul(roots).context("FEC size overflow")?; + let required_size = required_blocks.checked_mul(block_size).context("FEC size overflow")?; + ensure!( + storage.len() == required_size, + "FEC size mismatch: expected {required_size} bytes, got {}", + storage.len() + ); + + Ok(Some(FecConfig { data, storage, roots, rounds })) +} + +fn generate_fec( + partition: &mut [u8], + config: &FecConfig, + block_size: usize, + cancellation_token: &CancellationToken, +) -> Result<()> { + cancellation_token.check()?; + let rs = ReedSolomon::new(config.roots)?; + let rs_n = FEC_SYMBOLS - config.roots; + let matrix_len = block_size.checked_mul(rs_n).context("FEC matrix size overflow")?; + let parity_len = block_size.checked_mul(config.roots).context("FEC parity size overflow")?; + let mut matrix = zeroed_buffer(matrix_len, "FEC matrix")?; + let mut parity = zeroed_buffer(parity_len, "FEC parity")?; + let data_blocks = config.data.len() / block_size; + + for round in 0..config.rounds { + cancellation_token.check()?; + matrix.fill(0); + for symbol in 0..rs_n { + if symbol % 32 == 0 { + cancellation_token.check()?; + } + let block_index = symbol + .checked_mul(config.rounds) + .and_then(|value| value.checked_add(round)) + .context("FEC interleave offset overflow")?; + if block_index >= data_blocks { + continue; + } + let source_start = config + .data + .start + .checked_add( + block_index.checked_mul(block_size).context("FEC data offset overflow")?, + ) + .context("FEC data offset overflow")?; + let source_end = + source_start.checked_add(block_size).context("FEC data end overflow")?; + let source = &partition[source_start..source_end]; + for (byte_index, byte) in source.iter().copied().enumerate() { + matrix[byte_index * rs_n + symbol] = byte; + } + } + + parity.fill(0); + for byte_index in 0..block_size { + if byte_index % 256 == 0 { + cancellation_token.check()?; + } + let data_start = byte_index * rs_n; + let parity_start = byte_index * config.roots; + rs.encode( + &matrix[data_start..data_start + rs_n], + &mut parity[parity_start..parity_start + config.roots], + ); + } + + cancellation_token.check()?; + let destination_start = config + .storage + .start + .checked_add(round.checked_mul(parity_len).context("FEC write offset overflow")?) + .context("FEC write offset overflow")?; + let destination_end = + destination_start.checked_add(parity_len).context("FEC write end overflow")?; + partition[destination_start..destination_end].copy_from_slice(&parity); + } + cancellation_token.check() +} + +struct ReedSolomon { + alpha_to: [u8; FEC_SYMBOLS + 1], + index_of: [u8; FEC_SYMBOLS + 1], + generator: Vec, + roots: usize, +} + +impl ReedSolomon { + fn new(roots: usize) -> Result { + let mut alpha_to = [0u8; FEC_SYMBOLS + 1]; + let mut index_of = [0u8; FEC_SYMBOLS + 1]; + index_of[0] = FEC_SYMBOLS as u8; + let mut symbol = 1u16; + for (exponent, alpha) in alpha_to.iter_mut().take(FEC_SYMBOLS).enumerate() { + index_of[usize::from(symbol)] = exponent as u8; + *alpha = symbol as u8; + symbol <<= 1; + if symbol & 0x100 != 0 { + symbol ^= 0x11d; + } + symbol &= 0xff; + } + ensure!(symbol == 1, "FEC field polynomial is not primitive"); + + let mut generator = zeroed_buffer(roots + 1, "FEC generator polynomial")?; + generator[0] = 1; + for root in 0..roots { + generator[root + 1] = 1; + for coefficient in (1..=root).rev() { + generator[coefficient] = if generator[coefficient] == 0 { + generator[coefficient - 1] + } else { + let exponent = usize::from(index_of[usize::from(generator[coefficient])]); + generator[coefficient - 1] ^ alpha_to[(exponent + root) % FEC_SYMBOLS] + }; + } + let exponent = usize::from(index_of[usize::from(generator[0])]); + generator[0] = alpha_to[(exponent + root) % FEC_SYMBOLS]; + } + for coefficient in &mut generator { + *coefficient = index_of[usize::from(*coefficient)]; + } + Ok(Self { alpha_to, index_of, generator, roots }) + } + + fn encode(&self, data: &[u8], parity: &mut [u8]) { + parity.fill(0); + for byte in data { + let feedback = usize::from(self.index_of[usize::from(*byte ^ parity[0])]); + if feedback != FEC_SYMBOLS { + for (index, parity_byte) in parity.iter_mut().enumerate().take(self.roots).skip(1) { + let generator = usize::from(self.generator[self.roots - index]); + *parity_byte ^= self.alpha_to[(feedback + generator) % FEC_SYMBOLS]; + } + } + parity.copy_within(1..self.roots, 0); + parity[self.roots - 1] = if feedback == FEC_SYMBOLS { + 0 + } else { + let generator = usize::from(self.generator[0]); + self.alpha_to[(feedback + generator) % FEC_SYMBOLS] + }; + } + } +} + +fn extent_num_blocks(extent: Option<&Extent>, label: &str) -> Result> { + extent + .map(|extent| { + usize::try_from( + extent + .num_blocks + .with_context(|| format!("{label} extent block count is missing"))?, + ) + .with_context(|| format!("{label} extent block count is too large")) + }) + .transpose() +} + +fn extent_range( + extent: &Extent, + block_size: usize, + partition_len: usize, + label: &str, +) -> Result> { + let start_block = usize::try_from( + extent.start_block.with_context(|| format!("{label} extent start block is missing"))?, + ) + .with_context(|| format!("{label} extent start block is too large"))?; + let num_blocks = usize::try_from( + extent.num_blocks.with_context(|| format!("{label} extent block count is missing"))?, + ) + .with_context(|| format!("{label} extent block count is too large"))?; + let start = start_block.checked_mul(block_size).context("Verity extent offset overflow")?; + let len = num_blocks.checked_mul(block_size).context("Verity extent length overflow")?; + let end = start.checked_add(len).context("Verity extent end overflow")?; + ensure!(end <= partition_len, "{label} extent exceeds partition size"); + Ok(start..end) +} + +fn div_ceil(value: usize, divisor: usize) -> Result { + value + .checked_add(divisor - 1) + .context("Hash tree block count overflow") + .map(|rounded| rounded / divisor) +} + +fn zeroed_level(blocks: usize, block_size: usize) -> Result> { + let len = blocks.checked_mul(block_size).context("Hash tree level size overflow")?; + zeroed_buffer(len, "hash tree") +} + +fn zeroed_buffer(len: usize, label: &str) -> Result> { + let mut level = Vec::new(); + level.try_reserve_exact(len).with_context(|| format!("Unable to allocate {label}"))?; + level.resize(len, 0); + Ok(level) +} + +fn salted_hash(salt: &[u8], block: &[u8]) -> digest::Digest { + let mut context = digest::Context::new(&digest::SHA256); + context.update(salt); + context.update(block); + context.finish() +} + +fn ranges_overlap(left: &Range, right: &Range) -> bool { + left.start < right.end && right.start < left.end +} diff --git a/src/zucchini.rs b/src/zucchini.rs new file mode 100644 index 0000000..929b2bb --- /dev/null +++ b/src/zucchini.rs @@ -0,0 +1,155 @@ +use std::error; +use std::fmt; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Status { + InvalidArgument, + InvalidPatch, + UnsupportedElement, + WrongOutputSize, + ApplyError, + Cancelled, + AllocationFailure, + UnsupportedTarget, + Unknown(i32), +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct Error { + status: Status, + message: String, +} + +impl Error { + pub fn status(&self) -> Status { + self.status + } +} +impl fmt::Display for Error { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.message) + } +} +impl error::Error for Error {} +pub type Result = std::result::Result; + +/// Exclusive upper bound for buffers represented by Zucchini offsets. +pub const OFFSET_BOUND: usize = (u32::MAX / 2) as usize; + +/// Applies a Zucchini patch into a newly allocated, exact-size output buffer. +/// +pub fn apply(old: &[u8], patch: &[u8], output_size: usize) -> Result> { + if old.len() >= OFFSET_BOUND || patch.len() >= OFFSET_BOUND || output_size >= OFFSET_BOUND { + return Err(Error { + status: Status::InvalidArgument, + message: "image exceeds Zucchini offset bound".into(), + }); + } + let mut output = Vec::new(); + output.try_reserve_exact(output_size).map_err(|e| Error { + status: Status::AllocationFailure, + message: format!("unable to allocate Zucchini output: {e}"), + })?; + output.resize(output_size, 0); + #[cfg(otadump_zucchini)] + { + use std::ffi::{CStr, c_char}; + #[repr(C)] + struct NativeResult { + status: i32, + error: *const c_char, + } + unsafe extern "C" { + fn otadump_zucchini_apply( + old: *const u8, + old_len: usize, + patch: *const u8, + patch_len: usize, + new: *mut u8, + new_len: usize, + ) -> NativeResult; + } + // SAFETY: all slices remain alive and output is a distinct allocation for the call. + let result = unsafe { + otadump_zucchini_apply( + old.as_ptr(), + old.len(), + patch.as_ptr(), + patch.len(), + output.as_mut_ptr(), + output.len(), + ) + }; + if result.status == 0 { + return Ok(output); + } + let message = if result.error.is_null() { + format!("Zucchini failed with status {}", result.status) + } else { + unsafe { CStr::from_ptr(result.error) }.to_string_lossy().into_owned() + }; + Err(Error { + status: match result.status { + 1 => Status::InvalidArgument, + 2 => Status::InvalidPatch, + 3 => Status::UnsupportedElement, + 4 => Status::WrongOutputSize, + 5 => Status::ApplyError, + n => Status::Unknown(n), + }, + message, + }) + } + #[cfg(not(otadump_zucchini))] + { + let _ = (old, patch, output); + Err(Error { + status: Status::UnsupportedTarget, + message: "Zucchini is supported only on Linux x86_64 GNU hosts".into(), + }) + } +} + +/// Same as [`apply`], but cooperatively stops when `cancelled` returns true. +/// Native Zucchini is one non-interruptible operation; cancellation is checked +/// immediately before and after that call. +pub fn apply_with_cancel bool>( + old: &[u8], + patch: &[u8], + output_size: usize, + cancelled: F, +) -> Result> { + if cancelled() { + return Err(Error { + status: Status::Cancelled, + message: "Zucchini operation cancelled".into(), + }); + } + // ponytail: the upstream apply loop has no safe cancellation seam; the + // operation boundary is the deliberate cancellation ceiling. + let result = apply(old, patch, output_size); + if cancelled() { + return Err(Error { + status: Status::Cancelled, + message: "Zucchini operation cancelled".into(), + }); + } + result +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn apply_with_cancel_reports_cancelled() { + let error = apply_with_cancel(b"abcd", b"not a patch", 5, || true).unwrap_err(); + assert_eq!(error.status(), Status::Cancelled); + } + + #[test] + fn apply_reports_invalid_patch_when_not_cancelled() { + let error = apply_with_cancel(b"abcd", b"not a patch", 5, || false).unwrap_err(); + assert_eq!(error.status(), Status::InvalidPatch); + } +} diff --git a/tests/delta_extraction.rs b/tests/delta_extraction.rs new file mode 100644 index 0000000..83dcbd6 --- /dev/null +++ b/tests/delta_extraction.rs @@ -0,0 +1,244 @@ +use std::fs; +use std::path::{Path, PathBuf}; + +use otadump::ExtractOptions; +use prost::Message; +use ring::digest; +use tempfile::TempDir; + +#[derive(Clone, PartialEq, Message)] +struct Extent { + #[prost(uint64, optional, tag = "1")] + start_block: Option, + #[prost(uint64, optional, tag = "2")] + num_blocks: Option, +} + +#[derive(Clone, PartialEq, Message)] +struct PartitionInfo { + #[prost(uint64, optional, tag = "1")] + size: Option, + #[prost(bytes = "vec", optional, tag = "2")] + hash: Option>, +} + +#[derive(Clone, PartialEq, Message)] +struct InstallOperation { + #[prost(int32, tag = "1")] + operation_type: i32, + #[prost(uint64, optional, tag = "2")] + data_offset: Option, + #[prost(uint64, optional, tag = "3")] + data_length: Option, + #[prost(message, repeated, tag = "4")] + src_extents: Vec, + #[prost(uint64, optional, tag = "5")] + src_length: Option, + #[prost(message, repeated, tag = "6")] + dst_extents: Vec, + #[prost(uint64, optional, tag = "7")] + dst_length: Option, + #[prost(bytes = "vec", optional, tag = "8")] + data_sha256_hash: Option>, + #[prost(bytes = "vec", optional, tag = "9")] + src_sha256_hash: Option>, +} + +#[derive(Clone, PartialEq, Message)] +struct PartitionUpdate { + #[prost(string, tag = "1")] + partition_name: String, + #[prost(message, optional, tag = "6")] + old_partition_info: Option, + #[prost(message, optional, tag = "7")] + new_partition_info: Option, + #[prost(message, repeated, tag = "8")] + operations: Vec, +} + +#[derive(Clone, PartialEq, Message)] +struct DeltaArchiveManifest { + #[prost(uint32, optional, tag = "3")] + block_size: Option, + #[prost(uint32, optional, tag = "12")] + minor_version: Option, + #[prost(message, repeated, tag = "13")] + partitions: Vec, +} + +fn sha256(bytes: &[u8]) -> Vec { + digest::digest(&digest::SHA256, bytes).as_ref().to_vec() +} + +fn extent(start_block: u64, num_blocks: u64) -> Extent { + Extent { start_block: Some(start_block), num_blocks: Some(num_blocks) } +} + +fn partition_info(bytes: &[u8]) -> PartitionInfo { + PartitionInfo { size: Some(bytes.len() as u64), hash: Some(sha256(bytes)) } +} + +fn write_payload(path: &Path, manifest: DeltaArchiveManifest, data: &[u8]) { + let manifest = manifest.encode_to_vec(); + let mut payload = Vec::with_capacity(24 + manifest.len() + data.len()); + payload.extend_from_slice(b"CrAU"); + payload.extend_from_slice(&2u64.to_be_bytes()); + payload.extend_from_slice(&(manifest.len() as u64).to_be_bytes()); + payload.extend_from_slice(&0u32.to_be_bytes()); + payload.extend_from_slice(&manifest); + payload.extend_from_slice(data); + fs::write(path, payload).unwrap(); +} + +fn puffin_fixture(name: &str) -> Vec { + let path = PathBuf::from("tests/fixtures/puffin").join(name); + fs::read(path).unwrap() +} + +#[test] +fn full_replace_payload_extracts_expected_partition() { + let temporary = TempDir::new().unwrap(); + let output_dir = temporary.path().join("output"); + let expected = b"full payload data"; + let operation = InstallOperation { + operation_type: 0, + data_offset: Some(0), + data_length: Some(expected.len() as u64), + dst_extents: vec![extent(0, 2)], + data_sha256_hash: Some(sha256(expected)), + ..Default::default() + }; + let mut padded = expected.to_vec(); + padded.resize(32, 0); + let manifest = DeltaArchiveManifest { + block_size: Some(16), + minor_version: Some(0), + partitions: vec![PartitionUpdate { + partition_name: "vendor".into(), + old_partition_info: None, + new_partition_info: Some(partition_info(&padded)), + operations: vec![operation], + }], + }; + let payload = temporary.path().join("payload.bin"); + write_payload(&payload, manifest, expected); + + ExtractOptions::new().extract(&payload, &output_dir).unwrap(); + + assert_eq!(fs::read(output_dir.join("vendor.img")).unwrap(), padded); +} + +#[test] +fn source_copy_keeps_source_immutable_and_respects_extent_order() { + let temporary = TempDir::new().unwrap(); + let source_dir = temporary.path().join("source"); + let output_dir = temporary.path().join("output"); + fs::create_dir(&source_dir).unwrap(); + let source = b"AAAABBBBCCCCDDDD"; + fs::write(source_dir.join("system.img"), source).unwrap(); + let expected = b"AAAACCCC"; + let operation = InstallOperation { + operation_type: 4, + src_extents: vec![extent(2, 1), extent(0, 1)], + src_length: Some(8), + dst_extents: vec![extent(1, 1), extent(0, 1)], + dst_length: Some(8), + src_sha256_hash: Some(sha256(b"CCCCAAAA")), + ..Default::default() + }; + let manifest = DeltaArchiveManifest { + block_size: Some(4), + minor_version: Some(2), + partitions: vec![PartitionUpdate { + partition_name: "system".into(), + old_partition_info: Some(partition_info(source)), + new_partition_info: Some(partition_info(expected)), + operations: vec![operation], + }], + }; + let payload = temporary.path().join("payload.bin"); + write_payload(&payload, manifest, &[]); + + ExtractOptions::new().source_dir(&source_dir).extract(&payload, &output_dir).unwrap(); + + assert_eq!(fs::read(output_dir.join("system.img")).unwrap(), expected); + assert_eq!(fs::read(source_dir.join("system.img")).unwrap(), source); +} + +#[test] +fn missing_source_image_fails_without_publishing_output() { + let temporary = TempDir::new().unwrap(); + let source_dir = temporary.path().join("source"); + let output_dir = temporary.path().join("output"); + fs::create_dir(&source_dir).unwrap(); + let operation = InstallOperation { + operation_type: 4, + src_extents: vec![extent(0, 1)], + dst_extents: vec![extent(0, 1)], + ..Default::default() + }; + let manifest = DeltaArchiveManifest { + block_size: Some(4), + minor_version: Some(2), + partitions: vec![PartitionUpdate { + partition_name: "boot".into(), + old_partition_info: Some(partition_info(b"base")), + new_partition_info: Some(partition_info(b"base")), + operations: vec![operation], + }], + }; + let payload = temporary.path().join("payload.bin"); + write_payload(&payload, manifest, &[]); + + let error = ExtractOptions::new() + .source_dir(&source_dir) + .extract(&payload, &output_dir) + .unwrap_err() + .to_string(); + + assert!(error.contains("Unable to open source partition image"), "unexpected error: {error}"); + assert!(!output_dir.join("boot.img").exists()); + assert!(fs::read_dir(&output_dir).unwrap().next().is_none()); +} + +#[cfg(otadump_zucchini)] +#[test] +fn puffdiff_zucchini_golden_reconstructs_and_keeps_source() { + let source = puffin_fixture("deflates-sample1.bin"); + let target = puffin_fixture("deflates-sample2.bin"); + let patch = puffin_fixture("patch-1-to-2-zucchini.puf"); + + let temporary = TempDir::new().unwrap(); + let source_dir = temporary.path().join("source"); + let output_dir = temporary.path().join("output"); + fs::create_dir(&source_dir).unwrap(); + fs::write(source_dir.join("system.img"), &source).unwrap(); + + let manifest = DeltaArchiveManifest { + block_size: Some(1), + minor_version: Some(5), + partitions: vec![PartitionUpdate { + partition_name: "system".into(), + old_partition_info: Some(partition_info(&source)), + new_partition_info: Some(partition_info(&target)), + operations: vec![InstallOperation { + operation_type: 9, + data_offset: Some(0), + data_length: Some(patch.len() as u64), + src_extents: vec![extent(0, source.len() as u64)], + src_length: Some(source.len() as u64), + dst_extents: vec![extent(0, target.len() as u64)], + dst_length: Some(target.len() as u64), + data_sha256_hash: Some(sha256(&patch)), + src_sha256_hash: Some(sha256(&source)), + }], + }], + }; + let payload = temporary.path().join("payload.bin"); + write_payload(&payload, manifest, &patch); + + ExtractOptions::new().source_dir(&source_dir).extract(&payload, &output_dir).unwrap(); + + assert_eq!(fs::read(output_dir.join("system.img")).unwrap(), target); + assert_eq!(fs::read(source_dir.join("system.img")).unwrap(), source); +} diff --git a/tests/fixtures/puffin/LICENSE b/tests/fixtures/puffin/LICENSE new file mode 100644 index 0000000..f6d41c6 --- /dev/null +++ b/tests/fixtures/puffin/LICENSE @@ -0,0 +1,27 @@ +// Copyright 2017 The ChromiumOS Authors +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are +// met: +// +// * Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// * Redistributions in binary form must reproduce the above +// copyright notice, this list of conditions and the following disclaimer +// in the documentation and/or other materials provided with the +// distribution. +// * Neither the name of Google Inc. nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/tests/fixtures/puffin/LICENSE.zucchini b/tests/fixtures/puffin/LICENSE.zucchini new file mode 100644 index 0000000..a32e00c --- /dev/null +++ b/tests/fixtures/puffin/LICENSE.zucchini @@ -0,0 +1,27 @@ +// Copyright 2015 The Chromium Authors. All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are +// met: +// +// * Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// * Redistributions in binary form must reproduce the above +// copyright notice, this list of conditions and the following disclaimer +// in the documentation and/or other materials provided with the +// distribution. +// * Neither the name of Google Inc. nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/tests/fixtures/puffin/README.md b/tests/fixtures/puffin/README.md new file mode 100644 index 0000000..1bb1be5 --- /dev/null +++ b/tests/fixtures/puffin/README.md @@ -0,0 +1,38 @@ +# Puffin fixtures + +These fixtures come from ChromiumOS Puffin commit `343e23db1b4d81045e91a10244244893f5acd73b`. +The source files are `src/unittest_common.cc` and `src/patching_unittest.cc`. + +- `deflates-sample1.bin` is `kDeflatesSample1`. +- `deflates-sample2.bin` is `kDeflatesSample2`. +- `patch-1-to-2-zucchini.puf` reconstructs `deflates-sample2.bin` from `deflates-sample1.bin` through three actual deflate streams on each side. + +Only these three fixture payloads are retained because they are the minimal set +needed by the product-contract PUFFDIFF/ZUCCHINI extraction test. + +The files retain Puffin's BSD-3-Clause license in `LICENSE` and Zucchini's BSD-3-Clause license in `LICENSE.zucchini`. +Verify them with `sha256sum -c SHA256SUMS` from this directory. + +The Zucchini fixture uses Puffin at the commit above and AOSP Zucchini at `e256025d9d9c906caba814fd7b99a311ac8b8ec1`. +Build the pinned AOSP host `puffin` target with the pinned Zucchini source. +Generate and independently apply the PUF1 bytes with the upstream tool: + +```sh +PUFFIN=/path/to/aosp/out/host/linux-x86/bin/puffin +"$PUFFIN" \ + --operation=puffdiff \ + --src_file=deflates-sample1.bin \ + --dst_file=deflates-sample2.bin \ + --patch_file=patch-1-to-2-zucchini.puf \ + --src_file_type=raw \ + --dst_file_type=raw \ + --src_deflates_bit=16:50,80:10,96:18 \ + --dst_deflates_bit=0:50,72:80,152:18 \ + --patch_algorithm=1 +"$PUFFIN" \ + --operation=puffpatch \ + --src_file=deflates-sample1.bin \ + --dst_file=generated.bin \ + --patch_file=patch-1-to-2-zucchini.puf +cmp deflates-sample2.bin generated.bin +``` diff --git a/tests/fixtures/puffin/SHA256SUMS b/tests/fixtures/puffin/SHA256SUMS new file mode 100644 index 0000000..ee4413e --- /dev/null +++ b/tests/fixtures/puffin/SHA256SUMS @@ -0,0 +1,3 @@ +de0527c0d2d27f378cb79aec447172e0efcfff9cc45d0be05c313a87a2069280 deflates-sample1.bin +649b17ae72a6d503a1176604d876923270cd467ae1d458a54e58575eda98766e deflates-sample2.bin +826cd5b1a815cda2d13c2ca5733a9f9e99e0030d27178f2e9d44e47d73d4a9a8 patch-1-to-2-zucchini.puf diff --git a/tests/fixtures/puffin/deflates-sample1.bin b/tests/fixtures/puffin/deflates-sample1.bin new file mode 100644 index 0000000..3bfad73 Binary files /dev/null and b/tests/fixtures/puffin/deflates-sample1.bin differ diff --git a/tests/fixtures/puffin/deflates-sample2.bin b/tests/fixtures/puffin/deflates-sample2.bin new file mode 100644 index 0000000..a5ae9f5 Binary files /dev/null and b/tests/fixtures/puffin/deflates-sample2.bin differ diff --git a/tests/fixtures/puffin/patch-1-to-2-zucchini.puf b/tests/fixtures/puffin/patch-1-to-2-zucchini.puf new file mode 100644 index 0000000..948ff9d Binary files /dev/null and b/tests/fixtures/puffin/patch-1-to-2-zucchini.puf differ diff --git a/tests/native_artifact_fetch_test.py b/tests/native_artifact_fetch_test.py new file mode 100644 index 0000000..dae9f14 --- /dev/null +++ b/tests/native_artifact_fetch_test.py @@ -0,0 +1,150 @@ +import hashlib +import json +import os +import pathlib +import subprocess +import tarfile +import tempfile +import unittest + + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[1] +FETCH_SCRIPT = REPO_ROOT / "scripts" / "fetch-native-artifact.py" + + +def sha256_file(path: pathlib.Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def write_tree(root: pathlib.Path, entries: dict[str, bytes]) -> None: + for relpath, content in entries.items(): + target = root / relpath + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(content) + + +def write_bundle( + root: pathlib.Path, + entries: dict[str, bytes], + *, + include_traversal_entry: bool = False, +) -> pathlib.Path: + extract_root = root / "linux-x86_64-gnu" + write_tree(extract_root, entries) + archive = root / "bundle.tar.gz" + with tarfile.open(archive, "w:gz") as tar: + tar.add(extract_root, arcname="linux-x86_64-gnu") + if include_traversal_entry: + attack = root / "attack.txt" + attack.write_text("attack", encoding="utf-8") + tar.add(attack, arcname="../attack.txt") + return archive + + +def write_locks(root: pathlib.Path, archive: pathlib.Path, url: str) -> tuple[pathlib.Path, pathlib.Path]: + checksums = { + "lib/libotadump_zucchini.a": b"archive-bytes\n", + "SHA256SUMS": b"sha-entry\n", + "include/zucchini_ffi.h": b"#pragma once\n", + "licenses/LICENSE.zucchini": b"license\n", + "provenance.txt": b"provenance\n", + } + lock_path = root / "ARTIFACT_LOCK.sha256" + lock_path.write_text( + "\n".join(f"{hashlib.sha256(content).hexdigest()} {relpath}" for relpath, content in checksums.items()) + + "\n", + encoding="utf-8", + ) + bundle = { + "schema": 1, + "commit": "9534f6e24296b1682d53cdac19da0b81fc97fcbc", + "url": url, + "sha256": sha256_file(archive), + "extract_root": "linux-x86_64-gnu", + } + bundle_path = root / "ARTIFACT_BUNDLE_LOCK.json" + bundle_path.write_text(json.dumps(bundle), encoding="utf-8") + return bundle_path, lock_path + + +class NativeArtifactFetchTests(unittest.TestCase): + def setUp(self) -> None: + self.tempdir = tempfile.TemporaryDirectory() + self.root = pathlib.Path(self.tempdir.name) + self.entries = { + "lib/libotadump_zucchini.a": b"archive-bytes\n", + "SHA256SUMS": b"sha-entry\n", + "include/zucchini_ffi.h": b"#pragma once\n", + "licenses/LICENSE.zucchini": b"license\n", + "provenance.txt": b"provenance\n", + } + self.valid_url = ( + "https://raw.githubusercontent.com/AndroidDumps/otadump/" + "9534f6e24296b1682d53cdac19da0b81fc97fcbc/native/artifacts/zucchini/" + "zucchini-linux-x86_64-gnu.tar.gz" + ) + + def tearDown(self) -> None: + self.tempdir.cleanup() + + def run_fetch( + self, + bundle_lock: pathlib.Path, + checksum_lock: pathlib.Path, + out_dir: pathlib.Path, + cache_dir: pathlib.Path, + preseed: pathlib.Path, + ) -> subprocess.CompletedProcess[str]: + env = os.environ.copy() + env["OTADUMP_NATIVE_PRESEED"] = str(preseed) + env.pop("OTADUMP_NATIVE_OFFLINE", None) + env.pop("OTADUMP_NATIVE_CACHE", None) + return subprocess.run( + [ + "python3", + str(FETCH_SCRIPT), + "--bundle-lock", + str(bundle_lock), + "--checksum-lock", + str(checksum_lock), + "--out", + str(out_dir), + "--cache", + str(cache_dir), + ], + text=True, + capture_output=True, + env=env, + check=False, + ) + + def test_preseed_archive_fetch_populates_verified_output(self) -> None: + archive = write_bundle(self.root, self.entries) + bundle_lock, checksum_lock = write_locks(self.root, archive, self.valid_url) + out_dir = self.root / "out" / "linux-x86_64-gnu" + result = self.run_fetch(bundle_lock, checksum_lock, out_dir, self.root / "cache", archive) + + self.assertEqual(result.returncode, 0, result.stderr) + self.assertTrue((out_dir / "lib" / "libotadump_zucchini.a").is_file()) + + def test_traversal_entry_is_rejected(self) -> None: + archive = write_bundle(self.root, self.entries, include_traversal_entry=True) + bundle_lock, checksum_lock = write_locks(self.root, archive, self.valid_url) + result = self.run_fetch( + bundle_lock, + checksum_lock, + self.root / "out" / "linux-x86_64-gnu", + self.root / "cache", + archive, + ) + + self.assertNotEqual(result.returncode, 0) + self.assertIn("unsafe entry", result.stderr) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/python_smoke.py b/tests/python_smoke.py index 8c9d270..6be42a5 100644 --- a/tests/python_smoke.py +++ b/tests/python_smoke.py @@ -1,21 +1,90 @@ -from inspect import signature +from hashlib import sha256 from pathlib import Path from tempfile import TemporaryDirectory +from typing import Optional import otadump -def main() -> None: - assert issubclass(otadump.OtaDumpError, Exception) - assert "partitions" in str(signature(otadump.extract)) +def _varint(value: int) -> bytes: + encoded = bytearray() + while value >= 0x80: + encoded.append((value & 0x7F) | 0x80) + value >>= 7 + encoded.append(value) + return bytes(encoded) + + +def _integer(field: int, value: int) -> bytes: + return _varint(field << 3) + _varint(value) + + +def _bytes(field: int, value: bytes) -> bytes: + return _varint((field << 3) | 2) + _varint(len(value)) + value + + +def _extent(length: int) -> bytes: + return _integer(1, 0) + _integer(2, length) + +def _partition_info(data: bytes) -> bytes: + return _integer(1, len(data)) + _bytes(2, sha256(data).digest()) + + +def _payload(target: bytes, source: Optional[bytes] = None) -> bytes: + operation_type = 0 if source is None else 4 + operation = _integer(1, operation_type) + if source is None: + operation += _integer(2, 0) + _integer(3, len(target)) + operation += _bytes(8, sha256(target).digest()) + else: + operation += _bytes(4, _extent(len(source))) + _integer(5, len(source)) + operation += _bytes(9, sha256(source).digest()) + operation += _bytes(6, _extent(len(target))) + _integer(7, len(target)) + + partition = _bytes(1, b"system") + if source is not None: + partition += _bytes(6, _partition_info(source)) + partition += _bytes(7, _partition_info(target)) + _bytes(8, operation) + manifest = _integer(3, 1) + _bytes(13, partition) + operation_data = target if source is None else b"" + return ( + b"CrAU" + + (2).to_bytes(8, "big") + + len(manifest).to_bytes(8, "big") + + bytes(4) + + manifest + + operation_data + ) + + +def main() -> None: with TemporaryDirectory() as temporary_directory: temporary_path = Path(temporary_directory) - payload_file = temporary_path / "invalid.bin" - payload_file.write_bytes(b"invalid") + full_target = b"full partition image" + full_payload = temporary_path / "full.bin" + full_payload.write_bytes(_payload(full_target)) + full_output = temporary_path / "full-output" + otadump.extract(full_payload, full_output) + assert (full_output / "system.img").read_bytes() == full_target + + source = b"source partition image" + source_dir = temporary_path / "source" + source_dir.mkdir() + source_image = source_dir / "system.img" + source_image.write_bytes(source) + delta_payload = temporary_path / "delta.bin" + delta_payload.write_bytes(_payload(source, source)) + delta_output = temporary_path / "delta-output" + otadump.extract(delta_payload, delta_output, source_dir=source_dir) + assert (delta_output / "system.img").read_bytes() == source + assert source_image.read_bytes() == source + + invalid_payload = temporary_path / "invalid.bin" + invalid_payload.write_bytes(b"invalid") try: - otadump.extract(payload_file, temporary_path / "output") + otadump.extract(invalid_payload, temporary_path / "invalid-output") except otadump.OtaDumpError as error: assert str(error) == "Invalid payload file" else: