libavb_rs: unlock challenge generation

Implement `cert_generate_unlock_challenge()`.

Bug: b/320543206
Test: atest libavb_rs_test libavb_rs_uuid_test libavb_rs_unittest libavb_rs_uuid_unittest
Change-Id: I0ad0302cd230f9d19b60cb3c8ff693d933db9633

4 files changed, 211 insertions, 11 deletions

diff --git a/rust/src/cert.rs b/rust/src/cert.rs
index cc43367..2d673ad 100644
--- a/rust/src/cert.rs
+++ b/rust/src/cert.rs
@@ -83,9 +83,11 @@
 //! # Internally, the device calls `cert_validate_unlock_credential()` to verify the credential.
 //! ```
 
-use crate::{error::io_enum_to_result, ops, IoError, IoResult, Ops};
-use avb_bindgen::avb_cert_validate_vbmeta_public_key;
-use core::pin::pin;
+use crate::{error::io_enum_to_result, ops, IoError, IoResult, Ops, PublicKeyForPartitionInfo};
+use avb_bindgen::{avb_cert_generate_unlock_challenge, avb_cert_validate_vbmeta_public_key};
+use core::{ffi::CStr, pin::pin};
+#[cfg(feature = "uuid")]
+use uuid::Uuid;
 
 /// libavb_cert permanent attributes.
 pub use avb_bindgen::AvbCertPermanentAttributes as CertPermanentAttributes;
@@ -249,15 +251,29 @@ pub fn cert_validate_vbmeta_public_key(
 /// The user can sign the resulting token via `avbtool make_cert_unlock_credential`.
 ///
 /// # Arguments
-/// * `cert_ops`: the `CertOps` callback implementations.
+/// * `cert_ops`: the `CertOps` callback implementations; base `Ops` are not required here.
 ///
 /// # Returns
 /// The challenge to sign with the PUK, or `IoError` on `cert_ops` failure.
-pub fn cert_generate_unlock_challenge(
-    _cert_ops: &mut dyn CertOps,
-) -> IoResult<CertUnlockChallenge> {
-    // TODO(b/320543206): implement
-    Err(IoError::NotImplemented)
+pub fn cert_generate_unlock_challenge(cert_ops: &mut dyn CertOps) -> IoResult<CertUnlockChallenge> {
+    // `OpsBridge` requires a full `Ops` object, so we wrap `cert_ops` in a do-nothing `Ops`
+    // implementation. This is simpler than teaching `OpsBridge` to handle the cert-only case.
+    let mut ops = CertOnlyOps { cert_ops };
+    let ops_bridge = pin!(ops::OpsBridge::new(&mut ops));
+    let mut challenge = CertUnlockChallenge::default();
+    io_enum_to_result(
+        // SAFETY:
+        // * `ops_bridge.init_and_get_c_ops()` gives us a valid `AvbOps` with cert.
+        // * `challenge` is a valid C-compatible `CertUnlockChallenge`.
+        // * this function does not retain references to any of these arguments.
+        unsafe {
+            avb_cert_generate_unlock_challenge(
+                ops_bridge.init_and_get_c_ops().cert_ops,
+                &mut challenge,
+            )
+        },
+    )?;
+    Ok(challenge)
 }
 
 /// Validates a signed credential for authenticated unlock.
@@ -283,3 +299,77 @@ pub fn cert_validate_unlock_credential(
     // TODO(b/320543206): implement
     Err(IoError::NotImplemented)
 }
+
+/// An `Ops` implementation that only provides the `cert_ops()` callback.
+struct CertOnlyOps<'a> {
+    cert_ops: &'a mut dyn CertOps,
+}
+
+impl<'a> Ops<'static> for CertOnlyOps<'a> {
+    fn read_from_partition(
+        &mut self,
+        _partition: &CStr,
+        _offset: i64,
+        _buffer: &mut [u8],
+    ) -> IoResult<usize> {
+        Err(IoError::NotImplemented)
+    }
+
+    fn validate_vbmeta_public_key(
+        &mut self,
+        _public_key: &[u8],
+        _public_key_metadata: Option<&[u8]>,
+    ) -> IoResult<bool> {
+        Err(IoError::NotImplemented)
+    }
+
+    fn read_rollback_index(&mut self, _rollback_index_location: usize) -> IoResult<u64> {
+        Err(IoError::NotImplemented)
+    }
+
+    fn write_rollback_index(
+        &mut self,
+        _rollback_index_location: usize,
+        _index: u64,
+    ) -> IoResult<()> {
+        Err(IoError::NotImplemented)
+    }
+
+    fn read_is_device_unlocked(&mut self) -> IoResult<bool> {
+        Err(IoError::NotImplemented)
+    }
+
+    #[cfg(feature = "uuid")]
+    fn get_unique_guid_for_partition(&mut self, _partition: &CStr) -> IoResult<Uuid> {
+        Err(IoError::NotImplemented)
+    }
+
+    fn get_size_of_partition(&mut self, _partition: &CStr) -> IoResult<u64> {
+        Err(IoError::NotImplemented)
+    }
+
+    fn read_persistent_value(&mut self, _name: &CStr, _value: &mut [u8]) -> IoResult<usize> {
+        Err(IoError::NotImplemented)
+    }
+
+    fn write_persistent_value(&mut self, _name: &CStr, _value: &[u8]) -> IoResult<()> {
+        Err(IoError::NotImplemented)
+    }
+
+    fn erase_persistent_value(&mut self, _name: &CStr) -> IoResult<()> {
+        Err(IoError::NotImplemented)
+    }
+
+    fn validate_public_key_for_partition(
+        &mut self,
+        _partition: &CStr,
+        _public_key: &[u8],
+        _public_key_metadata: Option<&[u8]>,
+    ) -> IoResult<PublicKeyForPartitionInfo> {
+        Err(IoError::NotImplemented)
+    }
+
+    fn cert_ops(&mut self) -> Option<&mut dyn CertOps> {
+        Some(self.cert_ops)
+    }
+}
diff --git a/rust/src/ops.rs b/rust/src/ops.rs
index 5263471..d8c1b16 100644
--- a/rust/src/ops.rs
+++ b/rust/src/ops.rs
@@ -330,7 +330,7 @@ impl<'o, 'p> OpsBridge<'o, 'p> {
                 read_permanent_attributes: Some(read_permanent_attributes),
                 read_permanent_attributes_hash: Some(read_permanent_attributes_hash),
                 set_key_version: Some(set_key_version),
-                get_random: None, // TODO(b/320543206): implement.
+                get_random: Some(get_random),
             },
             rust_ops: ops,
             _pin: PhantomPinned,
