Merge "keystore2_test: Enable stress test"
diff --git a/diced/open_dice/Android.bp b/diced/open_dice/Android.bp
new file mode 100644
index 0000000..be8388f
--- /dev/null
+++ b/diced/open_dice/Android.bp
@@ -0,0 +1,36 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_defaults {
+    name: "libdiced_open_dice_defaults",
+    crate_name: "diced_open_dice",
+    srcs: ["src/lib.rs"],
+    static_libs: [
+        "libopen_dice_cbor",
+    ],
+    vendor_available: true,
+    apex_available: [
+        "//apex_available:platform",
+        "com.android.virt",
+    ],
+}
+
+rust_library_rlib {
+    name: "libdiced_open_dice_nostd",
+    defaults: ["libdiced_open_dice_defaults"],
+    rustlibs: [
+        "libopen_dice_cbor_bindgen_nostd",
+    ],
+}
+
+rust_library_rlib {
+    name: "libdiced_open_dice",
+    defaults: ["libdiced_open_dice_defaults"],
+    rustlibs: [
+        "libopen_dice_cbor_bindgen",
+    ],
+    features: [
+        "std",
+    ],
+}
\ No newline at end of file
diff --git a/diced/open_dice/src/dice.rs b/diced/open_dice/src/dice.rs
new file mode 100644
index 0000000..c08cc40
--- /dev/null
+++ b/diced/open_dice/src/dice.rs
@@ -0,0 +1,112 @@
+// Copyright 2023, 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.
+
+//! Structs and functions about the types used in DICE.
+//! This module mirrors the content in open-dice/include/dice/dice.h
+
+use core::ptr;
+use open_dice_cbor_bindgen::{
+    DiceConfigType, DiceInputValues, DiceMode, DICE_HASH_SIZE, DICE_HIDDEN_SIZE,
+    DICE_INLINE_CONFIG_SIZE,
+};
+
+/// The size of a DICE hash.
+pub const HASH_SIZE: usize = DICE_HASH_SIZE as usize;
+/// The size of the DICE hidden value.
+pub const HIDDEN_SIZE: usize = DICE_HIDDEN_SIZE as usize;
+/// The size of a DICE inline config.
+const INLINE_CONFIG_SIZE: usize = DICE_INLINE_CONFIG_SIZE as usize;
+
+/// Array type of hashes used by DICE.
+pub type Hash = [u8; HASH_SIZE];
+/// Array type of additional input.
+pub type Hidden = [u8; HIDDEN_SIZE];
+/// Array type of inline configuration values.
+pub type InlineConfig = [u8; INLINE_CONFIG_SIZE];
+
+/// Configuration descriptor for DICE input values.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum Config<'a> {
+    /// Reference to an inline descriptor.
+    Inline(&'a InlineConfig),
+    /// Reference to a free form descriptor that will be hashed by the implementation.
+    Descriptor(&'a [u8]),
+}
+
+impl Config<'_> {
+    fn dice_config_type(&self) -> DiceConfigType {
+        match self {
+            Self::Inline(_) => DiceConfigType::kDiceConfigTypeInline,
+            Self::Descriptor(_) => DiceConfigType::kDiceConfigTypeDescriptor,
+        }
+    }
+
+    fn inline_config(&self) -> InlineConfig {
+        match self {
+            Self::Inline(inline) => **inline,
+            Self::Descriptor(_) => [0u8; INLINE_CONFIG_SIZE],
+        }
+    }
+
+    fn descriptor_ptr(&self) -> *const u8 {
+        match self {
+            Self::Descriptor(descriptor) => descriptor.as_ptr(),
+            _ => ptr::null(),
+        }
+    }
+
+    fn descriptor_size(&self) -> usize {
+        match self {
+            Self::Descriptor(descriptor) => descriptor.len(),
+            _ => 0,
+        }
+    }
+}
+
+/// Wrap of `DiceInputValues`.
+#[derive(Clone, Debug)]
+pub struct InputValues(DiceInputValues);
+
+impl InputValues {
+    /// Creates a new `InputValues`.
+    pub fn new(
+        code_hash: &Hash,
+        code_descriptor: Option<&[u8]>,
+        config: Config,
+        authority_hash: &Hash,
+        authority_descriptor: Option<&[u8]>,
+        mode: DiceMode,
+        hidden: Option<&Hidden>,
+    ) -> Self {
+        Self(DiceInputValues {
+            code_hash: *code_hash,
+            code_descriptor: code_descriptor.map_or(ptr::null(), |d| d.as_ptr()),
+            code_descriptor_size: code_descriptor.map_or(0, |d| d.len()),
+            config_type: config.dice_config_type(),
+            config_value: config.inline_config(),
+            config_descriptor: config.descriptor_ptr(),
+            config_descriptor_size: config.descriptor_size(),
+            authority_hash: *authority_hash,
+            authority_descriptor: authority_descriptor.map_or(ptr::null(), |d| d.as_ptr()),
+            authority_descriptor_size: authority_descriptor.map_or(0, |d| d.len()),
+            mode,
+            hidden: hidden.map_or([0; HIDDEN_SIZE], |h| *h),
+        })
+    }
+
+    /// Returns a raw pointer to the wrapped `DiceInputValues`.
+    pub fn as_ptr(&self) -> *const DiceInputValues {
+        &self.0 as *const DiceInputValues
+    }
+}
diff --git a/diced/open_dice/src/lib.rs b/diced/open_dice/src/lib.rs
new file mode 100644
index 0000000..96e2569
--- /dev/null
+++ b/diced/open_dice/src/lib.rs
@@ -0,0 +1,22 @@
+// Copyright 2023, 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.
+
+//! Implements safe wrappers around the public API of libopen-dice for
+//! both std and nostd usages.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+mod dice;
+
+pub use dice::{Config, Hash, Hidden, InlineConfig, InputValues, HASH_SIZE, HIDDEN_SIZE};
diff --git a/diced/open_dice_cbor/Android.bp b/diced/open_dice_cbor/Android.bp
index a84190a..3a7f04c 100644
--- a/diced/open_dice_cbor/Android.bp
+++ b/diced/open_dice_cbor/Android.bp
@@ -22,6 +22,7 @@
     srcs: ["lib.rs"],
 
     rustlibs: [
+        "libdiced_open_dice",
         // For ZVec
         "libkeystore2_crypto_rust",
         "libopen_dice_bcc_bindgen",
@@ -46,6 +47,7 @@
     test_suites: ["general-tests"],
     auto_gen_config: true,
     rustlibs: [
+        "libdiced_open_dice",
         "libdiced_sample_inputs",
         "libkeystore2_crypto_rust",
         "libopen_dice_bcc_bindgen",
diff --git a/diced/open_dice_cbor/lib.rs b/diced/open_dice_cbor/lib.rs
index ffb8a48..e5bda89 100644
--- a/diced/open_dice_cbor/lib.rs
+++ b/diced/open_dice_cbor/lib.rs
@@ -24,29 +24,25 @@
 //!     code_hash: [3u8, dice::HASH_SIZE],
 //!     config: dice::ConfigOwned::Descriptor("My descriptor".as_bytes().to_vec()),
 //!     authority_hash: [0u8, dice::HASH_SIZE],
-//!     mode: dice::Mode::Normal,
+//!     mode: dice::DiceMode::kDiceModeNormal,
 //!     hidden: [0u8, dice::HIDDEN_SIZE],
 //! };
 //! let (cdi_attest, cdi_seal, cert_chain) = context
 //!     .main_flow(&parent_cdi_attest, &parent_cdi_seal, &input_values)?;
 //! ```
 
+use diced_open_dice::InlineConfig;
+pub use diced_open_dice::{Config, Hash, Hidden, HASH_SIZE, HIDDEN_SIZE};
 use keystore2_crypto::{zvec, ZVec};
 use open_dice_bcc_bindgen::BccMainFlow;
+pub use open_dice_cbor_bindgen::DiceMode;
 use open_dice_cbor_bindgen::{
-    DiceConfigType, DiceDeriveCdiCertificateId, DiceDeriveCdiPrivateKeySeed,
-    DiceGenerateCertificate, DiceHash, DiceInputValues, DiceKdf, DiceKeypairFromSeed, DiceMainFlow,
-    DiceMode, DiceResult, DiceSign, DiceVerify, DICE_CDI_SIZE, DICE_HASH_SIZE, DICE_HIDDEN_SIZE,
-    DICE_ID_SIZE, DICE_INLINE_CONFIG_SIZE, DICE_PRIVATE_KEY_SEED_SIZE, DICE_PRIVATE_KEY_SIZE,
+    DiceDeriveCdiCertificateId, DiceDeriveCdiPrivateKeySeed, DiceGenerateCertificate, DiceHash,
+    DiceInputValues, DiceKdf, DiceKeypairFromSeed, DiceMainFlow, DiceResult, DiceSign, DiceVerify,
+    DICE_CDI_SIZE, DICE_ID_SIZE, DICE_PRIVATE_KEY_SEED_SIZE, DICE_PRIVATE_KEY_SIZE,
     DICE_PUBLIC_KEY_SIZE, DICE_SIGNATURE_SIZE,
 };
 use open_dice_cbor_bindgen::{
-    DiceConfigType_kDiceConfigTypeDescriptor as DICE_CONFIG_TYPE_DESCRIPTOR,
-    DiceConfigType_kDiceConfigTypeInline as DICE_CONFIG_TYPE_INLINE,
-    DiceMode_kDiceModeDebug as DICE_MODE_DEBUG,
-    DiceMode_kDiceModeMaintenance as DICE_MODE_RECOVERY,
-    DiceMode_kDiceModeNormal as DICE_MODE_NORMAL,
-    DiceMode_kDiceModeNotInitialized as DICE_MODE_NOT_CONFIGURED,
     DiceResult_kDiceResultBufferTooSmall as DICE_RESULT_BUFFER_TOO_SMALL,
     DiceResult_kDiceResultInvalidInput as DICE_RESULT_INVALID_INPUT,
     DiceResult_kDiceResultOk as DICE_RESULT_OK,
@@ -54,12 +50,6 @@
 };
 use std::ffi::{c_void, NulError};
 
-/// The size of a DICE hash.
-pub const HASH_SIZE: usize = DICE_HASH_SIZE as usize;
-/// The size of the DICE hidden value.
-pub const HIDDEN_SIZE: usize = DICE_HIDDEN_SIZE as usize;
-/// The size of a DICE inline config.
-pub const INLINE_CONFIG_SIZE: usize = DICE_INLINE_CONFIG_SIZE as usize;
 /// The size of a private key seed.
 pub const PRIVATE_KEY_SEED_SIZE: usize = DICE_PRIVATE_KEY_SEED_SIZE as usize;
 /// The size of a CDI.
@@ -122,50 +112,11 @@
     }
 }
 
