[apkverify] Introduce enums for signature and content digest algorithms

Test: atest libapkverify.test
Bug: 246254355
Change-Id: I18adcb8116341f52a8ca6a7489c589a0d2096c2e
diff --git a/libs/apkverify/Android.bp b/libs/apkverify/Android.bp
index 50d7a60..9bb8f8e 100644
--- a/libs/apkverify/Android.bp
+++ b/libs/apkverify/Android.bp
@@ -13,9 +13,11 @@
         "libbyteorder",
         "libbytes",
         "liblog_rust",
+        "libnum_traits",
         "libopenssl",
         "libzip",
     ],
+    proc_macros: ["libnum_derive"],
 }
 
 rust_library {
diff --git a/libs/apkverify/src/algorithms.rs b/libs/apkverify/src/algorithms.rs
new file mode 100644
index 0000000..f15f0a4
--- /dev/null
+++ b/libs/apkverify/src/algorithms.rs
@@ -0,0 +1,107 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * 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
+ *
+ *      http://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.
+ */
+
+//! Algorithms used for APK Signature Scheme.
+
+use anyhow::{bail, Result};
+use num_derive::FromPrimitive;
+use openssl::hash::MessageDigest;
+use std::cmp::Ordering;
+
+/// [Signature Algorithm IDs]: https://source.android.com/docs/security/apksigning/v2#signature-algorithm-ids
+///
+/// Some of the algorithms are not implemented. See b/197052981.
+#[derive(Clone, Debug, Eq, FromPrimitive)]
+#[repr(u32)]
+pub enum SignatureAlgorithmID {
+    RsaPssWithSha256 = 0x0101,
+    RsaPssWithSha512 = 0x0102,
+    RsaPkcs1V15WithSha256 = 0x0103,
+    RsaPkcs1V15WithSha512 = 0x0104,
+    EcdsaWithSha256 = 0x0201,
+    EcdsaWithSha512 = 0x0202,
+    DsaWithSha256 = 0x0301,
+    VerityRsaPkcs1V15WithSha256 = 0x0421,
+    VerityEcdsaWithSha256 = 0x0423,
+    VerityDsaWithSha256 = 0x0425,
+}
+
+impl Ord for SignatureAlgorithmID {
+    /// Ranks the signature algorithm according to the corresponding content
+    /// digest algorithm's rank.
+    fn cmp(&self, other: &Self) -> Ordering {
+        self.to_content_digest_algorithm().cmp(&other.to_content_digest_algorithm())
+    }
+}
+
+impl PartialOrd for SignatureAlgorithmID {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl PartialEq for SignatureAlgorithmID {
+    fn eq(&self, other: &Self) -> bool {
+        self.cmp(other) == Ordering::Equal
+    }
+}
+
+impl SignatureAlgorithmID {
+    pub(crate) fn to_content_digest_algorithm(&self) -> ContentDigestAlgorithm {
+        match self {
+            SignatureAlgorithmID::RsaPssWithSha256
+            | SignatureAlgorithmID::RsaPkcs1V15WithSha256
+            | SignatureAlgorithmID::EcdsaWithSha256
+            | SignatureAlgorithmID::DsaWithSha256 => ContentDigestAlgorithm::ChunkedSha256,
+            SignatureAlgorithmID::RsaPssWithSha512
+            | SignatureAlgorithmID::RsaPkcs1V15WithSha512
+            | SignatureAlgorithmID::EcdsaWithSha512 => ContentDigestAlgorithm::ChunkedSha512,
+            SignatureAlgorithmID::VerityRsaPkcs1V15WithSha256
+            | SignatureAlgorithmID::VerityEcdsaWithSha256
+            | SignatureAlgorithmID::VerityDsaWithSha256 => {
+                ContentDigestAlgorithm::VerityChunkedSha256
+            }
+        }
+    }
+}
+
+/// The rank of the content digest algorithm in this enum is used to help pick
+/// v4 apk digest.
+/// According to APK Signature Scheme v4, [apk digest] is the first available
+/// content digest of the highest rank (rank N).
+///
+/// This rank was also used for step 3a of the v3 signature verification.
+///
+/// [apk digest]: https://source.android.com/docs/security/features/apksigning/v4#apk-digest
+/// [v3 verification]: https://source.android.com/docs/security/apksigning/v3#v3-verification
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
+pub(crate) enum ContentDigestAlgorithm {
+    ChunkedSha256 = 1,
+    VerityChunkedSha256,
+    ChunkedSha512,
+}
+
+impl ContentDigestAlgorithm {
+    pub(crate) fn new_message_digest(&self) -> Result<MessageDigest> {
+        match self {
+            ContentDigestAlgorithm::ChunkedSha256 => Ok(MessageDigest::sha256()),
+            ContentDigestAlgorithm::ChunkedSha512 => Ok(MessageDigest::sha512()),
+            ContentDigestAlgorithm::VerityChunkedSha256 => {
+                bail!("TODO(b/197052981): CONTENT_DIGEST_VERITY_CHUNKED_SHA256 is not implemented")
+            }
+        }
+    }
+}
diff --git a/libs/apkverify/src/lib.rs b/libs/apkverify/src/lib.rs
index c5aa9e5..f6c1a21 100644
--- a/libs/apkverify/src/lib.rs
+++ b/libs/apkverify/src/lib.rs
@@ -16,6 +16,7 @@
 
 //! Verifies APK/APEX signing with v2/v3 scheme
 
+mod algorithms;
 mod bytes_ext;
 mod sigutil;
 #[allow(dead_code)]
diff --git a/libs/apkverify/src/sigutil.rs b/libs/apkverify/src/sigutil.rs
index 6a7ce32..ea6d63a 100644
--- a/libs/apkverify/src/sigutil.rs
+++ b/libs/apkverify/src/sigutil.rs
@@ -16,19 +16,21 @@
 
 //! Utilities for Signature Verification
 
-use anyhow::{anyhow, bail, ensure, Error, Result};
+use anyhow::{anyhow, ensure, Error, Result};
 use byteorder::{LittleEndian, ReadBytesExt};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
+use num_traits::FromPrimitive;
 use openssl::hash::{DigestBytes, Hasher, MessageDigest};
 use std::cmp::min;
 use std::io::{self, Cursor, ErrorKind, Read, Seek, SeekFrom, Take};
 
+use crate::algorithms::SignatureAlgorithmID;
 use crate::ziputil::{set_central_directory_offset, zip_sections};
 
 const APK_SIG_BLOCK_MIN_SIZE: u32 = 32;
 const APK_SIG_BLOCK_MAGIC: u128 = 0x3234206b636f6c4220676953204b5041;
 
-// TODO(jooyung): introduce type
+// TODO(b/246254355): Migrates usages of raw signature algorithm id to the enum.
 pub const SIGNATURE_RSA_PSS_WITH_SHA256: u32 = 0x0101;
 pub const SIGNATURE_RSA_PSS_WITH_SHA512: u32 = 0x0102;
 pub const SIGNATURE_RSA_PKCS1_V1_5_WITH_SHA256: u32 = 0x0103;
@@ -40,13 +42,6 @@
 pub const SIGNATURE_VERITY_ECDSA_WITH_SHA256: u32 = 0x0423;
 pub const SIGNATURE_VERITY_DSA_WITH_SHA256: u32 = 0x0425;
 
-// TODO(jooyung): introduce type
-const CONTENT_DIGEST_CHUNKED_SHA256: u32 = 1;
-const CONTENT_DIGEST_CHUNKED_SHA512: u32 = 2;
-const CONTENT_DIGEST_VERITY_CHUNKED_SHA256: u32 = 3;
-#[allow(unused)]
-const CONTENT_DIGEST_SHA256: u32 = 4;
-
 const CHUNK_SIZE_BYTES: u64 = 1024 * 1024;
 
 /// The [APK structure] has four major sections:
@@ -97,6 +92,9 @@
     ///    order the chunks appear in the APK.
     /// (see https://source.android.com/security/apksigning/v2#integrity-protected-contents)
     pub fn compute_digest(&mut self, signature_algorithm_id: u32) -> Result<Vec<u8>> {
+        // TODO(b/246254355): Passes the enum SignatureAlgorithmID directly to this method.
+        let signature_algorithm_id = SignatureAlgorithmID::from_u32(signature_algorithm_id)
+            .ok_or_else(|| anyhow!("Unsupported algorithm ID: {}", signature_algorithm_id))?;
         let digester = Digester::new(signature_algorithm_id)?;
 
         let mut digests_of_chunks = BytesMut::new();
@@ -163,30 +161,22 @@
 }
 
 struct Digester {
-    algorithm: MessageDigest,
+    message_digest: MessageDigest,
 }
 
 const CHUNK_HEADER_TOP: &[u8] = &[0x5a];
 const CHUNK_HEADER_MID: &[u8] = &[0xa5];
 
 impl Digester {
-    fn new(signature_algorithm_id: u32) -> Result<Digester> {
-        let digest_algorithm_id = to_content_digest_algorithm(signature_algorithm_id)?;
-        let algorithm = match digest_algorithm_id {
-            CONTENT_DIGEST_CHUNKED_SHA256 => MessageDigest::sha256(),
-            CONTENT_DIGEST_CHUNKED_SHA512 => MessageDigest::sha512(),
-            // TODO(jooyung): implement
-            CONTENT_DIGEST_VERITY_CHUNKED_SHA256 => {
-                bail!("TODO(b/190343842): CONTENT_DIGEST_VERITY_CHUNKED_SHA256: not implemented")
-            }
-            _ => bail!("Unknown digest algorithm: {}", digest_algorithm_id),
-        };
-        Ok(Digester { algorithm })
+    fn new(signature_algorithm_id: SignatureAlgorithmID) -> Result<Digester> {
+        let message_digest =
+            signature_algorithm_id.to_content_digest_algorithm().new_message_digest()?;
+        Ok(Digester { message_digest })
     }
 
     // v2/v3 digests are computed after prepending "header" byte and "size" info.
     fn digest(&self, data: &[u8], header: &[u8], size: u32) -> Result<DigestBytes> {
-        let mut hasher = Hasher::new(self.algorithm)?;
+        let mut hasher = Hasher::new(self.message_digest)?;
         hasher.update(header)?;
         hasher.update(&size.to_le_bytes())?;
         hasher.update(data)?;
@@ -259,55 +249,6 @@
     Err(Error::new(io::Error::from(ErrorKind::NotFound)).context(context))
 }
 
-pub fn is_supported_signature_algorithm(algorithm_id: u32) -> bool {
-    matches!(
-        algorithm_id,
-        SIGNATURE_RSA_PSS_WITH_SHA256
-            | SIGNATURE_RSA_PSS_WITH_SHA512
-            | SIGNATURE_RSA_PKCS1_V1_5_WITH_SHA256
-            | SIGNATURE_RSA_PKCS1_V1_5_WITH_SHA512
-            | SIGNATURE_ECDSA_WITH_SHA256
-            | SIGNATURE_ECDSA_WITH_SHA512
-            | SIGNATURE_DSA_WITH_SHA256
-            | SIGNATURE_VERITY_RSA_PKCS1_V1_5_WITH_SHA256
-            | SIGNATURE_VERITY_ECDSA_WITH_SHA256
-            | SIGNATURE_VERITY_DSA_WITH_SHA256
-    )
-}
-
-fn to_content_digest_algorithm(algorithm_id: u32) -> Result<u32> {
-    match algorithm_id {
-        SIGNATURE_RSA_PSS_WITH_SHA256
-        | SIGNATURE_RSA_PKCS1_V1_5_WITH_SHA256
-        | SIGNATURE_ECDSA_WITH_SHA256
-        | SIGNATURE_DSA_WITH_SHA256 => Ok(CONTENT_DIGEST_CHUNKED_SHA256),
-        SIGNATURE_RSA_PSS_WITH_SHA512
-        | SIGNATURE_RSA_PKCS1_V1_5_WITH_SHA512
-        | SIGNATURE_ECDSA_WITH_SHA512 => Ok(CONTENT_DIGEST_CHUNKED_SHA512),
-        SIGNATURE_VERITY_RSA_PKCS1_V1_5_WITH_SHA256
-        | SIGNATURE_VERITY_ECDSA_WITH_SHA256
-        | SIGNATURE_VERITY_DSA_WITH_SHA256 => Ok(CONTENT_DIGEST_VERITY_CHUNKED_SHA256),
-        _ => bail!("Unknown signature algorithm: {}", algorithm_id),
-    }
-}
-
-/// This method is used to help pick v4 apk digest. According to APK Signature
-/// Scheme v4, apk digest is the first available content digest of the highest
-/// rank (rank N).
-///
-/// This rank was also used for step 3a of the v3 signature verification.
-///
-/// [v3 verification]: https://source.android.com/docs/security/apksigning/v3#v3-verification
-pub fn get_signature_algorithm_rank(algo: u32) -> Result<u32> {
-    let content_digest = to_content_digest_algorithm(algo)?;
-    match content_digest {
-        CONTENT_DIGEST_CHUNKED_SHA256 => Ok(0),
-        CONTENT_DIGEST_VERITY_CHUNKED_SHA256 => Ok(1),
-        CONTENT_DIGEST_CHUNKED_SHA512 => Ok(2),
-        _ => bail!("Unknown digest algorithm: {}", content_digest),
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/libs/apkverify/src/v3.rs b/libs/apkverify/src/v3.rs
index d429ba1..c86696f 100644
--- a/libs/apkverify/src/v3.rs
+++ b/libs/apkverify/src/v3.rs
@@ -20,6 +20,7 @@
 
 use anyhow::{anyhow, bail, ensure, Context, Result};
 use bytes::Bytes;
+use num_traits::FromPrimitive;
 use openssl::hash::MessageDigest;
 use openssl::pkey::{self, PKey};
 use openssl::rsa::Padding;
@@ -30,6 +31,7 @@
 use std::ops::Range;
 use std::path::Path;
 
+use crate::algorithms::SignatureAlgorithmID;
 use crate::bytes_ext::{BytesExt, LengthPrefixed, ReadFromBytes};
 use crate::sigutil::*;
 
@@ -71,6 +73,7 @@
 
 #[derive(Debug)]
 struct Signature {
+    /// TODO(b/246254355): Change the type of signature_algorithm_id to SignatureAlgorithmID
     signature_algorithm_id: u32,
     signature: LengthPrefixed<Bytes>,
 }
@@ -143,8 +146,8 @@
         Ok(self
             .signatures
             .iter()
-            .filter(|sig| is_supported_signature_algorithm(sig.signature_algorithm_id))
-            .max_by_key(|sig| get_signature_algorithm_rank(sig.signature_algorithm_id).unwrap())
+            .filter(|sig| SignatureAlgorithmID::from_u32(sig.signature_algorithm_id).is_some())
+            .max_by_key(|sig| SignatureAlgorithmID::from_u32(sig.signature_algorithm_id).unwrap())
             .ok_or_else(|| anyhow!("No supported signatures found"))?)
     }