[sw/rom] Enable sec_mmio in keymgr driver

This commit enables sec_mmio functionality in the keymgr driver. The
following values are now tracked by sec_mmio:

* KEYMGR_WORKING_STATE_REG_OFFSET
* KEYMGR_RESEED_INTERVAL_SHADOWED_REG_OFFSET
* KEYMGR_SEALING_SW_BINDING_0_REG_OFFSET[0:7]
* KEYMGR_ATTEST_SW_BINDING_0_REG_OFFSET[0:7]
* KEYMGR_SW_BINDING_REGWEN_REG_OFFSET
* KEYMGR_MAX_CREATOR_KEY_VER_REGWEN_REG_OFFSET
* KEYMGR_MAX_CREATOR_KEY_VER_SHADOWED_REG_OFFSET
* KEYMGR_MAX_OWNER_INT_KEY_VER_REGWEN_REG_OFFSET
* KEYMGR_MAX_OWNER_INT_KEY_VER_SHADOWED_REG_OFFSET

The max key version register configuration was moved into a separate
function to be able to lock down the creator versus owner intermediate
key intermediate values.

The keymgr functional test was updated to enable sec_mmio expectation
checks during the simulated boot sequence.

Signed-off-by: Miguel Osorio <miguelosorio@google.com>
diff --git a/sw/device/silicon_creator/lib/drivers/keymgr.c b/sw/device/silicon_creator/lib/drivers/keymgr.c
index 954d180..44723e5 100644
--- a/sw/device/silicon_creator/lib/drivers/keymgr.c
+++ b/sw/device/silicon_creator/lib/drivers/keymgr.c
@@ -7,6 +7,7 @@
 #include "sw/device/lib/base/freestanding/assert.h"
 #include "sw/device/lib/base/memory.h"
 #include "sw/device/silicon_creator/lib/base/abs_mmio.h"
+#include "sw/device/silicon_creator/lib/base/sec_mmio.h"
 
 #include "hw/top_earlgrey/sw/autogen/top_earlgrey.h"
 #include "keymgr_regs.h"  // Generated.
@@ -34,7 +35,7 @@
  * @return `kErrorOk` if the key manager is at the `expected_state` and the
  * status is idle or success.
  */
