s2pa 0.1.0

// Copyright 2022 Adobe. All rights reserved.
// This file is licensed to you under the Apache License,
// Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
// or the MIT license (http://opensource.org/licenses/MIT),
// at your option.

// Unless required by applicable law or agreed to in writing,
// this software is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR REPRESENTATIONS OF ANY KIND, either express or
// implied. See the LICENSE-MIT and LICENSE-APACHE files for the
// specific language governing permissions and limitations under
// each license.

use std::{
    fs::File,
    io::{Cursor, Read, Seek, SeekFrom},
    ops::RangeInclusive,
    path::Path,
};

use range_set::RangeSet;
use serde::{Deserialize, Serialize};
// direct sha functions
use sha2::{Digest, Sha256, Sha384, Sha512};

use crate::{utils::io_utils::stream_len, Error, Result};

const MAX_HASH_BUF: usize = 256 * 1024 * 1024; // cap memory usage to 256MB

#[derive(Clone, Serialize, Deserialize, Debug, PartialEq, Eq)]
/// Defines a hash range to be used with `hash_stream_by_alg`
pub struct HashRange {
    start: u64,
    length: u64,

    #[serde(skip)]
    bmff_offset: Option<u64>, /* optional tracking of offset positions to include in BMFF_V2 hashes in BE format */
}

impl HashRange {
    pub fn new(start: u64, length: u64) -> Self {
        HashRange {
            start,
            length,
            bmff_offset: None,
        }
    }

    /// update the start value
    #[allow(dead_code)]
    pub fn set_start(&mut self, start: u64) {
        self.start = start;
    }

    /// return start as usize
    pub fn start(&self) -> u64 {
        self.start
    }

    /// return length as usize
    pub fn length(&self) -> u64 {
        self.length
    }

    pub fn set_length(&mut self, length: u64) {
        self.length = length;
    }

    // set offset for BMFF_V2 to be hashed in addition to data
    pub fn set_bmff_offset(&mut self, offset: u64) {
        self.bmff_offset = Some(offset);
    }

    // get option offset for BMFF_V2 hash
    pub fn bmff_offset(&self) -> Option<u64> {
        self.bmff_offset
    }
}

/// Compare two byte vectors return true if match, false otherwise
pub fn vec_compare(va: &[u8], vb: &[u8]) -> bool {
    (va.len() == vb.len()) &&  // zip stops at the shortest
     va.iter()
       .zip(vb)
       .all(|(a,b)| a == b)
}

#[derive(Clone)]
pub enum Hasher {
    SHA256(Sha256),
    SHA384(Sha384),
    SHA512(Sha512),
}

impl Hasher {
    // update hash value with new data
    pub fn update(&mut self, data: &[u8]) {
        use Hasher::*;
        // update the hash
        match self {
            SHA256(ref mut d) => d.update(data),
            SHA384(ref mut d) => d.update(data),
            SHA512(ref mut d) => d.update(data),
        }
    }

    // consume hasher and return the final digest
    pub fn finalize(hasher_enum: Hasher) -> Vec<u8> {
        use Hasher::*;
        // return the hash
        match hasher_enum {
            SHA256(d) => d.finalize().to_vec(),
            SHA384(d) => d.finalize().to_vec(),
            SHA512(d) => d.finalize().to_vec(),
        }
    }
}

// Return hash bytes for desired hashing algorithm.
pub fn hash_by_alg(alg: &str, data: &[u8], exclusions: Option<Vec<HashRange>>) -> Vec<u8> {
    let mut reader = Cursor::new(data);

    hash_stream_by_alg(alg, &mut reader, exclusions, true).unwrap_or_default()
}

// Return hash inclusive bytes for desired hashing algorithm.
pub fn hash_by_alg_with_inclusions(alg: &str, data: &[u8], inclusions: Vec<HashRange>) -> Vec<u8> {
    let mut reader = Cursor::new(data);

    hash_stream_by_alg(alg, &mut reader, Some(inclusions), false).unwrap_or_default()
}

