From 409eadddb284b92c1b193b663383d7f75c8fd8a4 Mon Sep 17 00:00:00 2001
From: "James O. D. Hunt" <james.o.hunt@intel.com>
Date: Thu, 12 Oct 2023 10:40:27 +0100
Subject: [PATCH] runtime-rs: ch: Improve readability of guest protection
 checks

Improve the way `handle_guest_protection()` is structured by inverting
the logic and checking the value of the `confidential_guest` setting
before checking the guest protection. This makes the code easier to
understand.

> **Notes:**
>
> - This change also unconditionally saves the available guest protection
>   (where previously it was only saved when `confidential_guest=true`).
>   This explains the minor unit test fix.
>
> - This changes also errors if the CH driver finds an unexpected
>   protection (since only Intel TDX is currently tested).

Signed-off-by: James O. D. Hunt <james.o.hunt@intel.com>
---
 .../hypervisor/src/ch/inner_hypervisor.rs     | 48 ++++++++++++++-----
 1 file changed, 37 insertions(+), 11 deletions(-)

diff --git a/src/runtime-rs/crates/hypervisor/src/ch/inner_hypervisor.rs b/src/runtime-rs/crates/hypervisor/src/ch/inner_hypervisor.rs
index 8e48040087..d16857015b 100644
--- a/src/runtime-rs/crates/hypervisor/src/ch/inner_hypervisor.rs
+++ b/src/runtime-rs/crates/hypervisor/src/ch/inner_hypervisor.rs
@@ -55,6 +55,10 @@ pub enum GuestProtectionError {
     // "true".
     #[error("TDX guest protection available and must be used with Cloud Hypervisor (set 'confidential_guest=true')")]
     TDXProtectionMustBeUsedWithCH,
+
+    // TDX is the only tested CH protection currently.
+    #[error("Expected TDX protection, found {0}")]
+    ExpectedTDXProtection(GuestProtection),
 }
 
 impl CloudHypervisorInner {
@@ -482,17 +486,25 @@ impl CloudHypervisorInner {
             task::spawn_blocking(|| -> Result<GuestProtection> { get_guest_protection() })
                 .await??;
 
-        if protection == GuestProtection::NoProtection {
-            if confidential_guest {
-                return Err(anyhow!(GuestProtectionError::NoProtectionAvailable));
-            } else {
-                debug!(sl!(), "no guest protection available");
-            }
-        } else if confidential_guest {
-            self.guest_protection_to_use = protection.clone();
+        self.guest_protection_to_use = protection.clone();
 
-            info!(sl!(), "guest protection available and requested"; "guest-protection" => protection.to_string());
+        info!(sl!(), "guest protection {:?}", protection.to_string());
+
+        if confidential_guest {
+            if protection == GuestProtection::NoProtection {
+                // User wants protection, but none available.
+                return Err(anyhow!(GuestProtectionError::NoProtectionAvailable));
+            } else if let GuestProtection::Tdx(_) = protection {
+                info!(sl!(), "guest protection available and requested"; "guest-protection" => protection.to_string());
+            } else {
+                return Err(anyhow!(GuestProtectionError::ExpectedTDXProtection(
+                    protection
+                )));
+            }
+        } else if protection == GuestProtection::NoProtection {
+            debug!(sl!(), "no guest protection available");
         } else if let GuestProtection::Tdx(_) = protection {
+            // CH requires TDX protection to be used.
             return Err(anyhow!(GuestProtectionError::TDXProtectionMustBeUsedWithCH));
         } else {
             info!(sl!(), "guest protection available but not requested"; "guest-protection" => protection.to_string());
@@ -899,13 +911,27 @@ mod tests {
                 confidential_guest: false,
                 available_protection: Some(GuestProtection::Tdx(tdx_details.clone())),
                 result: Err(anyhow!(GuestProtectionError::TDXProtectionMustBeUsedWithCH)),
-                guest_protection_to_use: GuestProtection::NoProtection,
+                guest_protection_to_use: GuestProtection::Tdx(tdx_details.clone()),
             },
             TestData {
                 confidential_guest: true,
                 available_protection: Some(GuestProtection::Tdx(tdx_details.clone())),
                 result: Ok(()),
-                guest_protection_to_use: GuestProtection::Tdx(tdx_details.clone()),
+                guest_protection_to_use: GuestProtection::Tdx(tdx_details),
+            },
+            TestData {
+                confidential_guest: false,
+                available_protection: Some(GuestProtection::Pef),
+                result: Ok(()),
+                guest_protection_to_use: GuestProtection::NoProtection,
+            },
+            TestData {
+                confidential_guest: true,
+                available_protection: Some(GuestProtection::Pef),
+                result: Err(anyhow!(GuestProtectionError::ExpectedTDXProtection(
+                    GuestProtection::Pef
+                ))),
+                guest_protection_to_use: GuestProtection::Pef,
             },
         ];