-/// Configuration descriptor for dice input values.
-#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
-pub enum Config<'a> {
-    /// A reference to an inline descriptor.
-    Inline(&'a [u8; INLINE_CONFIG_SIZE]),
-    /// A reference to a free form descriptor that will be hashed by the implementation.
-    Descriptor(&'a [u8]),
-}
-
 enum ConfigOwned {
-    Inline([u8; INLINE_CONFIG_SIZE]),
+    Inline(InlineConfig),
     Descriptor(Vec<u8>),
 }
 
-impl Config<'_> {
-    fn get_type(&self) -> DiceConfigType {
-        match self {
-            Self::Inline(_) => DICE_CONFIG_TYPE_INLINE,
-            Self::Descriptor(_) => DICE_CONFIG_TYPE_DESCRIPTOR,
-        }
-    }
-
-    fn get_inline(&self) -> [u8; INLINE_CONFIG_SIZE] {
-        match self {
-            Self::Inline(inline) => **inline,
-            _ => [0u8; INLINE_CONFIG_SIZE],
-        }
-    }
-
-    fn get_descriptor_as_ptr(&self) -> *const u8 {
-        match self {
-            Self::Descriptor(descriptor) => descriptor.as_ptr(),
-            _ => std::ptr::null(),
-        }
-    }
-
-    fn get_descriptor_size(&self) -> usize {
-        match self {
-            Self::Descriptor(descriptor) => descriptor.len(),
-            _ => 0,
-        }
-    }
-}
-
 impl From<Config<'_>> for ConfigOwned {
     fn from(config: Config) -> Self {
         match config {
@@ -175,66 +126,41 @@
     }
 }
 
