[bug/5614743] fix: coderabbit fixes

Author: prbinu-nvidia

crates/scout/src/platform.rs

@@ -17,16 +17,23 @@
 
 use smbioslib::{SMBiosSystemInformation, table_load_from_device};
 
+pub(crate) fn is_host_from_product_names<'a>(
+    product_names: impl IntoIterator<Item = &'a str>,
+) -> bool {
+    !product_names
+        .into_iter()
+        .any(|name| name.to_ascii_lowercase().contains("bluefield"))
+}
+
 /// Returns `true` when scout is running on a managed host (as opposed to a DPU).
 pub(crate) fn is_host() -> bool {
     match table_load_from_device() {
-        Ok(data) => data.any(|sys_info: SMBiosSystemInformation| {
-            !sys_info
-                .product_name()
-                .to_string()
-                .to_lowercase()
-                .contains("bluefield")
-        }),
+        Ok(data) => {
+            let product_names = data
+                .map(|sys_info: SMBiosSystemInformation| sys_info.product_name().to_string())
+                .collect::<Vec<_>>();
+            is_host_from_product_names(product_names.iter().map(String::as_str))
+        }
         Err(_err) => true,
     }
 }
@@ -35,6 +42,26 @@ pub(crate) fn is_host() -> bool {
 mod tests {
     use super::*;
 
+    #[test]
+    fn is_host_from_product_names_cases() {
+        let cases: &[(&[&str], bool)] = &[
+            (&["DGX H100"], true),
+            (&["BlueField-3 DPU"], false),
+            (&["NVIDIA Bluefield 2"], false),
+            (&[""], true),
+            (&["DGX H100", "BlueField-3 DPU"], false),
+            (&["DGX H100", "Other Platform"], true),
+        ];
+
+        for (product_names, want_host) in cases {
+            assert_eq!(
+                is_host_from_product_names(product_names.iter().copied()),
+                *want_host,
+                "product_names={product_names:?}"
+            );
+        }
+    }
+
     #[test]
     fn is_host_returns_bool_without_panicking() {
         let _ = is_host();

crates/scout/src/register.rs

@@ -57,29 +57,21 @@ pub async fn run(
         crate::tpm::set_tpm_max_auth_fail()?;
 
         // create tss context
-        let mut tss_ctx = match attest::create_context_from_path(tpm_path) {
-            Ok(ctx) => ctx,
-            Err(e) => {
-                let err = CarbideClientError::TpmError(format!("Could not create context: {e}"));
-                if tpm::is_recoverable_tpm_client_error(&err) {
-                    tpm::recover_tpm_and_reboot(tpm_path)?;
-                }
-                return Err(err);
-            }
-        };
+        let mut tss_ctx = attest::create_context_from_path(tpm_path)
+            .map_err(|e| CarbideClientError::TpmError(format!("Could not create context: {e}")))?;
 
         // CHANGETO - supply context externally
         hardware_info.tpm_description = attest::get_tpm_description(&mut tss_ctx);
 
         let result = match attest::create_attest_key_info(&mut tss_ctx) {
             Ok(result) => result,
             Err(e) => {
-                let err =
-                    CarbideClientError::TpmError(format!("Could not create AttestKeyInfo: {e}"));
-                if tpm::is_recoverable_tpm_client_error(&err) {
+                if tpm::should_attempt_tpm_recovery_for_attest_key_failure(&*e) {
                     tpm::recover_tpm_and_reboot(tpm_path)?;
                 }
-                return Err(err);
+                return Err(CarbideClientError::TpmError(format!(
+                    "Could not create AttestKeyInfo: {e}"
+                )));
             }
         };
 

crates/scout/src/tpm.rs

@@ -15,8 +15,8 @@
  * limitations under the License.
  */
 
-use std::fs::File;
-use std::io::Write;
+use std::fs::{self, OpenOptions};
+use std::io::{ErrorKind, Write};
 use std::path::Path;
 use std::process::Command;
 
@@ -25,7 +25,7 @@ use tss_esapi::interface_types::session_handles::AuthSession;
 
 use crate::{CarbideClientError, attestation as attest};
 
-pub(crate) const TPM_RECOVERY_ATTEMPTED_PATH: &str = "/tmp/tpm_recovery_reboot_attempted";
+pub(crate) const TPM_RECOVERY_ATTEMPTED_PATH: &str = "/run/scout/tpm_recovery_reboot_attempted";
 
 // From https://superuser.com/questions/1404738/tpm-2-0-hardware-error-da-lockout-mode
 pub(crate) fn set_tpm_max_auth_fail() -> Result<(), CarbideClientError> {
@@ -87,34 +87,47 @@ pub(crate) fn clear_tpm(tpm_path: &str) -> Result<(), CarbideClientError> {
     Ok(())
 }
 
-pub(crate) fn is_recoverable_tpm_client_error(error: &CarbideClientError) -> bool {
-    match error {
-        CarbideClientError::TpmError(message) => {
-            message.contains("Could not create AttestKeyInfo")
-                || message.contains("Could not create context")
-                || message.contains("TPM2_Clear")
-        }
-        _ => false,
+/// Returns true when attestation-key setup failed after a TPM context was opened successfully.
+///
+/// Recovery is only attempted for this stage: context creation failures (bad path, missing device)
+/// are not recoverable via TPM clear.
+pub(crate) fn should_attempt_tpm_recovery_for_attest_key_failure(
+    source: &dyn std::error::Error,
+) -> bool {
+    let message = source.to_string().to_ascii_lowercase();
+    !message.contains("not supported")
+}
+
+fn claim_tpm_recovery_attempt() -> Result<(), CarbideClientError> {
+    if let Some(parent) = Path::new(TPM_RECOVERY_ATTEMPTED_PATH).parent() {
+        fs::create_dir_all(parent).map_err(CarbideClientError::StdIo)?;
     }
+
+    let mut marker = match OpenOptions::new()
+        .write(true)
+        .create_new(true)
+        .open(TPM_RECOVERY_ATTEMPTED_PATH)
+    {
+        Ok(file) => file,
+        Err(e) if e.kind() == ErrorKind::AlreadyExists => {
+            return Err(CarbideClientError::TpmError(
+                "TPM recovery was already attempted this boot cycle; refusing to loop".to_string(),
+            ));
+        }
+        Err(e) => return Err(CarbideClientError::StdIo(e)),
+    };
+    marker
+        .write_all(b"tpm recovery reboot requested\n")
+        .map_err(CarbideClientError::StdIo)
 }
 
 /// Clears the TPM and reboots the host once per boot cycle to recover from missing TPM material.
 pub(crate) fn recover_tpm_and_reboot(tpm_path: &str) -> Result<(), CarbideClientError> {
-    if Path::new(TPM_RECOVERY_ATTEMPTED_PATH).exists() {
-        return Err(CarbideClientError::TpmError(
-            "TPM recovery was already attempted this boot cycle; refusing to loop".to_string(),
-        ));
-    }
+    claim_tpm_recovery_attempt()?;
 
     tracing::warn!("Attempting automated TPM clear and reboot to recover attestation state");
     clear_tpm(tpm_path)?;
 
-    let mut marker =
-        File::create(TPM_RECOVERY_ATTEMPTED_PATH).map_err(CarbideClientError::StdIo)?;
-    marker
-        .write_all(b"tpm recovery reboot requested\n")
-        .map_err(CarbideClientError::StdIo)?;
-
     let output = Command::new("systemctl")
         .arg("reboot")
         .output()
@@ -135,14 +148,20 @@ mod tests {
     use super::*;
 
     #[test]
-    fn recoverable_tpm_errors_include_attest_key_info_failures() {
-        let err = CarbideClientError::TpmError("Could not create AttestKeyInfo: test".to_string());
-        assert!(is_recoverable_tpm_client_error(&err));
-    }
-
-    #[test]
-    fn non_tpm_client_errors_are_not_recoverable() {
-        let err = CarbideClientError::GenericError("transport failed".to_string());
-        assert!(!is_recoverable_tpm_client_error(&err));
+    fn attest_key_failure_recovery_classification_cases() {
+        let cases: &[(&str, bool)] = &[
+            ("handle already exists", true),
+            ("tpm corruption detected", true),
+            ("feature not supported on this device", false),
+        ];
+
+        for (message, want_recovery) in cases {
+            let err: Box<dyn std::error::Error> = Box::new(std::io::Error::other(*message));
+            assert_eq!(
+                should_attempt_tpm_recovery_for_attest_key_failure(&*err),
+                *want_recovery,
+                "message={message:?}"
+            );
+        }
     }
 }