// Return hash bytes for asset using desired hashing algorithm.
pub fn hash_asset_by_alg(
    alg: &str,
    asset_path: &Path,
    exclusions: Option<Vec<HashRange>>,
) -> Result<Vec<u8>> {
    let mut file = File::open(asset_path)?;
    hash_stream_by_alg(alg, &mut file, exclusions, true)
}

// Return hash inclusive bytes for asset using desired hashing algorithm.
pub fn hash_asset_by_alg_with_inclusions(
    alg: &str,
    asset_path: &Path,
    inclusions: Vec<HashRange>,
) -> Result<Vec<u8>> {
    let mut file = File::open(asset_path)?;
    hash_stream_by_alg(alg, &mut file, Some(inclusions), false)
}

/*  Returns hash bytes for a stream using desired hashing algorithm.  The function handles the many
    possible hash requirements of C2PA.  The function accepts a source stream 'data', an optional
    set of hash ranges 'hash_range' and a boolean to indicate whether the hash range is an exclusion
    or inclusion set of hash ranges.

    The basic case is to hash a stream without hash ranges:
    The data represents a single contiguous stream of bytes to be hash where D are data bytes

    to_be_hashed: [DDDDDDDDD...DDDDDDDDDD]

    The data is then chunked and hashed in groups to reduce memory
    footprint and increase performance.

    The most common case for C2PA is the use of an exclusion hash.  In this case the 'hash_range' indicate
    which byte ranges should be excluded shown here depicted with I for included bytes and  X for excluded bytes

    to_be_hashed: [IIIIXXXIIIIXXXXXIIIXXIII...IIII]

    In this case the data is split into a set of ranges covering the included bytes.  The set of ranged bytes
    are then chunked and hashed just like the default case.

    The opposite of this is when 'is_exclusion' is set to false indicating the 'hash_ranges' represent the bytes
    to include in the hash. Here are the bytes in 'data' are excluded except those explicitly referenced.

    to_be_hashed: [XXXXXXIIIIXXXXXIIXXXX...XXXX]

    Again a set of ranged bytes are created and hashed as described above.

    The last case is a special requirement for BMFF based assets (exclusion hashes only).  For this case we not
    only hash the data but also the location where the data was found in the asset.  To do this we add a special
    HashRange object to the hash ranges to indicate which locations in the stream require this special offset
    hash.  To make processing efficient we again split the data into ranges at not just the exclusion
    points but also for these markers.  The hashing loop knows to pause at these special marker ranges to insert
    the hash of the offset.  The stream sent to the hashing loop logically looks like this where M is the marker.
    to_be_hashed: [IIIIIXXXXXMIIIIIMXXXXXMXXXXIII...III]

    The data is again split into range sets breaking at the exclusion points and now also the markers.
*/
pub fn hash_stream_by_alg<R>(
    alg: &str,
    data: &mut R,
    hash_range: Option<Vec<HashRange>>,
    is_exclusion: bool,
) -> Result<Vec<u8>>
where
    R: Read + Seek + ?Sized,
{
    let mut bmff_v2_starts: Vec<u64> = Vec::new();

    use Hasher::*;
    let mut hasher_enum = match alg {
        "sha256" => SHA256(Sha256::new()),
        "sha384" => SHA384(Sha384::new()),
        "sha512" => SHA512(Sha512::new()),
        _ => {
            return Err(Error::UnsupportedType);
        }
    };

    let data_len = stream_len(data)?;
    data.rewind()?;

    if data_len < 1 {
        return Err(Error::OtherError("no data to hash".into()));
    }

    let ranges = match hash_range {
        Some(mut hr) if !hr.is_empty() => {
            // hash data skipping excluded regions
            // sort the exclusions
            hr.sort_by_key(|a| a.start());

            // verify structure of blocks
            let num_blocks = hr.len();
            let range_end = hr[num_blocks - 1].start() + hr[num_blocks - 1].length();
            let data_end = data_len - 1;

            // range extends past end of file so fail
            if data_len < range_end {
                return Err(Error::BadParam(
                    "The exclusion range exceed the data length".to_string(),
                ));
            }

            if is_exclusion {
                //build final ranges
                let mut ranges_vec: Vec<RangeInclusive<u64>> = Vec::new();
                let mut ranges = RangeSet::<[RangeInclusive<u64>; 1]>::from(0..=data_end);
                for exclusion in hr {
                    // add new BMFF V2 offset as a new range to be included so that we can
                    // pause to add the offset hash
                    if let Some(offset) = exclusion.bmff_offset() {
                        bmff_v2_starts.push(offset);
                        continue;
                    }

                    let end = exclusion.start() + exclusion.length() - 1;
                    let exclusion_start = exclusion.start();
                    ranges.remove_range(exclusion_start..=end);
                }

                // merge standard ranges and BMFF V2 ranges into single list
                if !bmff_v2_starts.is_empty() {
                    bmff_v2_starts.sort();

                    // split ranges at BMFF V2 offsets and insert offset value
                    for r in ranges.into_smallvec() {
                        // if bmff_v2 offset is within the range then split the range at the off set and both side to ranges_vec
                        let mut current_range = r;
                        for os in &bmff_v2_starts {
                            if current_range.contains(os) {
                                if *current_range.start() == *os {
                                    ranges_vec.push(RangeInclusive::new(*os, *os));
                                // offset
                                } else {
                                    ranges_vec
                                        .push(RangeInclusive::new(*current_range.start(), *os - 1)); // left side
                                    ranges_vec.push(RangeInclusive::new(*os, *os)); // offset
                                    current_range = RangeInclusive::new(*os, *current_range.end());
                                    // right side
                                }
                            }
                        }
                        ranges_vec.push(current_range);
                    }

                    // add in remaining BMFF V2 offsets that were not included in the ranges because of subsets
                    let range_start = RangeInclusive::new(0, 0);
                    let range_end = RangeInclusive::new(data_end, data_end);
                    let before_any_range = *ranges_vec.first().unwrap_or(&range_start).start();
                    let after_any_range = *ranges_vec.last().unwrap_or(&range_end).end();

                    for os in &bmff_v2_starts {
                        if !ranges_vec.iter().any(|r| r.contains(os))
                            && *os > before_any_range
                            && *os < after_any_range
                        {
                            ranges_vec.push(RangeInclusive::new(*os, *os));
                        }
                    }

                    // sort by start position
                    ranges_vec.sort_by(|a, b| {
                        let a_start = a.start();
                        let b_start = b.start();
                        a_start.cmp(b_start)
                    });

                    ranges_vec
                } else {
                    for r in ranges.into_smallvec() {
                        ranges_vec.push(r);
                    }
                    ranges_vec
                }
            } else {
                //build final ranges
                let mut ranges_vec: Vec<RangeInclusive<u64>> = Vec::new();
                for inclusion in hr {
                    let end = inclusion.start() + inclusion.length() - 1;
                    let inclusion_start = inclusion.start();

                    // add new BMFF V2 offset as a new range to be included so that we can
                    // pause to add the offset hash
                    if let Some(offset) = inclusion.bmff_offset() {
                        ranges_vec.push(RangeInclusive::new(offset, offset));
                        bmff_v2_starts.push(offset);
                    }

                    // add inclusion
                    ranges_vec.push(RangeInclusive::new(inclusion_start, end));
                }
                ranges_vec
            }
        }
        _ => {
            let mut ranges_vec: Vec<RangeInclusive<u64>> = Vec::new();
            let data_end = data_len - 1;
            ranges_vec.push(RangeInclusive::new(0_u64, data_end));

            ranges_vec
        }
    };

    if cfg!(feature = "no_interleaved_io") || cfg!(target_arch = "wasm32") {
        // hash the data for ranges
        for r in ranges {
            let start = r.start();
            let end = r.end();
            let mut chunk_left = end - start + 1;

            // check to see if this range is an BMFF V2 offset to include in the hash
            if bmff_v2_starts.contains(start) && end == start {
                hasher_enum.update(&start.to_be_bytes());
                continue;
            }

            // move to start of range
            data.seek(SeekFrom::Start(*start))?;

            loop {
                let mut chunk = vec![0u8; std::cmp::min(chunk_left as usize, MAX_HASH_BUF)];

                data.read_exact(&mut chunk)?;

                hasher_enum.update(&chunk);

                chunk_left -= chunk.len() as u64;
                if chunk_left == 0 {
                    break;
                }
            }
        }
    } else {
        // hash the data for ranges
        for r in ranges {
            let start = r.start();
            let end = r.end();
            let mut chunk_left = end - start + 1;

            // check to see if this range is an BMFF V2 offset to include in the hash
            if bmff_v2_starts.contains(start) && end == start {
                hasher_enum.update(&start.to_be_bytes());
                continue;
            }

            // move to start of range
            data.seek(SeekFrom::Start(*start))?;

            let mut chunk = vec![0u8; std::cmp::min(chunk_left as usize, MAX_HASH_BUF)];
            data.read_exact(&mut chunk)?;

            loop {
                let (tx, rx) = std::sync::mpsc::channel();

                chunk_left -= chunk.len() as u64;

                std::thread::spawn(move || {
                    hasher_enum.update(&chunk);
                    tx.send(hasher_enum).unwrap_or_default();
                });

                // are we done
                if chunk_left == 0 {
                    hasher_enum = match rx.recv() {
                        Ok(hasher) => hasher,
                        Err(_) => return Err(Error::ThreadReceiveError),
                    };
                    break;
                }

                // read next chunk while we wait for hash
                let mut next_chunk = vec![0u8; std::cmp::min(chunk_left as usize, MAX_HASH_BUF)];
                data.read_exact(&mut next_chunk)?;

                hasher_enum = match rx.recv() {
                    Ok(hasher) => hasher,
                    Err(_) => return Err(Error::ThreadReceiveError),
                };

                chunk = next_chunk;
            }
        }
    }

    // return the hash
    Ok(Hasher::finalize(hasher_enum))
}

