rust-apk 1.0.0

//! Digests: per-entry digests for JAR signing, and the chunked content
//! digests of APK Signature Scheme v2.

use std::io;
use std::num::NonZeroUsize;
use std::ops::Range;
use std::thread;

use der::oid::ObjectIdentifier;
use sha1::Sha1;
use sha2::{Digest, Sha256, Sha384, Sha512};

use crate::archive::{ApkLayout, ReadAt, ZipError, ZipLimit};

/// A digest used by JAR (v1) signatures, weakest first.
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub enum DigestAlgorithm {
    Sha1,
    Sha256,
    Sha384,
    Sha512,
}

impl DigestAlgorithm {
    const ALL: [Self; 4] = [Self::Sha1, Self::Sha256, Self::Sha384, Self::Sha512];

    pub(crate) fn hasher(self) -> EntryHasher {
        match self {
            Self::Sha1 => EntryHasher::Sha1(Sha1::new()),
            Self::Sha256 => EntryHasher::Sha256(Sha256::new()),
            Self::Sha384 => EntryHasher::Sha384(Sha384::new()),
            Self::Sha512 => EntryHasher::Sha512(Sha512::new()),
        }
    }

    pub fn digest(self, data: &[u8]) -> Vec<u8> {
        let mut hasher = self.hasher();
        hasher.update(data);
        hasher.finish()
    }

    /// The name used in JAR manifest attributes, e.g. `SHA-256-Digest`.
    pub(crate) fn jar_name(self) -> &'static str {
        match self {
            Self::Sha1 => "SHA1",
            Self::Sha256 => "SHA-256",
            Self::Sha384 => "SHA-384",
            Self::Sha512 => "SHA-512",
        }
    }

    pub(crate) fn from_jar_name(name: &str) -> Option<Self> {
        match name.to_ascii_uppercase().as_str() {
            "SHA-1" => Some(Self::Sha1),
            name => Self::ALL.into_iter().find(|a| a.jar_name() == name),
        }
    }

    pub(crate) fn oid(self) -> ObjectIdentifier {
        ObjectIdentifier::new_unwrap(match self {
            Self::Sha1 => "1.3.14.3.2.26",
            Self::Sha256 => "2.16.840.1.101.3.4.2.1",
            Self::Sha384 => "2.16.840.1.101.3.4.2.2",
            Self::Sha512 => "2.16.840.1.101.3.4.2.3",
        })
    }

    pub(crate) fn from_oid(oid: &ObjectIdentifier) -> Option<Self> {
        Self::ALL.into_iter().find(|a| a.oid() == *oid)
    }
}

/// An incremental [`DigestAlgorithm`], which can be written to.
#[derive(Clone)]
pub(crate) enum EntryHasher {
    Sha1(Sha1),
    Sha256(Sha256),
    Sha384(Sha384),
    Sha512(Sha512),
}

impl EntryHasher {
    pub fn update(&mut self, data: &[u8]) {
        match self {
            Self::Sha1(h) => h.update(data),
            Self::Sha256(h) => h.update(data),
            Self::Sha384(h) => h.update(data),
            Self::Sha512(h) => h.update(data),
        }
    }

    pub fn finish(self) -> Vec<u8> {
        match self {
            Self::Sha1(h) => h.finalize().to_vec(),
            Self::Sha256(h) => h.finalize().to_vec(),
            Self::Sha384(h) => h.finalize().to_vec(),
            Self::Sha512(h) => h.finalize().to_vec(),
        }
    }
}

impl io::Write for EntryHasher {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.update(buf);
        Ok(buf.len())
    }

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

/// A chunked content digest of APK Signature Scheme v2 and v3.
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub(crate) enum ContentDigest {
    ChunkedSha256,
    ChunkedSha512,
}

impl ContentDigest {
    const fn output_len(self) -> usize {
        match self {
            Self::ChunkedSha256 => 32,
            Self::ChunkedSha512 => 64,
        }
    }

    fn chunk(self, data: &[u8], out: &mut Vec<u8>) {
        fn hash<D: Digest>(data: &[u8], out: &mut Vec<u8>) {
            let mut hasher = D::new();
            hasher.update([0xa5]);
            hasher.update((data.len() as u32).to_le_bytes());
            hasher.update(data);
            out.extend_from_slice(&hasher.finalize());
        }
        match self {
            Self::ChunkedSha256 => hash::<Sha256>(data, out),
            Self::ChunkedSha512 => hash::<Sha512>(data, out),
        }
    }

    fn top(self, count: u32, chunks: &[u8]) -> Vec<u8> {
        fn hash<D: Digest>(count: u32, chunks: &[u8]) -> Vec<u8> {
            let mut hasher = D::new();
            hasher.update([0x5a]);
            hasher.update(count.to_le_bytes());
            hasher.update(chunks);
            hasher.finalize().to_vec()
        }
        match self {
            Self::ChunkedSha256 => hash::<Sha256>(count, chunks),
            Self::ChunkedSha512 => hash::<Sha512>(count, chunks),
        }
    }
}

const CHUNK_SIZE: u64 = 1 << 20;

/// One 1 MiB (or shorter, final) chunk of a digested section.
#[derive(Clone, Copy)]
enum Chunk {
    File { offset: u64, len: usize },
    Eocd,
}