-static rom_error_t check_expected_state(uint32_t expected_state) {
+static rom_error_t expected_state_check(uint32_t expected_state) {
   // Read and clear the status register by writing back the read value,
   // polling until the status is non-WIP.
   uint32_t op_status;
@@ -51,7 +52,9 @@
   uint32_t error_code = abs_mmio_read32(kBase + KEYMGR_ERR_CODE_REG_OFFSET);
   abs_mmio_write32(kBase + KEYMGR_ERR_CODE_REG_OFFSET, error_code);
 
-  uint32_t got_state = abs_mmio_read32(kBase + KEYMGR_WORKING_STATE_REG_OFFSET);
+  // Read the working state with sec_mmio so that we can check the expected
+  // value periodically.
+  uint32_t got_state = sec_mmio_read32(kBase + KEYMGR_WORKING_STATE_REG_OFFSET);
   if (op_status_field == KEYMGR_OP_STATUS_STATUS_VALUE_IDLE &&
       error_code == 0u && got_state == expected_state) {
     return kErrorOk;
@@ -60,37 +63,54 @@
 }
 
 rom_error_t keymgr_init(uint16_t entropy_reseed_interval) {
-  RETURN_IF_ERROR(check_expected_state(kKeymgrStateReset));
+  RETURN_IF_ERROR(expected_state_check(kKeymgrStateReset));
   uint32_t reg = bitfield_field32_write(
       0, KEYMGR_RESEED_INTERVAL_SHADOWED_VAL_FIELD, entropy_reseed_interval);
-  abs_mmio_write32_shadowed(kBase + KEYMGR_RESEED_INTERVAL_SHADOWED_REG_OFFSET,
+  sec_mmio_write32_shadowed(kBase + KEYMGR_RESEED_INTERVAL_SHADOWED_REG_OFFSET,
                             reg);
+  sec_mmio_write_increment(/*value=*/1);
   return kErrorOk;
 }
 
-void keymgr_set_next_stage_inputs(
+void keymgr_sw_binding_set(
     const keymgr_binding_value_t *binding_value_sealing,
-    const keymgr_binding_value_t *binding_value_attestation,
-    uint32_t max_key_ver) {
+    const keymgr_binding_value_t *binding_value_attestation) {
   // Write and lock (rw0c) the software binding value. This register is unlocked
   // by hardware upon a successful state transition.
-  // FIXME: Consider using sec_mmio module for the following register writes.
   for (size_t i = 0; i < ARRAYSIZE(binding_value_sealing->data); ++i) {
-    abs_mmio_write32(
+    sec_mmio_write32(
         kBase + KEYMGR_SEALING_SW_BINDING_0_REG_OFFSET + i * sizeof(uint32_t),
         binding_value_sealing->data[i]);
   }
   for (size_t i = 0; i < ARRAYSIZE(binding_value_attestation->data); ++i) {
-    abs_mmio_write32(
+    sec_mmio_write32(
         kBase + KEYMGR_ATTEST_SW_BINDING_0_REG_OFFSET + i * sizeof(uint32_t),
         binding_value_attestation->data[i]);
   }
-  abs_mmio_write32(kBase + KEYMGR_SW_BINDING_REGWEN_REG_OFFSET, 0);
+  sec_mmio_write32(kBase + KEYMGR_SW_BINDING_REGWEN_REG_OFFSET, 0);
+  sec_mmio_write_increment(/*value=*/17);
+}
 
+void keymgr_sw_binding_unlock_wait(void) {
+  while (!abs_mmio_read32(kBase + KEYMGR_SW_BINDING_REGWEN_REG_OFFSET)) {
+  }
+  sec_mmio_read32(kBase + KEYMGR_SW_BINDING_REGWEN_REG_OFFSET);
+}
+
+void keymgr_creator_max_ver_set(uint32_t max_key_ver) {
   // Write and lock (rw0c) the max key version.
-  abs_mmio_write32_shadowed(
+  sec_mmio_write32_shadowed(
       kBase + KEYMGR_MAX_CREATOR_KEY_VER_SHADOWED_REG_OFFSET, max_key_ver);
-  abs_mmio_write32(kBase + KEYMGR_MAX_CREATOR_KEY_VER_REGWEN_REG_OFFSET, 0);
+  sec_mmio_write32(kBase + KEYMGR_MAX_CREATOR_KEY_VER_REGWEN_REG_OFFSET, 0);
+  sec_mmio_write_increment(/*value=*/2);
+}
+
+void keymgr_owner_int_max_ver_set(uint32_t max_key_ver) {
+  // Write and lock (rw0c) the max key version.
+  sec_mmio_write32_shadowed(
+      kBase + KEYMGR_MAX_OWNER_INT_KEY_VER_SHADOWED_REG_OFFSET, max_key_ver);
+  sec_mmio_write32(kBase + KEYMGR_MAX_OWNER_INT_KEY_VER_REGWEN_REG_OFFSET, 0);
+  sec_mmio_write_increment(/*value=*/2);
 }
 
 void keymgr_advance_state(void) {
@@ -102,6 +122,6 @@
   abs_mmio_write32(kBase + KEYMGR_CONTROL_REG_OFFSET, reg);
 }
 
-rom_error_t keymgr_check_state(keymgr_state_t expected_state) {
-  return check_expected_state(expected_state);
+rom_error_t keymgr_state_check(keymgr_state_t expected_state) {
+  return expected_state_check(expected_state);
 }
diff --git a/sw/device/silicon_creator/lib/drivers/keymgr.h b/sw/device/silicon_creator/lib/drivers/keymgr.h
index 9ed6413..923c12e 100644
--- a/sw/device/silicon_creator/lib/drivers/keymgr.h
+++ b/sw/device/silicon_creator/lib/drivers/keymgr.h
@@ -58,17 +58,38 @@
 } keymgr_state_t;
 
 /**
- * Sets the key manager inputs.
+ * Sets the key manager software binding inputs.
  *
- * @param binding_value Software binding value from the manifest of the next
- * stage.
- * @param max_key_version Maximum key version from the manifest of the next
- * stage.
+ * @param binding_value_sealing Software binding for sealing value.
+ * @param binding_value_attestation Software binding for attestation value.
  */
-void keymgr_set_next_stage_inputs(
+void keymgr_sw_binding_set(
     const keymgr_binding_value_t *binding_value_sealing,
-    const keymgr_binding_value_t *binding_value_attestation,
-    uint32_t max_key_version);
+    const keymgr_binding_value_t *binding_value_attestation);
+
+/**
+ * Blocks until the software binding registers are unlocked.
+ *
+ * This function can be called after `keymgr_advance_state()` to wait for the
+ * software binding registers to become available for writing.
+ */
+void keymgr_sw_binding_unlock_wait(void);
+
+/**
+ * Sets the Silicon Creator max key version.
+ *
+ * @param max_key_ver Maximum key version associated with the Silicon Creator
+ * key manager stage.
+ */
+void keymgr_creator_max_ver_set(uint32_t max_key_ver);
+
+/**
+ * Sets the Silicon Owner Intermediate max key version.
+ *
+ * @param max_key_ver Maximum key version associated with the Silicon Onwer
+ * Intermediate key manager stage.
+ */
+void keymgr_owner_int_max_ver_set(uint32_t max_key_ver);
 
 /**
  * Initializes the key manager.
@@ -88,12 +109,16 @@
 /**
  * Advances the state of the key manager.
  *
- * The `keymgr_check_state()` function must be called before this function to
+ * The `keymgr_state_check()` function must be called before this function to
  * ensure the key manager is in the expected state and ready to receive op
  * commands.
  *
- * The caller is responsible for calling the `keymgr_check_state()` at a later
+ * The caller is responsible for calling the `keymgr_state_check()` at a later
  * time to ensure the advance transition completed without errors.
+ *
+ * Note: It is recommended to call `keymgr_sw_binding_unlock_wait()` before the
+ * secure mmio `sec_mmio_check_values()` function to make sure the internal
+ * state of the key manager is updated in the secure mmio expectations table.
  */
 void keymgr_advance_state(void);
 
@@ -104,7 +129,7 @@
  * @return `kErrorOk` if the key manager is in `expected_state` and the status
  * is idle or success; otherwise returns `kErrorKeymgrInternal`.
  */
-rom_error_t keymgr_check_state(keymgr_state_t expected_state);
+rom_error_t keymgr_state_check(keymgr_state_t expected_state);
 
 #ifdef __cplusplus
 }
diff --git a/sw/device/silicon_creator/lib/drivers/keymgr_functest.c b/sw/device/silicon_creator/lib/drivers/keymgr_functest.c
index bf1624f..eb76839 100644
--- a/sw/device/silicon_creator/lib/drivers/keymgr_functest.c
+++ b/sw/device/silicon_creator/lib/drivers/keymgr_functest.c
@@ -91,6 +91,11 @@
 
 const test_config_t kTestConfig;
 
+/** SEC MMIO error handler. */
+static void error_handler_cb(rom_error_t e) {
+  LOG_ERROR("Secure MMIO error: 0x%x", e);
+}
+
 /**
  * Writes `size` words of `data` into flash info page.
  *
@@ -138,36 +143,6 @@
                   ARRAYSIZE(kOwnerSecret));
 }
 
-/** Key manager configuration steps performed in mask ROM. */
-rom_error_t keymgr_rom_test(void) {
-  ASSERT_OK(keymgr_check_state(kKeymgrStateReset));
-  keymgr_set_next_stage_inputs(&kBindingValueRomExt, &kBindingValueRomExt,
-                               kMaxVerRomExt);
-  return kErrorOk;
-}
-
-/** Key manager configuration steps performed in ROM_EXT. */
-rom_error_t keymgr_rom_ext_test(void) {
-  ASSERT_OK(keymgr_check_state(kKeymgrStateReset));
-
-  const uint16_t kEntropyReseedInterval = 0x1234;
-  ASSERT_OK(keymgr_init(kEntropyReseedInterval));
-  keymgr_advance_state();
-  ASSERT_OK(keymgr_check_state(kKeymgrStateInit));
-
-  // FIXME: Check `kBindingValueRomExt` before advancing state.
-  keymgr_advance_state();
-  ASSERT_OK(keymgr_check_state(kKeymgrStateCreatorRootKey));
-
-  keymgr_set_next_stage_inputs(&kBindingValueBl0, &kBindingValueBl0,
-                               kMaxVerBl0);
-
-  // FIXME: Check `kBindingValueBl0` before advancing state.
-  keymgr_advance_state();
-  ASSERT_OK(keymgr_check_state(kKeymgrStateOwnerIntermediateKey));
-  return kErrorOk;
-}
-
 static const dif_pwrmgr_wakeup_reason_t kWakeUpReasonPor = {
     .types = 0,
     .request_sources = 0,
@@ -269,6 +244,47 @@
   wait_for_interrupt();
 }
 
+/** Key manager configuration steps performed in mask ROM. */
+rom_error_t keymgr_rom_test(void) {
+  ASSERT_OK(keymgr_state_check(kKeymgrStateReset));
+  keymgr_sw_binding_set(&kBindingValueRomExt, &kBindingValueRomExt);
+  keymgr_creator_max_ver_set(kMaxVerRomExt);
+  sec_mmio_check_values(/*rnd_offset=*/0);
+  sec_mmio_check_counters(/*expected_check_count=*/1);
+  return kErrorOk;
+}
+
+/** Key manager configuration steps performed in ROM_EXT. */
+rom_error_t keymgr_rom_ext_test(void) {
+  ASSERT_OK(keymgr_state_check(kKeymgrStateReset));
+
+  const uint16_t kEntropyReseedInterval = 0x1234;
+  ASSERT_OK(keymgr_init(kEntropyReseedInterval));
+
+  sec_mmio_check_values(/*rnd_offset=*/0);
+  keymgr_advance_state();
+  ASSERT_OK(keymgr_state_check(kKeymgrStateInit));
+
+  keymgr_advance_state();
+  ASSERT_OK(keymgr_state_check(kKeymgrStateCreatorRootKey));
+
+  // The software binding register lock is reset after advancing the key
+  // manager, so we need to call this function to update sec_mmio expectation
+  // table.
+  keymgr_sw_binding_unlock_wait();
+  sec_mmio_check_values(/*rnd_offset=*/0);
+
+  keymgr_sw_binding_set(&kBindingValueBl0, &kBindingValueBl0);
+  keymgr_owner_int_max_ver_set(kMaxVerBl0);
+  sec_mmio_check_values(/*rnd_offset=*/0);
+
+  keymgr_advance_state();
+  ASSERT_OK(keymgr_state_check(kKeymgrStateOwnerIntermediateKey));
+
+  sec_mmio_check_counters(/*expected_check_count=*/5);
+  return kErrorOk;
+}
+
 bool test_main(void) {
   rom_error_t result = kErrorOk;
 
@@ -311,7 +327,7 @@
   } else {
     LOG_INFO("Powered up for the second time, actuate keymgr");
 
-    // Initialize kmac for key manager
+    sec_mmio_init(error_handler_cb);
     init_kmac_for_keymgr();
 
     EXECUTE_TEST(result, keymgr_rom_test);
diff --git a/sw/device/silicon_creator/lib/drivers/keymgr_unittest.cc b/sw/device/silicon_creator/lib/drivers/keymgr_unittest.cc
index 4eb2c7d..dfe19ae 100644
--- a/sw/device/silicon_creator/lib/drivers/keymgr_unittest.cc
+++ b/sw/device/silicon_creator/lib/drivers/keymgr_unittest.cc
@@ -9,6 +9,7 @@
 
 #include "gtest/gtest.h"
 #include "sw/device/silicon_creator/lib/base/mock_abs_mmio.h"
+#include "sw/device/silicon_creator/lib/base/mock_sec_mmio.h"
 #include "sw/device/silicon_creator/lib/error.h"
 
 #include "hw/top_earlgrey/sw/autogen/top_earlgrey.h"
@@ -33,7 +34,7 @@
     EXPECT_ABS_READ32(base_ + KEYMGR_ERR_CODE_REG_OFFSET, err_code);
     EXPECT_ABS_WRITE32(base_ + KEYMGR_ERR_CODE_REG_OFFSET, err_code);
 
-    EXPECT_ABS_READ32(base_ + KEYMGR_WORKING_STATE_REG_OFFSET, km_state);
+    EXPECT_SEC_READ32(base_ + KEYMGR_WORKING_STATE_REG_OFFSET, km_state);
   }
   uint32_t base_ = TOP_EARLGREY_KEYMGR_BASE_ADDR;
   SwBindingCfg cfg_ = {
@@ -42,6 +43,7 @@
       .binding_value_attestation = {9, 10, 11, 12, 13, 14, 15},
   };
   mask_rom_test::MockAbsMmio mmio_;
+  mask_rom_test::MockSecMmio sec_mmio_;
 };
 
 TEST_F(KeymgrTest, Initialize) {
@@ -49,56 +51,74 @@
                     KEYMGR_WORKING_STATE_STATE_VALUE_RESET,
                     /*err_code=*/0u);
 
-  EXPECT_ABS_WRITE32_SHADOWED(
+  EXPECT_SEC_WRITE32_SHADOWED(
       base_ + KEYMGR_RESEED_INTERVAL_SHADOWED_REG_OFFSET, 0u);
-
+  EXPECT_SEC_WRITE_INCREMENT(1);
   EXPECT_EQ(keymgr_init(0u), kErrorOk);
 }
 
-TEST_F(KeymgrTest, SetNextStageInputs) {
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_0_REG_OFFSET,
+TEST_F(KeymgrTest, SwBindingValuesSet) {
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_0_REG_OFFSET,
                      cfg_.binding_value_sealing.data[0]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_1_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_1_REG_OFFSET,
                      cfg_.binding_value_sealing.data[1]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_2_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_2_REG_OFFSET,
                      cfg_.binding_value_sealing.data[2]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_3_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_3_REG_OFFSET,
                      cfg_.binding_value_sealing.data[3]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_4_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_4_REG_OFFSET,
                      cfg_.binding_value_sealing.data[4]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_5_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_5_REG_OFFSET,
                      cfg_.binding_value_sealing.data[5]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_6_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_6_REG_OFFSET,
                      cfg_.binding_value_sealing.data[6]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_7_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_SEALING_SW_BINDING_7_REG_OFFSET,
                      cfg_.binding_value_sealing.data[7]);
 
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_0_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_0_REG_OFFSET,
                      cfg_.binding_value_attestation.data[0]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_1_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_1_REG_OFFSET,
                      cfg_.binding_value_attestation.data[1]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_2_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_2_REG_OFFSET,
                      cfg_.binding_value_attestation.data[2]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_3_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_3_REG_OFFSET,
                      cfg_.binding_value_attestation.data[3]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_4_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_4_REG_OFFSET,
                      cfg_.binding_value_attestation.data[4]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_5_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_5_REG_OFFSET,
                      cfg_.binding_value_attestation.data[5]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_6_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_6_REG_OFFSET,
                      cfg_.binding_value_attestation.data[6]);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_7_REG_OFFSET,
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_ATTEST_SW_BINDING_7_REG_OFFSET,
                      cfg_.binding_value_attestation.data[7]);
 
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_SW_BINDING_REGWEN_REG_OFFSET, 0);
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_SW_BINDING_REGWEN_REG_OFFSET, 0);
+  EXPECT_SEC_WRITE_INCREMENT(17);
+  keymgr_sw_binding_set(&cfg_.binding_value_sealing,
+                        &cfg_.binding_value_attestation);
+}
 
-  EXPECT_ABS_WRITE32_SHADOWED(
+TEST_F(KeymgrTest, SwBindingUnlockWait) {
+  EXPECT_ABS_READ32(base_ + KEYMGR_SW_BINDING_REGWEN_REG_OFFSET, 1);
+  EXPECT_SEC_READ32(base_ + KEYMGR_SW_BINDING_REGWEN_REG_OFFSET, 1);
+  keymgr_sw_binding_unlock_wait();
+}
+
+TEST_F(KeymgrTest, SetCreatorMaxVerKey) {
+  EXPECT_SEC_WRITE32_SHADOWED(
       base_ + KEYMGR_MAX_CREATOR_KEY_VER_SHADOWED_REG_OFFSET, cfg_.max_key_ver);
-  EXPECT_ABS_WRITE32(base_ + KEYMGR_MAX_CREATOR_KEY_VER_REGWEN_REG_OFFSET, 0);
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_MAX_CREATOR_KEY_VER_REGWEN_REG_OFFSET, 0);
+  EXPECT_SEC_WRITE_INCREMENT(2);
+  keymgr_creator_max_ver_set(cfg_.max_key_ver);
+}
 
-  keymgr_set_next_stage_inputs(&cfg_.binding_value_sealing,
-                               &cfg_.binding_value_attestation,
-                               cfg_.max_key_ver);
+TEST_F(KeymgrTest, SetOwnerIntMaxVerKey) {
+  EXPECT_SEC_WRITE32_SHADOWED(
+      base_ + KEYMGR_MAX_OWNER_INT_KEY_VER_SHADOWED_REG_OFFSET,
+      cfg_.max_key_ver);
+  EXPECT_SEC_WRITE32(base_ + KEYMGR_MAX_OWNER_INT_KEY_VER_REGWEN_REG_OFFSET, 0);
+  EXPECT_SEC_WRITE_INCREMENT(2);
+  keymgr_owner_int_max_ver_set(cfg_.max_key_ver);
 }
 
 TEST_F(KeymgrTest, AdvanceState) {
@@ -117,28 +137,28 @@
   ExpectStatusCheck(KEYMGR_OP_STATUS_STATUS_VALUE_IDLE,
                     KEYMGR_WORKING_STATE_STATE_VALUE_CREATOR_ROOT_KEY,
                     /*err_code=*/0u);
-  EXPECT_EQ(keymgr_check_state(kKeymgrStateCreatorRootKey), kErrorOk);
+  EXPECT_EQ(keymgr_state_check(kKeymgrStateCreatorRootKey), kErrorOk);
 }
 
 TEST_F(KeymgrTest, CheckStateInvalidResponse) {
   ExpectStatusCheck(KEYMGR_OP_STATUS_STATUS_VALUE_IDLE,
                     KEYMGR_WORKING_STATE_STATE_VALUE_INVALID,
                     /*err_code=*/0u);
-  EXPECT_EQ(keymgr_check_state(kKeymgrStateCreatorRootKey),
+  EXPECT_EQ(keymgr_state_check(kKeymgrStateCreatorRootKey),
             kErrorKeymgrInternal);
 
   // Any non-idle status is expected to fail.
   ExpectStatusCheck(KEYMGR_OP_STATUS_STATUS_VALUE_DONE_ERROR,
                     KEYMGR_WORKING_STATE_STATE_VALUE_CREATOR_ROOT_KEY,
                     /*err_code=*/0u);
-  EXPECT_EQ(keymgr_check_state(kKeymgrStateCreatorRootKey),
+  EXPECT_EQ(keymgr_state_check(kKeymgrStateCreatorRootKey),
             kErrorKeymgrInternal);
 
   // Any non-zero error code is expected to fail.
   ExpectStatusCheck(KEYMGR_OP_STATUS_STATUS_VALUE_IDLE,
                     KEYMGR_WORKING_STATE_STATE_VALUE_CREATOR_ROOT_KEY,
                     /*err_code=*/1u);
-  EXPECT_EQ(keymgr_check_state(kKeymgrStateCreatorRootKey),
+  EXPECT_EQ(keymgr_state_check(kKeymgrStateCreatorRootKey),
             kErrorKeymgrInternal);
 }
 
diff --git a/sw/device/silicon_creator/lib/drivers/meson.build b/sw/device/silicon_creator/lib/drivers/meson.build
index bd91398..3dd9aec 100644
--- a/sw/device/silicon_creator/lib/drivers/meson.build
+++ b/sw/device/silicon_creator/lib/drivers/meson.build
@@ -74,6 +74,7 @@
     ],
     dependencies: [
       sw_silicon_creator_lib_base_abs_mmio,
+      sw_silicon_creator_lib_base_sec_mmio,
     ],
   ),
 )
@@ -109,6 +110,7 @@
       sw_lib_dif_pwrmgr,
       sw_lib_dif_otp_ctrl,
       sw_lib_flash_ctrl,
+      sw_silicon_creator_lib_base_sec_mmio,
       sw_silicon_creator_lib_driver_keymgr,
       sw_silicon_creator_lib_driver_lifecycle,
     ],
diff --git a/sw/device/silicon_creator/mask_rom/mask_rom.c b/sw/device/silicon_creator/mask_rom/mask_rom.c
index d904e2e..854e966 100644
--- a/sw/device/silicon_creator/mask_rom/mask_rom.c
+++ b/sw/device/silicon_creator/mask_rom/mask_rom.c
@@ -142,10 +142,9 @@
  * @return rom_error_t Result of the operation.
  */
 static rom_error_t mask_rom_boot(const manifest_t *manifest) {
-  RETURN_IF_ERROR(keymgr_check_state(kKeymgrStateReset));
-  keymgr_set_next_stage_inputs(&manifest->binding_value,
-                               &manifest->binding_value,
-                               manifest->max_key_version);
+  RETURN_IF_ERROR(keymgr_state_check(kKeymgrStateReset));
+  keymgr_sw_binding_set(&manifest->binding_value, &manifest->binding_value);
+  keymgr_creator_max_ver_set(manifest->max_key_version);
 
   // Unlock execution of ROM_EXT executable code (text) sections.
   RETURN_IF_ERROR(epmp_state_check(&epmp));