apk/archive/

writer.rs

//! Deterministic, aligned archive writing on top of the `zip` crate.

use std::collections::HashSet;
use std::fs::File;
use std::io::{self, Read, Seek, Write};
use std::path::Path;

use zip::write::SimpleFileOptions;
use zip::{CompressionMethod, DateTime, ZipArchive, ZipWriter};

use super::is_signed_entry;
use super::{Content, ZipError, ZipLimit};
use crate::digest::{DigestAlgorithm, EntryHasher};
use crate::error::{Error, FileOperation};
use crate::types::ArchivePath;

/// How entry data is compressed.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Compression {
    Stored,
    Deflated,
}

/// How one entry is written.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct EntryOptions {
    pub compression: Compression,
    /// The alignment of stored data within the file. Compressed data is
    /// never aligned (it cannot be memory-mapped).
    pub alignment: u16,
}

impl EntryOptions {
    /// Uncompressed with 4-byte alignment, zipalign's default for stored entries.
    pub const STORED: Self = Self { compression: Compression::Stored, alignment: 4 };
    /// Deflate-compressed.
    pub const DEFLATED: Self = Self { compression: Compression::Deflated, alignment: 1 };

    /// Uncompressed and aligned to `page_size`, so Android can map the data
    /// directly (native libraries with `extractNativeLibs="false"`).
    pub const fn page_aligned(page_size: u16) -> Self {
        Self { compression: Compression::Stored, alignment: page_size }
    }

    /// Stored for formats that are already compressed (as aapt2 does), and
    /// deflated otherwise.
    pub fn for_path(path: &ArchivePath) -> Self {
        const PRECOMPRESSED: &[&str] = &[
            "jpg", "jpeg", "png", "gif", "webp", "wav", "mp2", "mp3", "ogg", "aac", "mpg", "mpeg", "mid", "midi",
            "smf", "jet", "rtttl", "imy", "xmf", "mp4", "m4a", "m4v", "3gp", "3gpp", "3g2", "3gpp2", "amr", "awb",
            "wma", "wmv", "webm", "mkv", "zip", "jar", "apk", "gz", "zst", "br", "xz", "bz2", "7z", "ktx2", "basis",
        ];
        match path.extension() {
            Some(ext) if PRECOMPRESSED.iter().any(|p| p.eq_ignore_ascii_case(ext)) => Self::STORED,
            _ => Self::DEFLATED,
        }
    }

    fn zip_options(self) -> SimpleFileOptions {
        let (method, alignment) = match self.compression {
            Compression::Stored => (CompressionMethod::Stored, self.alignment.max(1)),
            Compression::Deflated => (CompressionMethod::Deflated, 1),
        };
        SimpleFileOptions::default()
            .compression_method(method)
            .last_modified_time(DateTime::default())
            .unix_permissions(0o644)
            .large_file(false)
            .with_alignment(alignment)
    }
}

/// Writes APK entries with reproducible metadata and per-entry alignment.
///
/// Entry names are unique, timestamps are fixed at 1980-01-01, and
/// permissions are `0644`, so identical inputs produce identical archives.
/// When digesting is enabled, the digest of every entry JAR signing covers
/// is recorded as it is written.
pub(crate) struct ApkWriter<W: Write + Seek> {
    zip: ZipWriter<W>,
    names: HashSet<String>,
    digest: Option<DigestAlgorithm>,
    digests: Vec<(String, Vec<u8>)>,
}

impl<W: Write + Seek> ApkWriter<W> {
    pub fn new(sink: W) -> Self {
        Self { zip: ZipWriter::new(sink), names: HashSet::new(), digest: None, digests: Vec::new() }
    }

    /// Records a digest of each entry's uncompressed data, for JAR signing
    /// without re-reading the archive.
    pub fn digest_entries(mut self, algorithm: DigestAlgorithm) -> Self {
        self.digest = Some(algorithm);
        self
    }

    /// The name and digest of every digested entry written so far.
    pub fn digests(&self) -> &[(String, Vec<u8>)] {
        &self.digests
    }

    /// Adds an entry with the contents of `data`.
    pub fn add(&mut self, path: &ArchivePath, data: &[u8], options: EntryOptions) -> Result<(), ZipError> {
        self.add_reader(path, data, options).map(drop)
    }

    /// Adds an entry streamed from `reader` and returns its size.
    pub fn add_reader(
        &mut self,
        path: &ArchivePath,
        reader: impl Read,
        options: EntryOptions,
    ) -> Result<u64, ZipError> {
        self.start(path.as_str())?;
        self.zip.start_file(path.as_str(), options.zip_options())?;
        self.stream(path.as_str(), reader)
    }

