Changing traits to support opaque keys for kek and kak

Modifying traits to support non-explicit keys for key encryption and
key agreement keys.

Bug: 253948020
Test: build.py, run rust unittests
Change-Id: I1a266cdcdada1a5e187465c402dac6b92decde75
diff --git a/keys.rs b/keys.rs
index 54dd14c..6e16e06 100644
--- a/keys.rs
+++ b/keys.rs
@@ -226,7 +226,7 @@
 //       if the IP block never releases the key. KeyMaterial type fixes that issue by including
 //       Opaque keys, but RawKeys are not included in KeyMaterial.
 impl kmr_ta::device::RetrieveKeyMaterial for TrustyKeys {
-    fn root_kek(&self, context: &[u8]) -> Result<crypto::RawKeyMaterial, Error> {
+    fn root_kek(&self, context: &[u8]) -> Result<crypto::OpaqueOr<crypto::hmac::Key>, Error> {
         let context = TrustyKekContext::from_raw(context)?;
         let hwkey_session = Hwkey::open().map_err(|e| {
             km_err!(SecureHwCommunicationFailed, "failed to connect to hwkey: {:?}", e)
@@ -257,7 +257,7 @@
                     })?;
             }
         }
-        Ok(crypto::RawKeyMaterial(key_buffer.to_vec()))
+        Ok(crypto::hmac::Key::new(key_buffer.to_vec()).into())
     }
 
     fn kek_context(&self) -> Result<Vec<u8>, Error> {
@@ -265,7 +265,7 @@
             .to_raw()
     }
 