@@ -1298,3 +1298,46 @@ unsafe fn try_set_key_version(
     cert_ops.set_key_version(rollback_index_location, key_version);
     Ok(())
 }
+
+/// Wraps a callback to convert the given `IoResult<>` to `None` for libavb.
+///
+/// See corresponding `try_*` function docs.
+unsafe extern "C" fn get_random(
+    cert_ops: *mut AvbCertOps,
+    num_bytes: usize,
+    output: *mut u8,
+) -> AvbIOResult {
+    result_to_io_enum(
+        // SAFETY: see corresponding `try_*` function safety documentation.
+        unsafe { try_get_random(cert_ops, num_bytes, output) },
+    )
+}
+
+/// Bounces the C callback into the user-provided Rust implementation.
+///
+/// # Safety
+/// * `cert_ops` must have been created via `ScopedAvbOps`.
+/// * `output` must point to a valid buffer of size `num_bytes` that we have exclusive access to.
+unsafe fn try_get_random(
+    cert_ops: *mut AvbCertOps,
+    num_bytes: usize,
+    output: *mut u8,
+) -> IoResult<()> {
+    check_nonnull(output)?;
+    if num_bytes == 0 {
+        return Ok(());
+    }
+
+    // SAFETY:
+    // * we've checked that the pointer is non-NULL.
+    // * libavb gives us a properly-allocated `output` with size `num_bytes`.
+    // * we only access the contents via the returned slice.
+    // * the returned slice is not held past the scope of this callback.
+    let output = unsafe { slice::from_raw_parts_mut(output, num_bytes) };
+
+    // SAFETY:
+    // * we only use `cert_ops` objects created via `ScopedAvbOps` as required.
+    // * `cert_ops` is only extracted once and is dropped at the end of the callback.
+    let cert_ops = unsafe { as_cert_ops(cert_ops) }?;
+    cert_ops.get_random(output)
+}
diff --git a/rust/tests/cert_tests.rs b/rust/tests/cert_tests.rs
index e18dc83..5921b7a 100644
--- a/rust/tests/cert_tests.rs
+++ b/rust/tests/cert_tests.rs
@@ -21,7 +21,8 @@ use crate::{
     verify_one_image_one_vbmeta,
 };
 use avb::{
-    CertPermanentAttributes, SlotVerifyError, CERT_PIK_VERSION_LOCATION, CERT_PSK_VERSION_LOCATION,
+    cert_generate_unlock_challenge, CertPermanentAttributes, IoError, SlotVerifyError,
+    CERT_PIK_VERSION_LOCATION, CERT_PSK_VERSION_LOCATION,
 };
 use hex::decode;
 use std::{collections::HashMap, fs};
