qg/src/digest.rs

// Copyright 2023 The Pigweed Authors
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not
// use this file except in compliance with the License. You may obtain a copy of
// the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations under
// the License.

use std::str::FromStr;

use sha2::Digest as Sha2Digest;

use crate::Error;

#[derive(Debug, PartialEq)]
#[allow(clippy::module_name_repetitions)]
pub enum ExpectedDigest {
    Ignore,
    Sha256(String),
}

impl FromStr for ExpectedDigest {
    type Err = Error;

    /// Parses a string of the format "algorithm:digest" or the special value
    /// "ignore".
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s == "ignore" {
            return Ok(Self::Ignore);
        }

        let parts: Vec<_> = s.split(':').collect();
        if parts.len() != 2 || parts[1].is_empty() {
            // TODO(frolv): Missing digest string.
            return Err(Error::GenericErrorPlaceholder);
        }

        match *parts.first().expect("guaranteed to have two elements") {
            "sha256" => {
                let hash = parts[1];
                if !hash.chars().all(|c| c.is_ascii_hexdigit()) {
                    // TODO(frolv): Invalid sha256 digest.
                    return Err(Error::GenericErrorPlaceholder);
                }
                Ok(Self::Sha256(hash.to_string()))
            }
            _ => Err(Error::GenericErrorPlaceholder),
        }
    }
}

impl ExpectedDigest {
    pub fn verifier(&self) -> Verifier {
        match self {
            Self::Ignore => Verifier::Ignore,
            Self::Sha256(s) => Verifier::Sha256 {
                checksum: sha2::Sha256::new(),
                expected: s.as_str(),
            },
        }
    }
}

pub enum Verifier<'a> {
    Ignore,

    Sha256 {
        checksum: sha2::Sha256,
        expected: &'a str,
    },
}

impl<'a> Verifier<'a> {
    pub fn update(&mut self, bytes: &[u8]) {
        match self {
            Self::Ignore => (),
            Self::Sha256 { checksum, .. } => checksum.update(bytes),
        }
    }

    pub fn reset(&mut self) {
        match self {
            Self::Ignore => (),
            Self::Sha256 { checksum, .. } => checksum.reset(),
        }
    }

    pub fn verify(self) -> bool {
        match self {
            Self::Ignore => true,
            Self::Sha256 { checksum, expected } => format!("{:x}", checksum.finalize()) == expected,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_valid() {
        assert_eq!(
            "sha256:0123456789abcdefABCDEF"
                .parse::<ExpectedDigest>()
                .unwrap(),
            ExpectedDigest::Sha256("0123456789abcdefABCDEF".into())
        );
        assert_eq!(
            "ignore".parse::<ExpectedDigest>().unwrap(),
            ExpectedDigest::Ignore
        );
    }

    #[test]
    fn test_parse_invalid() {
        assert!(matches!(
            "".parse::<ExpectedDigest>().unwrap_err(),
            Error::GenericErrorPlaceholder,
        ));
        assert!(matches!(
            ":".parse::<ExpectedDigest>().unwrap_err(),
            Error::GenericErrorPlaceholder,
        ));
        assert!(matches!(
            "abcdef".parse::<ExpectedDigest>().unwrap_err(),
            Error::GenericErrorPlaceholder,
        ));
        assert!(matches!(
            ":abcdef".parse::<ExpectedDigest>().unwrap_err(),
            Error::GenericErrorPlaceholder,
        ));
        assert!(matches!(
            "md6:abcdef".parse::<ExpectedDigest>().unwrap_err(),
            Error::GenericErrorPlaceholder,
        ));
        assert!(matches!(
            "sha256:".parse::<ExpectedDigest>().unwrap_err(),
            Error::GenericErrorPlaceholder,
        ));
        assert!(matches!(
            "sha256:abc:def".parse::<ExpectedDigest>().unwrap_err(),
            Error::GenericErrorPlaceholder,
        ));
        assert!(matches!(
            "sha256:0123456789abcdefABCDEFg"
                .parse::<ExpectedDigest>()
                .unwrap_err(),
            Error::GenericErrorPlaceholder,
        ));
    }
}