-/// DICE modes as defined here:
-/// https://pigweed.googlesource.com/open-dice/+/refs/heads/main/docs/specification.md#mode-value-details
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-pub enum Mode {
-    /// See documentation linked above.
-    NotConfigured = 0,
-    /// See documentation linked above.
-    Normal = 1,
-    /// See documentation linked above.
-    Debug = 2,
-    /// See documentation linked above.
-    Recovery = 3,
-}
-
-impl Mode {
-    fn get_internal(&self) -> DiceMode {
-        match self {
-            Self::NotConfigured => DICE_MODE_NOT_CONFIGURED,
-            Self::Normal => DICE_MODE_NORMAL,
-            Self::Debug => DICE_MODE_DEBUG,
-            Self::Recovery => DICE_MODE_RECOVERY,
-        }
-    }
-}
-
 /// This trait allows API users to supply DICE input values without copying.
 pub trait InputValues {
     /// Returns the code hash.
-    fn code_hash(&self) -> &[u8; HASH_SIZE];
+    fn code_hash(&self) -> &Hash;
     /// Returns the config.
     fn config(&self) -> Config;
     /// Returns the authority hash.
-    fn authority_hash(&self) -> &[u8; HASH_SIZE];
+    fn authority_hash(&self) -> &Hash;
     /// Returns the authority descriptor.
     fn authority_descriptor(&self) -> Option<&[u8]>;
     /// Returns the mode.
-    fn mode(&self) -> Mode;
+    fn mode(&self) -> DiceMode;
     /// Returns the hidden value.
-    fn hidden(&self) -> &[u8; HIDDEN_SIZE];
+    fn hidden(&self) -> &Hidden;
 }
 
 /// An owning convenience type implementing `InputValues`.
 pub struct InputValuesOwned {
-    code_hash: [u8; HASH_SIZE],
+    code_hash: Hash,
     config: ConfigOwned,
-    authority_hash: [u8; HASH_SIZE],
+    authority_hash: Hash,
     authority_descriptor: Option<Vec<u8>>,
-    mode: Mode,
-    hidden: [u8; HIDDEN_SIZE],
+    mode: DiceMode,
+    hidden: Hidden,
 }
 
 impl InputValuesOwned {
     /// Construct a new instance of InputValuesOwned.
     pub fn new(
-        code_hash: [u8; HASH_SIZE],
+        code_hash: Hash,
         config: Config,
-        authority_hash: [u8; HASH_SIZE],
+        authority_hash: Hash,
         authority_descriptor: Option<Vec<u8>>,
-        mode: Mode,
-        hidden: [u8; HIDDEN_SIZE],
+        mode: DiceMode,
+        hidden: Hidden,
     ) -> Self {
         Self {
             code_hash,
@@ -248,7 +174,7 @@
 }
 
 impl InputValues for InputValuesOwned {
-    fn code_hash(&self) -> &[u8; HASH_SIZE] {
+    fn code_hash(&self) -> &Hash {
         &self.code_hash
     }
     fn config(&self) -> Config {
@@ -257,16 +183,16 @@
             ConfigOwned::Descriptor(descriptor) => Config::Descriptor(descriptor.as_slice()),
         }
     }
-    fn authority_hash(&self) -> &[u8; HASH_SIZE] {
+    fn authority_hash(&self) -> &Hash {
         &self.authority_hash
     }
     fn authority_descriptor(&self) -> Option<&[u8]> {
         self.authority_descriptor.as_deref()
     }
-    fn mode(&self) -> Mode {
+    fn mode(&self) -> DiceMode {
         self.mode
     }
-    fn hidden(&self) -> &[u8; HIDDEN_SIZE] {
+    fn hidden(&self) -> &Hidden {
         &self.hidden
     }
 }
@@ -275,24 +201,16 @@
 where
     F: FnOnce(*const DiceInputValues) -> Result<R>,
 {
-    let input_values = DiceInputValues {
-        code_hash: *input_values.code_hash(),
-        code_descriptor: std::ptr::null(),
-        code_descriptor_size: 0,
-        config_type: input_values.config().get_type(),
-        config_value: input_values.config().get_inline(),
-        config_descriptor: input_values.config().get_descriptor_as_ptr(),
-        config_descriptor_size: input_values.config().get_descriptor_size(),
-        authority_hash: *input_values.authority_hash(),
-        authority_descriptor: input_values
-            .authority_descriptor()
-            .map_or_else(std::ptr::null, <[u8]>::as_ptr),
-        authority_descriptor_size: input_values.authority_descriptor().map_or(0, <[u8]>::len),
-        mode: input_values.mode().get_internal(),
-        hidden: *input_values.hidden(),
-    };
-
-    f(&input_values as *const DiceInputValues)
+    let input_values = diced_open_dice::InputValues::new(
+        input_values.code_hash(),
+        None, // code_descriptor
+        input_values.config(),
+        input_values.authority_hash(),
+        input_values.authority_descriptor(),
+        input_values.mode(),
+        Some(input_values.hidden()),
+    );
+    f(input_values.as_ptr())
 }
 
 /// Multiple of the open dice function required preallocated output buffer
diff --git a/diced/src/diced_client_test.rs b/diced/src/diced_client_test.rs
index 3915508..f579a4e 100644
--- a/diced/src/diced_client_test.rs
+++ b/diced/src/diced_client_test.rs
@@ -52,11 +52,7 @@
         diced_utils::ResidentArtifacts::new(&former.cdiAttest, &former.cdiSeal, &former.bcc.data)
             .unwrap();
 
-    let input_values: Vec<diced_utils::InputValues> =
-        input_values.iter().map(|v| v.into()).collect();
-
-    let artifacts =
-        artifacts.execute_steps(input_values.iter().map(|v| v as &dyn dice::InputValues)).unwrap();
+    let artifacts = artifacts.execute_steps(input_values.iter()).unwrap();
     let (cdi_attest, cdi_seal, bcc) = artifacts.into_tuple();
     let from_former = diced_utils::make_bcc_handover(
         cdi_attest[..].try_into().unwrap(),
@@ -96,11 +92,7 @@
     )
     .unwrap();
 
-    let input_values: Vec<diced_utils::InputValues> =
-        input_values.iter().map(|v| v.into()).collect();
-
-    let artifacts =
-        artifacts.execute_steps(input_values.iter().map(|v| v as &dyn dice::InputValues)).unwrap();
+    let artifacts = artifacts.execute_steps(input_values.iter()).unwrap();
     let (cdi_attest, cdi_seal, bcc) = artifacts.into_tuple();
     let from_former = diced_utils::make_bcc_handover(
         cdi_attest[..].try_into().unwrap(),
@@ -160,11 +152,8 @@
     .unwrap();
 
     let client = [client];
-    let input_values: Vec<diced_utils::InputValues> =
-        input_values.iter().chain(client.iter()).map(|v| v.into()).collect();
 
-    let artifacts =
-        artifacts.execute_steps(input_values.iter().map(|v| v as &dyn dice::InputValues)).unwrap();
+    let artifacts = artifacts.execute_steps(input_values.iter().chain(client.iter())).unwrap();
     let (cdi_attest, cdi_seal, bcc) = artifacts.into_tuple();
     let from_former = diced_utils::make_bcc_handover(
         cdi_attest[..].try_into().unwrap(),
diff --git a/diced/src/hal_node.rs b/diced/src/hal_node.rs
index 69cf4ac..d2e8b0c 100644
--- a/diced/src/hal_node.rs
+++ b/diced/src/hal_node.rs
@@ -187,10 +187,8 @@
         // `ResidentHal::new`
         run_forked(move || {
             let artifacts = artifacts.with_artifacts(|a| ResidentArtifacts::new_from(a))?;
-            let input_values: Vec<utils::InputValues> =
-                input_values.iter().map(|v| v.into()).collect();
             let artifacts = artifacts
-                .execute_steps(input_values.iter().map(|v| v as &dyn dice::InputValues))
+                .execute_steps(input_values)
                 .context("In ResidentHal::get_effective_artifacts:")?;
             f(artifacts)
         })
@@ -279,11 +277,8 @@
         *artifacts = run_forked(|| {
             let new_artifacts =
                 artifacts_clone.with_artifacts(|a| ResidentArtifacts::new_from(a))?;
-            let input_values: Vec<utils::InputValues> =
-                input_values.iter().map(|v| v.into()).collect();
-
             let new_artifacts = new_artifacts
-                .execute_steps(input_values.iter().map(|v| v as &dyn dice::InputValues))
+                .execute_steps(input_values)
                 .context("In ResidentHal::get_effective_artifacts:")?;
             artifacts_clone.update(&new_artifacts)
         })?;
@@ -414,11 +409,7 @@
             make_input_values("component 2 code", "component 2", "component 2 authority")?,
             make_input_values("component 3 code", "component 3", "component 3 authority")?,
         ];
-
-        let input_values: Vec<utils::InputValues> = input_values.iter().map(|v| v.into()).collect();
-
-        let new_artifacts =
-            artifacts.execute_steps(input_values.iter().map(|v| v as &dyn dice::InputValues))?;
+        let new_artifacts = artifacts.execute_steps(input_values)?;
 
         let result = utils::make_bcc_handover(
             new_artifacts.cdi_attest(),
diff --git a/diced/src/resident_node.rs b/diced/src/resident_node.rs
index 99a6dc9..0bd5d0d 100644
--- a/diced/src/resident_node.rs
+++ b/diced/src/resident_node.rs
@@ -24,7 +24,7 @@
 use anyhow::{Context, Result};
 use dice::{ContextImpl, OpenDiceCborContext};
 use diced_open_dice_cbor as dice;
-use diced_utils::{self as utils, InputValues, ResidentArtifacts};
+use diced_utils::{self as utils, ResidentArtifacts};
 use std::collections::HashMap;
 use std::convert::TryInto;
 use std::sync::RwLock;
@@ -61,17 +61,12 @@
 
         let client_arr = [client];
 
-        let input_values: Vec<utils::InputValues> = demotion_db
+        let input_values = demotion_db
             .get(&client_arr[0])
             .map(|v| v.iter())
             .unwrap_or_else(|| client_arr.iter())
-            .chain(input_values.iter())
-            .map(|v| v.into())
-            .collect();
-
-        artifacts
-            .execute_steps(input_values.iter().map(|v| v as &dyn dice::InputValues))
-            .context("In get_effective_artifacts:")
+            .chain(input_values.iter());
+        artifacts.execute_steps(input_values).context("In get_effective_artifacts:")
     }
 }
 
@@ -173,18 +168,10 @@
     fn demote_self(&self, input_values: &[BinderInputValues]) -> Result<()> {
         let mut artifacts = self.artifacts.write().unwrap();
 
-        let input_values = input_values
-            .iter()
-            .map(|v| {
-                v.try_into().with_context(|| format!("Failed to convert input values: {:#?}", v))
-            })
-            .collect::<Result<Vec<InputValues>>>()
-            .context("In ResidentNode::demote_self:")?;
-
         *artifacts = artifacts
             .try_clone()
             .context("In ResidentNode::demote_self: Failed to clone resident artifacts")?
-            .execute_steps(input_values.iter().map(|v| v as &dyn dice::InputValues))
+            .execute_steps(input_values)
             .context("In ResidentNode::demote_self:")?;
         Ok(())
     }
diff --git a/diced/src/sample_inputs.rs b/diced/src/sample_inputs.rs
index ff239ed..14e0efa 100644
--- a/diced/src/sample_inputs.rs
+++ b/diced/src/sample_inputs.rs
@@ -24,7 +24,7 @@
 use diced_utils::cbor;
 use diced_utils::InputValues;
 use keystore2_crypto::ZVec;
-use std::convert::{TryFrom, TryInto};
+use std::convert::TryInto;
 use std::io::Write;
 
 /// Sample UDS used to perform the root dice flow by `make_sample_bcc_and_cdis`.
@@ -82,12 +82,7 @@
     let input_values_vector = get_input_values_vector();
 
     let (cdi_attest, cdi_seal, mut cert) = dice_ctx
-        .main_flow(
-            UDS,
-            UDS,
-            &InputValues::try_from(&input_values_vector[0])
-                .context("In make_sample_bcc_and_cdis: Trying to convert input values. (0)")?,
-        )
+        .main_flow(UDS, UDS, &InputValues::from(&input_values_vector[0]))
         .context("In make_sample_bcc_and_cdis: Trying to run first main flow.")?;
 
     let mut bcc: Vec<u8> = vec![];
@@ -108,8 +103,7 @@
                 "In make_sample_bcc_and_cdis: Failed to convert cdi_seal to array reference. (1)",
             )?,
             &bcc,
-            &InputValues::try_from(&input_values_vector[1])
-                .context("In make_sample_bcc_and_cdis: Trying to convert input values. (1)")?,
+            &InputValues::from(&input_values_vector[1]),
         )
         .context("In make_sample_bcc_and_cdis: Trying to run first bcc main flow.")?;
     dice_ctx
@@ -121,20 +115,19 @@
                 "In make_sample_bcc_and_cdis: Failed to convert cdi_seal to array reference. (2)",
             )?,
             &bcc,
-            &InputValues::try_from(&input_values_vector[2])
-                .context("In make_sample_bcc_and_cdis: Trying to convert input values. (2)")?,
+            &InputValues::from(&input_values_vector[2]),
         )
         .context("In make_sample_bcc_and_cdis: Trying to run second bcc main flow.")
 }
 
 fn make_input_values(
-    code_hash: &[u8; dice::HASH_SIZE],
-    authority_hash: &[u8; dice::HASH_SIZE],
+    code_hash: &dice::Hash,
+    authority_hash: &dice::Hash,
     config_name: &str,
     config_version: u64,
     config_resettable: bool,
     mode: Mode,
-    hidden: &[u8; dice::HIDDEN_SIZE],
+    hidden: &dice::Hidden,
 ) -> Result<BinderInputValues> {
     Ok(BinderInputValues {
         codeHash: *code_hash,
diff --git a/diced/src/utils.rs b/diced/src/utils.rs
index 37f78ac..05702ad 100644
--- a/diced/src/utils.rs
+++ b/diced/src/utils.rs
@@ -19,7 +19,7 @@
     Mode::Mode as BinderMode,
 };
 use anyhow::{Context, Result};
-use dice::ContextImpl;
+use dice::{ContextImpl, DiceMode, Hash, Hidden};
 use diced_open_dice_cbor as dice;
 use keystore2_crypto::ZVec;
 use std::convert::TryInto;
@@ -47,7 +47,7 @@
 }
 
 impl dice::InputValues for InputValues<'_> {
-    fn code_hash(&self) -> &[u8; dice::HASH_SIZE] {
+    fn code_hash(&self) -> &Hash {
         &self.0.codeHash
     }
 
@@ -55,7 +55,7 @@
         dice::Config::Descriptor(self.0.config.desc.as_slice())
     }
 
-    fn authority_hash(&self) -> &[u8; dice::HASH_SIZE] {
+    fn authority_hash(&self) -> &Hash {
         &self.0.authorityHash
     }
 
@@ -63,17 +63,17 @@
         self.0.authorityDescriptor.as_deref()
     }
 
-    fn mode(&self) -> dice::Mode {
+    fn mode(&self) -> DiceMode {
         match self.0.mode {
-            BinderMode::NOT_INITIALIZED => dice::Mode::NotConfigured,
-            BinderMode::NORMAL => dice::Mode::Normal,
-            BinderMode::DEBUG => dice::Mode::Debug,
-            BinderMode::RECOVERY => dice::Mode::Recovery,
-            _ => dice::Mode::NotConfigured,
+            BinderMode::NOT_INITIALIZED => DiceMode::kDiceModeNotInitialized,
+            BinderMode::NORMAL => DiceMode::kDiceModeNormal,
+            BinderMode::DEBUG => DiceMode::kDiceModeDebug,
+            BinderMode::RECOVERY => DiceMode::kDiceModeMaintenance,
+            _ => DiceMode::kDiceModeNotInitialized,
         }
     }
 
-    fn hidden(&self) -> &[u8; dice::HIDDEN_SIZE] {
+    fn hidden(&self) -> &Hidden {
         // If `self` was created using try_from the length was checked and this cannot panic.
         &self.0.hidden
     }
@@ -153,7 +153,7 @@
         (cdi_attest, cdi_seal, bcc)
     }
 
-    fn execute_step(self, input_values: &dyn dice::InputValues) -> Result<Self> {
+    fn execute_step(self, input_values: InputValues) -> Result<Self> {
         let ResidentArtifacts { cdi_attest, cdi_seal, bcc } = self;
 
         let (cdi_attest, cdi_seal, bcc) = dice::OpenDiceCborContext::new()
@@ -165,7 +165,7 @@
                     format!("Trying to convert cdi_seal. (length: {})", cdi_seal.len())
                 })?,
                 &bcc,
-                input_values,
+                &input_values,
             )
             .context("In ResidentArtifacts::execute_step:")?;
         Ok(ResidentArtifacts { cdi_attest, cdi_seal, bcc })
@@ -173,13 +173,14 @@
 
     /// Iterate through the iterator of dice input values performing one
     /// BCC main flow step on each element.
-    pub fn execute_steps<'a, Iter>(self, input_values: Iter) -> Result<Self>
+    pub fn execute_steps<'a, T, I>(self, input_values: I) -> Result<Self>
     where
-        Iter: IntoIterator<Item = &'a dyn dice::InputValues>,
+        T: Into<InputValues<'a>>,
+        I: IntoIterator<Item = T>,
     {
         input_values
             .into_iter()
-            .try_fold(self, |acc, input_values| acc.execute_step(input_values))
+            .try_fold(self, |acc, input| acc.execute_step(input.into()))
             .context("In ResidentArtifacts::execute_step:")
     }
 }