@@ -58,6 +59,28 @@ fn build_test_cert_ops_one_image_one_vbmeta<'a>() -> TestOps<'a> {
     ops
 }
 
+/// Enough fake RNG data to generate a single unlock challenge.
+const UNLOCK_CHALLENGE_FAKE_RNG: [u8; 16] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15];
+
+/// Returns a `TestOps` with only enough configuration to generate a single unlock challenge.
+///
+/// The generated unlock challenge will have:
+/// * permanent attributes sourced from the contents of `TEST_CERT_PERMANENT_ATTRIBUTES_PATH`.
+/// * RNG sourced from `UNLOCK_CHALLENGE_FAKE_RNG`.
+fn build_test_cert_ops_unlock_challenge_only<'a>() -> TestOps<'a> {
+    let mut ops = TestOps::default();
+
+    // Permanent attributes are needed for the embedded product ID.
+    let perm_attr_bytes = fs::read(TEST_CERT_PERMANENT_ATTRIBUTES_PATH).unwrap();
+    ops.cert_permanent_attributes =
+        Some(CertPermanentAttributes::read_from(&perm_attr_bytes[..]).unwrap());
+
+    // Fake RNG for unlock challenge generation.
+    ops.cert_fake_rng = UNLOCK_CHALLENGE_FAKE_RNG.into();
+
+    ops
+}
+
 #[test]
 fn cert_verify_succeeds() {
     let mut ops = build_test_cert_ops_one_image_one_vbmeta();
@@ -130,3 +153,43 @@ fn cert_verify_fails_with_bad_permanent_attributes_hash() {
 
     assert_eq!(result.unwrap_err(), SlotVerifyError::PublicKeyRejected);
 }
+
+#[test]
+fn cert_generate_unlock_challenge_succeeds() {
+    let mut ops = build_test_cert_ops_unlock_challenge_only();
+
+    let challenge = cert_generate_unlock_challenge(&mut ops).unwrap();
+
+    // Make sure the challenge token used our cert callback data correctly.
+    assert_eq!(
+        challenge.product_id_hash,
+        &decode(TEST_CERT_PRODUCT_ID_HASH_HEX).unwrap()[..]
+    );
+    assert_eq!(challenge.challenge, UNLOCK_CHALLENGE_FAKE_RNG);
+}
+
+#[test]
+fn cert_generate_unlock_challenge_fails_without_permanent_attributes() {
+    let mut ops = build_test_cert_ops_unlock_challenge_only();
+
+    // Challenge generation should fail without the product ID provided by the permanent attributes.
+    ops.cert_permanent_attributes = None;
+
+    assert_eq!(
+        cert_generate_unlock_challenge(&mut ops).unwrap_err(),
+        IoError::Io
+    );
+}
+
+#[test]
+fn cert_generate_unlock_challenge_fails_insufficient_rng() {
+    let mut ops = build_test_cert_ops_unlock_challenge_only();
+
+    // Remove a byte of RNG so there isn't enough.
+    ops.cert_fake_rng.pop();
+
+    assert_eq!(
+        cert_generate_unlock_challenge(&mut ops).unwrap_err(),
+        IoError::Io
+    );
+}
diff --git a/rust/tests/test_data.rs b/rust/tests/test_data.rs
index 683b869..3db13d5 100644
--- a/rust/tests/test_data.rs
+++ b/rust/tests/test_data.rs
@@ -65,3 +65,7 @@ pub const TEST_CERT_PSK_VERSION: u64 = 42;
 // $ sha256sum external/avb/test/data/cert_permanent_attributes.bin
 pub const TEST_CERT_PERMANENT_ATTRIBUTES_HASH_HEX: &str =
     "55419e1affff153b58f65ce8a5313a71d2a83a00d0abae10a25b9a8e493d04f7";
+
+// $ sha256sum external/avb/test/data/cert_product_id.bin
+pub const TEST_CERT_PRODUCT_ID_HASH_HEX: &str =
+    "374708fff7719dd5979ec875d56cd2286f6d3cf7ec317a3b25632aab28ec37bb";