-    fn kak(&self) -> Result<crypto::aes::Key, Error> {
+    fn kak(&self) -> Result<crypto::OpaqueOr<crypto::aes::Key>, Error> {
         let hwkey_session = Hwkey::open().map_err(|e| {
             km_err!(SecureHwCommunicationFailed, "failed to connect to HwKey: {:?}", e)
         })?;
@@ -276,7 +276,7 @@
             .get_keyslot_data(keyslot, &mut key_buffer)
             .map_err(|e| km_err!(SecureHwCommunicationFailed, "failed to retrieve kak: {:?}", e))?;
         // TODO: check whether `key_buffer` needs truncating to size of `_kak`.
-        Ok(crypto::aes::Key::Aes256(key_buffer))
+        Ok(crypto::aes::Key::Aes256(key_buffer).into())
     }
 
     fn timestamp_token_mac_input(&self, token: &TimeStampToken) -> Result<Vec<u8>, Error> {
@@ -305,6 +305,7 @@
     fn kak_call_returns_key() {
         let trusty_keys = TrustyKeys;
         let kak = trusty_keys.kak().expect("Couldn't retrieve kak");
+        let kak = kmr_common::explicit!(kak).expect("kak should be an explicit key");
 
         expect!(matches!(kak, crypto::aes::Key::Aes256(_)), "Should have received an AES 256 key");
 
@@ -321,11 +322,13 @@
     fn kak_two_calls_returns_same_key() {
         let trusty_keys = TrustyKeys;
 
-        let kak1 = match trusty_keys.kak().expect("Couldn't retrieve kak") {
+        let kak = trusty_keys.kak().expect("Couldn't retrieve kak");
+        let kak1 = match kmr_common::explicit!(kak).expect("kak should be an explicit key") {
             crypto::aes::Key::Aes256(key) => key,
             _ => panic!("Wrong type of key received"),
         };
-        let kak2 = match trusty_keys.kak().expect("Couldn't retrieve kak") {
+        let kak = trusty_keys.kak().expect("Couldn't retrieve kak");
+        let kak2 = match kmr_common::explicit!(kak).expect("kak should be an explicit key") {
             crypto::aes::Key::Aes256(key) => key,
             _ => panic!("Wrong type of key received"),
         };
@@ -338,6 +341,7 @@
         let kek = trusty_keys
             .root_kek(&trusty_keys.kek_context().expect("Couldn't get kek context"))
             .expect("Couldn't get kek");
+        let kek = kmr_common::explicit!(kek).expect("kek should be an explicit key");
 
         // Getting an all 0 key encryption key by chance is not likely if we got a connection to
         // HWKey
@@ -351,12 +355,14 @@
     #[test]
     fn kek_two_calls_returns_same_key() {
         let trusty_keys = TrustyKeys;
-        let kek1 = trusty_keys
+        let kek1 = kmr_common::explicit!(trusty_keys
             .root_kek(&trusty_keys.kek_context().expect("Couldn't get kek context"))
-            .expect("Couldn't get kek");
-        let kek2 = trusty_keys
+            .expect("Couldn't get kek"))
+        .expect("kek should be an explicit key");
+        let kek2 = kmr_common::explicit!(trusty_keys
             .root_kek(&trusty_keys.kek_context().expect("Couldn't get kek context"))
-            .expect("Couldn't get kek");
+            .expect("Couldn't get kek"))
+        .expect("kek should be an explicit key");
 
         expect_eq!(kek1.0, kek2.0, "Calls to root_kek should return the same key");
     }
@@ -384,12 +390,14 @@
         )
         .unwrap();
         let trusty_keys = TrustyKeys;
-        let kek1 = trusty_keys
+        let kek1 = kmr_common::explicit!(trusty_keys
             .root_kek(&context1.to_raw().expect("Couldn't serialize kek1 context"))
-            .expect("Couldn't get kek");
-        let kek2 = trusty_keys
+            .expect("Couldn't get kek"))
+        .expect("kek should be an explicit key");
+        let kek2 = kmr_common::explicit!(trusty_keys
             .root_kek(&context2.to_raw().expect("Couldn't serialize kek2 context"))
-            .expect("Couldn't get kek");
+            .expect("Couldn't get kek"))
+        .expect("kek should be an explicit key");
 
         expect_ne!(kek1.0, kek2.0, "kek keys should be different");
     }
@@ -401,12 +409,14 @@
                 .unwrap();
         let context2 = TrustyKekContext::new(false, None, None).unwrap();
         let trusty_keys = TrustyKeys;
-        let kek1 = trusty_keys
+        let kek1 = kmr_common::explicit!(trusty_keys
             .root_kek(&context1.to_raw().expect("Couldn't serialize kek1 context"))
-            .expect("Couldn't get kek");
-        let kek2 = trusty_keys
+            .expect("Couldn't get kek"))
+        .expect("kek should be an explicit key");
+        let kek2 = kmr_common::explicit!(trusty_keys
             .root_kek(&context2.to_raw().expect("Couldn't serialize kek2 context"))
-            .expect("Couldn't get kek");
+            .expect("Couldn't get kek"))
+        .expect("kek should be an explicit key");
 
         expect_ne!(kek1.0, kek2.0, "kek keys should be different");
     }
diff --git a/keys/legacy.rs b/keys/legacy.rs
index 7051d56..0b0ab95 100644
--- a/keys/legacy.rs
+++ b/keys/legacy.rs
@@ -7,8 +7,8 @@
 use kmr_common::{
     crypto,
     crypto::{aes, OpaqueKeyMaterial, OpaqueOr},
-    get_bool_tag_value, get_opt_tag_value, get_tag_value, keyblob, km_err, tag, try_to_vec,
-    vec_try, Error, FallibleAllocExt,
+    explicit, get_bool_tag_value, get_opt_tag_value, get_tag_value, keyblob, km_err, tag,
+    try_to_vec, vec_try, Error, FallibleAllocExt,
 };
 use kmr_ta::device;
 use kmr_wire::{
@@ -71,13 +71,14 @@
     /// Derive the key encryption key for a keyblob.
     fn derive_kek(
         &self,
-        root_kek: &kmr_common::crypto::RawKeyMaterial,
+        root_kek: &OpaqueOr<kmr_common::crypto::hmac::Key>,
         encrypted_keyblob: &legacy::EncryptedKeyBlob,
         hidden: &[KeyParam],
         sdd_data: Option<(SecureDeletionData, u32)>,
     ) -> Result<crypto::aes::Key, Error> {
         let info = self.build_derivation_info(encrypted_keyblob, hidden, sdd_data)?;
-        let raw_key = self.hkdf.hkdf(&[], &root_kek.0, &info, 256 / 8)?;
+        // Trusty uses explicit keys for legacy keyblobs.
+        let raw_key = self.hkdf.hkdf(&[], &explicit!(root_kek)?.0, &info, 256 / 8)?;
         let aes_key = crypto::aes::Key::Aes256(
             raw_key.try_into().map_err(|_e| km_err!(UnknownError, "unexpected HKDF output len"))?,
         );