diff --git a/keystore-engine/keystore2_engine.cpp b/keystore-engine/keystore2_engine.cpp
index 69caf51..dc90be3 100644
--- a/keystore-engine/keystore2_engine.cpp
+++ b/keystore-engine/keystore2_engine.cpp
@@ -146,11 +146,13 @@
         return nullptr;
     }
 
-    rsa->n = BN_dup(public_rsa->n);
-    rsa->e = BN_dup(public_rsa->e);
-    if (rsa->n == nullptr || rsa->e == nullptr) {
+    bssl::UniquePtr<BIGNUM> n(BN_dup(RSA_get0_n(public_rsa)));
+    bssl::UniquePtr<BIGNUM> e(BN_dup(RSA_get0_e(public_rsa)));
+    if (n == nullptr || e == nullptr || !RSA_set0_key(rsa.get(), n.get(), e.get(), nullptr)) {
         return nullptr;
     }
+    OWNERSHIP_TRANSFERRED(n);
+    OWNERSHIP_TRANSFERRED(e);
 
     bssl::UniquePtr<EVP_PKEY> result(EVP_PKEY_new());
     if (result.get() == nullptr || !EVP_PKEY_assign_RSA(result.get(), rsa.get())) {
@@ -420,19 +422,19 @@
         Keystore2KeyBackend{response.metadata.key, response.iSecurityLevel});
 
     bssl::UniquePtr<EVP_PKEY> result;
-    switch (EVP_PKEY_type(pkey->type)) {
+    switch (EVP_PKEY_id(pkey.get())) {
     case EVP_PKEY_RSA: {
-        bssl::UniquePtr<RSA> public_rsa(EVP_PKEY_get1_RSA(pkey.get()));
-        result = wrap_rsa(key_backend, public_rsa.get());
+        RSA* public_rsa = EVP_PKEY_get0_RSA(pkey.get());
+        result = wrap_rsa(key_backend, public_rsa);
         break;
     }
     case EVP_PKEY_EC: {
-        bssl::UniquePtr<EC_KEY> public_ecdsa(EVP_PKEY_get1_EC_KEY(pkey.get()));
-        result = wrap_ecdsa(key_backend, public_ecdsa.get());
+        EC_KEY* public_ecdsa = EVP_PKEY_get0_EC_KEY(pkey.get());
+        result = wrap_ecdsa(key_backend, public_ecdsa);
         break;
     }
     default:
-        LOG(ERROR) << AT << "Unsupported key type " << EVP_PKEY_type(pkey->type);
+        LOG(ERROR) << AT << "Unsupported key type " << EVP_PKEY_id(pkey.get());
         return nullptr;
     }
 
diff --git a/keystore2/src/rkpd_client.rs b/keystore2/src/rkpd_client.rs
index be1f4fb..d0d036c 100644
--- a/keystore2/src/rkpd_client.rs
+++ b/keystore2/src/rkpd_client.rs
@@ -590,6 +590,9 @@
 
         let new_key =
             get_rkpd_attestation_key(&SecurityLevel::TRUSTED_ENVIRONMENT, key_id).unwrap();
+
+        // Restore original key so that we don't leave RKPD with invalid blobs.
+        assert!(store_rkpd_attestation_key(&sec_level, &new_blob, &key.keyBlob).is_ok());
         assert_eq!(new_key.keyBlob, new_blob);
     }