// verify the hash using the specified algorithm
pub fn verify_by_alg(
    alg: &str,
    hash: &[u8],
    data: &[u8],
    exclusions: Option<Vec<HashRange>>,
) -> bool {
    // hash with the same algorithm as target
    let data_hash = hash_by_alg(alg, data, exclusions);
    vec_compare(hash, &data_hash)
}

// verify the hash using the specified algorithm
pub fn verify_asset_by_alg(
    alg: &str,
    hash: &[u8],
    asset_path: &Path,
    exclusions: Option<Vec<HashRange>>,
) -> bool {
    // hash with the same algorithm as target
    if let Ok(data_hash) = hash_asset_by_alg(alg, asset_path, exclusions) {
        vec_compare(hash, &data_hash)
    } else {
        false
    }
}

pub fn verify_stream_by_alg<R>(
    alg: &str,
    hash: &[u8],
    reader: &mut R,
    hash_range: Option<Vec<HashRange>>,
    is_exclusion: bool,
) -> bool
where
    R: Read + Seek + ?Sized,
{
    if let Ok(data_hash) = hash_stream_by_alg(alg, reader, hash_range, is_exclusion) {
        vec_compare(hash, &data_hash)
    } else {
        false
    }
}

// Used by Merkle tree calculations to generate the pair wise hash
pub fn concat_and_hash(alg: &str, left: &[u8], right: Option<&[u8]>) -> Vec<u8> {
    let mut temp = left.to_vec();

    if let Some(r) = right {
        temp.append(&mut r.to_vec())
    }

    hash_by_alg(alg, &temp, None)
}