    /// Adds an entry from in-memory or file [`Content`]. A file must not
    /// change size while it is copied.
    pub fn add_content(&mut self, path: &ArchivePath, content: &Content, options: EntryOptions) -> Result<u64, Error> {
        match content {
            Content::Bytes(bytes) => Ok(self.add_reader(path, &bytes[..], options)?),
            Content::File(source) => self.add_file(path, source, options),
        }
    }

    fn add_file(&mut self, path: &ArchivePath, source: &Path, options: EntryOptions) -> Result<u64, Error> {
        let file = File::open(source).map_err(Error::file(FileOperation::Open, source))?;
        let expected = file.metadata().map_err(Error::file(FileOperation::Inspect, source))?.len();
        if expected >= u64::from(u32::MAX) {
            return Err(ZipError::LimitExceeded { limit: ZipLimit::Size }.into());
        }
        // Reading one byte past the expected length detects growth.
        let actual = self.add_reader(path, file.take(expected + 1), options)?;
        if actual != expected {
            return Err(ZipError::SourceChanged { name: path.to_string(), expected, actual }.into());
        }
        Ok(actual)
    }

    /// Copies entry `index` of `archive`. Stored data is realigned to
    /// `alignment`; compressed data is copied without recompressing it.
    /// Metadata is normalized as for new entries.
    pub fn copy<R: Read + Seek>(
        &mut self,
        archive: &mut ZipArchive<R>,
        index: usize,
        alignment: u16,
    ) -> Result<(), ZipError> {
        let (name, encrypted, stored) = {
            let entry = archive.by_index_data(index)?;
            let name = String::from_utf8(entry.name_raw().to_vec()).map_err(|_| ZipError::NonUtf8Name { index })?;
            (name, entry.encrypted(), entry.compression() == CompressionMethod::Stored)
        };
        if encrypted {
            return Err(ZipError::EncryptedEntry { name });
        }
        // Names end up in JAR manifests, where control characters could forge
        // attributes, so they must be safe archive paths. As for Android,
        // only a trailing `/` marks a directory.
        let directory = name.ends_with('/');
        ArchivePath::try_new(name.strip_suffix('/').unwrap_or(&name))
            .map_err(|source| ZipError::UnsafeName { name: name.clone(), source })?;
        self.start(&name)?;
        if stored && !directory {
            let options = EntryOptions { compression: Compression::Stored, alignment }.zip_options();
            self.zip.start_file(name.as_str(), options)?;
            // Reading checks the data against its CRC.
            let data = archive.by_index(index)?;
            self.stream(&name, data)?;
            return Ok(());
        }
        if let Some(algorithm) = self.digest.filter(|_| is_signed_entry(&name)) {
            let mut hasher = algorithm.hasher();
            io::copy(&mut archive.by_index(index)?, &mut hasher)?;
            self.digests.push((name, hasher.finish()));
        }
        let raw = archive.by_index_raw(index)?;
        self.zip.raw_copy_file_touch(raw, DateTime::default(), Some(0o644))?;
        Ok(())
    }

    /// Sets the archive comment.
    pub fn set_comment(&mut self, comment: &[u8]) -> Result<(), ZipError> {
        Ok(self.zip.set_raw_comment(comment.into())?)
    }

    fn start(&mut self, name: &str) -> Result<(), ZipError> {
        if self.names.len() >= usize::from(u16::MAX) - 1 {
            return Err(ZipError::LimitExceeded { limit: ZipLimit::EntryCount });
        }
        if !self.names.insert(name.to_owned()) {
            return Err(ZipError::DuplicateEntry { name: name.to_owned() });
        }
        Ok(())
    }

    /// Writes the data of the entry just started, digesting it if enabled.
    fn stream(&mut self, name: &str, mut reader: impl Read) -> Result<u64, ZipError> {
        let mut hasher = self.digest.filter(|_| is_signed_entry(name)).map(DigestAlgorithm::hasher);
        let copied = io::copy(&mut reader, &mut Tee { zip: &mut self.zip, hasher: hasher.as_mut() });
        let size = match copied {
            Ok(size) => size,
            Err(error) => {
                let _ = self.zip.abort_file();
                return Err(error.into());
            }
        };
        if let Some(hasher) = hasher {
            self.digests.push((name.to_owned(), hasher.finish()));
        }
        Ok(size)
    }

    /// Writes the central directory and returns the sink.
    pub fn finish(self) -> Result<W, ZipError> {
        Ok(self.zip.finish()?)
    }
}