/// Computes `algorithms` over the APK in `source`, as if its central directory
/// started at `layout.entries_end` (that is, with any signing block removed).
///
/// Chunks are hashed in parallel; each worker reads its chunks positionally.
pub(crate) fn content_digests(
    source: &(impl ReadAt + ?Sized),
    layout: &ApkLayout,
    algorithms: &[ContentDigest],
) -> Result<Vec<(ContentDigest, Vec<u8>)>, ZipError> {
    let eocd = layout.eocd_with_central_directory_at(layout.entries_end)?;
    let mut chunks = Vec::new();
    for section in [0..layout.entries_end, layout.central_directory.clone()] {
        split(section, &mut chunks);
    }
    chunks.push(Chunk::Eocd);
    let count = u32::try_from(chunks.len()).map_err(|_| ZipError::LimitExceeded { limit: ZipLimit::Size })?;

    let workers = thread::available_parallelism().map_or(1, NonZeroUsize::get).min(chunks.len().div_ceil(4)).max(1);
    let per_worker = chunks.len().div_ceil(workers);
    let digest_range = |chunks: &[Chunk]| -> io::Result<Vec<Vec<u8>>> {
        let mut buffer = vec![0u8; CHUNK_SIZE as usize];
        let mut outputs: Vec<Vec<u8>> =
            algorithms.iter().map(|a| Vec::with_capacity(chunks.len() * a.output_len())).collect();
        for chunk in chunks {
            let data = match *chunk {
                Chunk::File { offset, len } => {
                    source.read_exact_at(&mut buffer[..len], offset)?;
                    &buffer[..len]
                }
                Chunk::Eocd => &eocd[..],
            };
            for (algorithm, output) in algorithms.iter().zip(&mut outputs) {
                algorithm.chunk(data, output);
            }
        }
        Ok(outputs)
    };
    let parts: Vec<Vec<Vec<u8>>> = if workers == 1 {
        vec![digest_range(&chunks)?]
    } else {
        thread::scope(|scope| {
            let handles: Vec<_> = chunks.chunks(per_worker).map(|range| scope.spawn(|| digest_range(range))).collect();
            handles.into_iter().map(|h| h.join().expect("digest worker panicked")).collect::<io::Result<_>>()
        })?
    };
    Ok(algorithms
        .iter()
        .enumerate()
        .map(|(i, &algorithm)| {
            let concatenated: Vec<u8> = parts.iter().flat_map(|part| part[i].iter().copied()).collect();
            (algorithm, algorithm.top(count, &concatenated))
        })
        .collect())
}

fn split(section: Range<u64>, chunks: &mut Vec<Chunk>) {
    let mut offset = section.start;
    while offset < section.end {
        let len = (section.end - offset).min(CHUNK_SIZE);
        chunks.push(Chunk::File { offset, len: len as usize });
        offset += len;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::archive::{ApkWriter, EntryOptions};
    use crate::types::ArchivePath;

    /// The v2 digest computed directly from the specification.
    fn reference<D: Digest>(bytes: &[u8], layout: &ApkLayout) -> Vec<u8> {
        let mut eocd = bytes[layout.eocd_offset as usize..].to_vec();
        eocd[16..20].copy_from_slice(&(layout.entries_end as u32).to_le_bytes());
        let sections = [
            &bytes[..layout.entries_end as usize],
            &bytes[layout.central_directory.start as usize..layout.central_directory.end as usize],
            &eocd[..],
        ];
        let mut digests = Vec::new();
        let mut count = 0u32;
        for chunk in sections.iter().flat_map(|s| s.chunks(1024 * 1024)) {
            let mut hash = D::new();
            hash.update([0xa5]);
            hash.update((chunk.len() as u32).to_le_bytes());
            hash.update(chunk);
            digests.extend_from_slice(&hash.finalize());
            count += 1;
        }
        let mut hash = D::new();
        hash.update([0x5a]);
        hash.update(count.to_le_bytes());
        hash.update(digests);
        hash.finalize().to_vec()
    }

    #[test]
    fn parallel_digest_matches_the_specification() {
        let mut writer = ApkWriter::new(io::Cursor::new(Vec::new()));
        let noise: Vec<u8> = (0..5_500_000u32).map(|i| (i.wrapping_mul(2_654_435_761) >> 13) as u8).collect();
        writer.add(&ArchivePath::try_new("assets/noise.bin").unwrap(), &noise, EntryOptions::STORED).unwrap();
        let bytes = writer.finish().unwrap().into_inner();
        let layout = ApkLayout::read(&mut io::Cursor::new(&bytes)).unwrap();
        let digests =
            content_digests(&bytes, &layout, &[ContentDigest::ChunkedSha256, ContentDigest::ChunkedSha512]).unwrap();
        assert_eq!(digests[0], (ContentDigest::ChunkedSha256, reference::<Sha256>(&bytes, &layout)));
        assert_eq!(digests[1], (ContentDigest::ChunkedSha512, reference::<Sha512>(&bytes, &layout)));
    }

    #[test]
    fn chunk_boundaries_match_the_reference_for_both_digests() {
        let archive = |len| {
            let mut writer = ApkWriter::new(io::Cursor::new(Vec::new()));
            writer.add(&ArchivePath::manifest(), &vec![0x5b; len], EntryOptions::STORED).unwrap();
            writer.finish().unwrap().into_inner()
        };
        let prefix = ApkLayout::read(&mut io::Cursor::new(archive(0))).unwrap().entries_end as usize;
        for section_len in [prefix, prefix + 1, (1 << 20) - 1, 1 << 20, (1 << 20) + 1] {
            let bytes = archive(section_len - prefix);
            let layout = ApkLayout::read(&mut io::Cursor::new(&bytes)).unwrap();
            assert_eq!(layout.entries_end as usize, section_len);
            let digests =
                content_digests(&bytes, &layout, &[ContentDigest::ChunkedSha256, ContentDigest::ChunkedSha512])
                    .unwrap();
            assert_eq!(digests[0].1, reference::<Sha256>(&bytes, &layout), "SHA-256, {section_len}");
            assert_eq!(digests[1].1, reference::<Sha512>(&bytes, &layout), "SHA-512, {section_len}");
        }
    }
}