From c86534455577b250f110e3a481a67078ed45d31d Mon Sep 17 00:00:00 2001
From: Vadim Bendebury <vbendeb@google.com>
Date: Tue, 26 Nov 2024 15:06:33 -0800
Subject: [PATCH] [provisioning]: allow the owner to specify perso success
 condition

Certain SKUs require personalization success validation which extends
beyond confirming that the ROM_EXT in slot B has started.

Let's allow the user to specify a string which when issued on the chip
console indicates perso success.

This will require longer timeouts after post perso reboot.

Tested the following:
  - when owner string is not provided
      . success and failure due to UDS cert problems are reported as
        before
  - when owner string is provided
    . failure due to UDS cert problems is still reported
    . success when fw string is detected is reported
    . in case fw string is not detected timeout results in an error

Signed-off-by: Vadim Bendebury <vbendeb@google.com>
---
 sw/host/provisioning/ft/src/main.rs    | 14 +++++++--
 sw/host/provisioning/ft_lib/src/lib.rs | 42 +++++++++++++++++---------
 2 files changed, 40 insertions(+), 16 deletions(-)

diff --git a/sw/host/provisioning/ft/src/main.rs b/sw/host/provisioning/ft/src/main.rs
index a608fa5933602..f33083b77c9d8 100644
--- a/sw/host/provisioning/ft/src/main.rs
+++ b/sw/host/provisioning/ft/src/main.rs
@@ -17,7 +17,7 @@ use p256::NistP256;
 use cert_lib::{CaConfig, CaKey, CaKeyType};
 use ft_lib::response::PersonalizeResponse;
 use ft_lib::{
-    check_rom_ext_boot_up, run_ft_personalize, run_sram_ft_individualize, test_exit, test_unlock,
+    check_slot_b_boot_up, run_ft_personalize, run_sram_ft_individualize, test_exit, test_unlock,
 };
 use opentitanlib::backend;
 use opentitanlib::console::spi::SpiConsoleDevice;
@@ -114,6 +114,10 @@ struct Opts {
     /// Name of the SPI interface to connect to the OTTF console.
     #[arg(long, default_value = "BOOTSTRAP")]
     console_spi: String,
+
+    /// Owner's firmware string indicating successful start up.
+    #[arg(long)]
+    owner_success_text: Option<String>,
 }
 
 fn main() -> Result<()> {
@@ -305,7 +309,13 @@ fn main() -> Result<()> {
         &mut response,
     )?;
 
-    check_rom_ext_boot_up(&transport, &opts.init, opts.timeout, &mut response)?;
+    check_slot_b_boot_up(
+        &transport,
+        &opts.init,
+        opts.timeout,
+        &mut response,
+        opts.owner_success_text,
+    )?;
     log::info!("Provisioning Done");
     let doc = if opts.provisioning_data.pretty {
         serde_json::to_string_pretty(&response)?
diff --git a/sw/host/provisioning/ft_lib/src/lib.rs b/sw/host/provisioning/ft_lib/src/lib.rs
index 65cdad6d9ebb3..a0e38ce20bf8b 100644
--- a/sw/host/provisioning/ft_lib/src/lib.rs
+++ b/sw/host/provisioning/ft_lib/src/lib.rs
@@ -528,11 +528,12 @@ pub fn run_ft_personalize(
     Ok(())
 }
 
-pub fn check_rom_ext_boot_up(
+pub fn check_slot_b_boot_up(
     transport: &TransportWrapper,
     init: &InitializeTest,
     timeout: Duration,
     response: &mut PersonalizeResponse,
+    owner_fw_success_string: Option<String>,
 ) -> Result<()> {
     transport.reset_target(init.bootstrap.options.reset_delay, true)?;
     let uart_console = transport.uart("console")?;
@@ -548,26 +549,39 @@ pub fn check_rom_ext_boot_up(
             .map(|s| s.as_str())
             .unwrap_or("unknown"),
     );
+    let t0 = Instant::now();
 
-    // Timeout for waiting for a potential error message indicating invalid UDS certificate.
-    // This value is tested on fpga cw340 and could be potentially fine-tuned.
-    const UDS_CERT_INVALID_TIMEOUT: Duration = Duration::from_millis(200);
+    // Timeout for waiting for a potential error message indicating invalid UDS
+    // certificate or for the Owner firmware successful startup.
+    //
+    // These values were tested on fpga cw340 and could be potentially fine-tuned.
+    let slot_b_startup_timeout: Duration =
+        Duration::from_millis(if owner_fw_success_string.is_none() {
+            200
+        } else {
+            1500
+        });
 
-    let t0 = Instant::now();
-    let result = UartConsole::wait_for(
-        &*uart_console,
-        r".*UDS certificate not valid.",
-        UDS_CERT_INVALID_TIMEOUT,
-    );
     response.stats.log_elapsed_time("rom_ext-done", t0);
+    let rom_ext_failure_msg = r"Invalid UDS certificate detected!";
+    let anchor_text = if let Some(owner_anchor) = &owner_fw_success_string {
+        format!(r"({}|{})", rom_ext_failure_msg, owner_anchor)
+    } else {
+        rom_ext_failure_msg.to_string()
+    };
+
+    let result =
+        UartConsole::wait_for(&*uart_console, anchor_text.as_str(), slot_b_startup_timeout);
 
     match result {
-        Ok(_captures) => {
-            // Error message found.
-            bail!("Invalid UDS certificate detected!");
+        Ok(captures) => {
+            if captures[0] == *rom_ext_failure_msg {
+                // Error message found.
+                bail!("Invalid UDS certificate detected!");
+            }
         }
         Err(e) => {
-            if e.to_string().contains("Timed Out") {
+            if owner_fw_success_string.is_none() && e.to_string().contains("Timed Out") {
                 // Error message not found after timeout. This is the expected behavior.
             } else {
                 // An unexpected error occurred while waiting for the console output.