/// Feeds written bytes to the archive and, optionally, a digest.
struct Tee<'a, W: Write + Seek> {
    zip: &'a mut ZipWriter<W>,
    hasher: Option<&'a mut EntryHasher>,
}

impl<W: Write + Seek> Write for Tee<'_, W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let written = self.zip.write(buf)?;
        if let Some(hasher) = self.hasher.as_deref_mut() {
            hasher.update(&buf[..written]);
        }
        Ok(written)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.zip.flush()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::archive::{ApkLayout, open};
    use std::io::Cursor;

    fn path(p: &str) -> ArchivePath {
        ArchivePath::try_new(p).unwrap()
    }

    fn data_offset(bytes: &[u8], name: &str) -> u64 {
        let mut archive = ZipArchive::new(Cursor::new(bytes)).unwrap();
        archive.by_name(name).unwrap().data_start().unwrap()
    }

    #[test]
    fn aligns_stored_entries_and_is_reproducible() {
        let build = || {
            let mut writer = ApkWriter::new(Cursor::new(Vec::new()));
            writer.add(&path("AndroidManifest.xml"), b"<manifest/>", EntryOptions::DEFLATED).unwrap();
            writer.add(&path("resources.arsc"), b"table", EntryOptions::STORED).unwrap();
            writer.add(&path("lib/arm64-v8a/libmain.so"), &[7; 1000], EntryOptions::page_aligned(16384)).unwrap();
            writer.finish().unwrap().into_inner()
        };
        let bytes = build();
        assert_eq!(bytes, build(), "identical inputs must produce identical archives");
        assert_eq!(data_offset(&bytes, "resources.arsc") % 4, 0);
        assert_eq!(data_offset(&bytes, "lib/arm64-v8a/libmain.so") % 16384, 0);
        assert_eq!(ApkLayout::read(&mut Cursor::new(&bytes)).unwrap().entry_count(), 3);
    }

    #[test]
    fn digests_signed_entries_and_rejects_duplicates() {
        let mut writer = ApkWriter::new(Cursor::new(Vec::new())).digest_entries(DigestAlgorithm::Sha256);
        writer.add(&path("a.txt"), b"abc", EntryOptions::DEFLATED).unwrap();
        writer.add(&path("META-INF/CERT.SF"), b"not digested", EntryOptions::STORED).unwrap();
        assert!(matches!(
            writer.add(&path("a.txt"), b"again", EntryOptions::STORED),
            Err(ZipError::DuplicateEntry { name }) if name == "a.txt"
        ));
        assert_eq!(writer.digests(), [("a.txt".to_owned(), DigestAlgorithm::Sha256.digest(b"abc"))]);
    }

    #[test]
    fn copies_realign_stored_data_and_keep_compressed_data() {
        let mut source = ZipWriter::new(Cursor::new(Vec::new()));
        source
            .start_file("lib/x86/libm.so", SimpleFileOptions::default().compression_method(CompressionMethod::Stored))
            .unwrap();
        source.write_all(&[1; 100]).unwrap();
        // Preserve the original encoder settings rather than recompressing at the default level.
        let options =
            SimpleFileOptions::default().compression_method(CompressionMethod::Deflated).compression_level(Some(1));
        source.start_file("notes.txt", options).unwrap();
        source.write_all(&[b'n'; 500]).unwrap();
        let source = source.finish().unwrap().into_inner();
        let (mut archive, _) = open(Cursor::new(&source)).unwrap();
        let mut compressed = Vec::new();
        archive.by_index_raw(1).unwrap().read_to_end(&mut compressed).unwrap();

        let mut copy = ApkWriter::new(Cursor::new(Vec::new())).digest_entries(DigestAlgorithm::Sha256);
        copy.copy(&mut archive, 0, 4096).unwrap();
        copy.copy(&mut archive, 1, 4).unwrap();
        let digests = copy.digests().to_vec();
        let copied = copy.finish().unwrap().into_inner();

        assert_eq!(data_offset(&copied, "lib/x86/libm.so") % 4096, 0);
        let (mut copied_archive, _) = open(Cursor::new(&copied)).unwrap();
        let mut copied_compressed = Vec::new();
        copied_archive.by_index_raw(1).unwrap().read_to_end(&mut copied_compressed).unwrap();
        assert_eq!(copied_compressed, compressed, "deflated bytes changed");
        let mut text = Vec::new();
        copied_archive.by_name("notes.txt").unwrap().read_to_end(&mut text).unwrap();
        assert_eq!(text, [b'n'; 500]);
        assert_eq!(digests[1], ("notes.txt".to_owned(), DigestAlgorithm::Sha256.digest(&[b'n'; 500])));
    }
}