Merge changes from topic "prepare-for-system-tee-services" into main

* changes:
  Only allow pVMs to request access to tee services
  virtmgr: also support non-vendor tee services
diff --git a/android/virtmgr/src/aidl.rs b/android/virtmgr/src/aidl.rs
index d53138f..3c5408c 100644
--- a/android/virtmgr/src/aidl.rs
+++ b/android/virtmgr/src/aidl.rs
@@ -1630,9 +1630,8 @@
         | "virtualizationservice_data_file" // files created by VS / VirtMgr
         | "vendor_microdroid_file" // immutable dm-verity protected partition (/vendor/etc/avf/microdroid/.*)
          => Ok(()),
-        // It is difficult to require specific label types for vendor initiated VM's files, so we
-        // allow anything with a vendor prefix.
-        t if calling_partition == CallingPartition::Vendor && t.starts_with("vendor_")  => Ok(()),
+        // It is difficult to require specific label types for vendor initiated VM's files.
+        _ if calling_partition == CallingPartition::Vendor => Ok(()),
         _ => bail!("Label {} is not allowed", context),
     }
 }
diff --git a/android/vm/src/main.rs b/android/vm/src/main.rs
index ff846a1..7178de5 100644
--- a/android/vm/src/main.rs
+++ b/android/vm/src/main.rs
@@ -22,8 +22,6 @@
     CpuOptions::CpuTopology::CpuTopology, IVirtualizationService::IVirtualizationService,
     PartitionType::PartitionType, VirtualMachineAppConfig::DebugLevel::DebugLevel,
 };
-#[cfg(not(llpvm_changes))]
-use anyhow::anyhow;
 use anyhow::{bail, Context, Error};
 use binder::{ProcessState, Strong};
 use clap::{Args, Parser};
@@ -220,7 +218,6 @@
     instance: PathBuf,
 
     /// Path to file containing instance_id. Required iff llpvm feature is enabled.
-    #[cfg(llpvm_changes)]
     #[arg(long = "instance-id-file")]
     instance_id: PathBuf,
 
@@ -255,26 +252,8 @@
         }
     }
 
-    fn instance_id(&self) -> Result<PathBuf, Error> {
-        cfg_if::cfg_if! {
-            if #[cfg(llpvm_changes)] {
-                Ok(self.instance_id.clone())
-            } else {
-                Err(anyhow!("LLPVM feature is disabled, --instance_id flag not supported"))
-            }
-        }
-    }
-
-    fn set_instance_id(&mut self, instance_id_file: PathBuf) -> Result<(), Error> {
-        cfg_if::cfg_if! {
-            if #[cfg(llpvm_changes)] {
-                self.instance_id = instance_id_file;
-                Ok(())
-            } else {
-                let _ = instance_id_file;
-                Err(anyhow!("LLPVM feature is disabled, --instance_id flag not supported"))
-            }
-        }
+    fn set_instance_id(&mut self, instance_id_file: PathBuf) {
+        self.instance_id = instance_id_file;
     }
 }
 
diff --git a/android/vm/src/run.rs b/android/vm/src/run.rs
index 8385fb4..1033164 100644
--- a/android/vm/src/run.rs
+++ b/android/vm/src/run.rs
@@ -87,8 +87,8 @@
         )?;
     }
 
-    let instance_id = if cfg!(llpvm_changes) {
-        let id_file = config.instance_id()?;
+    let instance_id = {
+        let id_file = config.instance_id;
         if id_file.exists() {
             let mut id = [0u8; 64];
             let mut instance_id_file = File::open(id_file)?;
@@ -100,9 +100,6 @@
             instance_id_file.write_all(&id)?;
             id
         }
-    } else {
-        // if llpvm feature flag is disabled, instance_id is not used.
-        [0u8; 64]
     };
 
     let storage = if let Some(ref path) = config.microdroid.storage {
@@ -254,10 +251,8 @@
         ..Default::default()
     };
 
-    if cfg!(llpvm_changes) {
-        app_config.set_instance_id(work_dir.join("instance_id"))?;
-        println!("instance_id file path: {}", app_config.instance_id()?.display());
-    }
+    app_config.set_instance_id(work_dir.join("instance_id"));
+    println!("instance_id file path: {}", app_config.instance_id.display());
 
     command_run_app(app_config)
 }
diff --git a/build/apex/Android.bp b/build/apex/Android.bp
index 20f44fe..79675cb 100644
--- a/build/apex/Android.bp
+++ b/build/apex/Android.bp
@@ -251,11 +251,6 @@
     srcs: [
         "sign_virt_apex.py",
     ],
-    version: {
-        py3: {
-            embedded_launcher: true,
-        },
-    },
     required: [
         // sign_virt_apex should be runnable from outside the source tree,
         // therefore, any required tool should be listed in build/make/core/Makefile as well.
@@ -324,11 +319,6 @@
     srcs: [
         "replace_bytes.py",
     ],
-    version: {
-        py3: {
-            embedded_launcher: true,
-        },
-    },
 }
 
 // Encapsulate the contributions made by the com.android.virt to the bootclasspath.
diff --git a/guest/microdroid_manager/src/main.rs b/guest/microdroid_manager/src/main.rs
index d665c87..4537834 100644
--- a/guest/microdroid_manager/src/main.rs
+++ b/guest/microdroid_manager/src/main.rs
@@ -244,13 +244,14 @@
 fn verify_payload_with_instance_img(
     metadata: &Metadata,
     dice: &DiceDriver,
+    state: &mut VmInstanceState,
 ) -> Result<MicrodroidData> {
     let mut instance = InstanceDisk::new().context("Failed to load instance.img")?;
     let saved_data = instance.read_microdroid_data(dice).context("Failed to read identity data")?;
 
     if is_strict_boot() {
         // Provisioning must happen on the first boot and never again.
-        if is_new_instance_legacy() {
+        if Path::new(AVF_NEW_INSTANCE).exists() {
             ensure!(
                 saved_data.is_none(),
                 MicrodroidError::PayloadInvalidConfig(
@@ -286,12 +287,14 @@
             );
             info!("Saved data is verified.");
         }
+        *state = VmInstanceState::PreviouslySeen;
         saved_data
     } else {
         info!("Saving verified data.");
         instance
             .write_microdroid_data(&extracted_data, dice)
             .context("Failed to write identity data")?;
+        *state = VmInstanceState::NewlyCreated;
         extracted_data
     };
     Ok(instance_data)
@@ -321,13 +324,14 @@
             .context("Failed to load DICE from driver")?
     };
 
+    let mut state = VmInstanceState::Unknown;
     // Microdroid skips checking payload against instance image iff the device supports
-    // secretkeeper. In that case Microdroid use VmSecret::V2, which provide protection against
-    // rollback of boot images and packages.
+    // secretkeeper. In that case Microdroid use VmSecret::V2, which provides instance state
+    // and protection against rollback of boot images and packages.
     let instance_data = if should_defer_rollback_protection() {
         verify_payload(&metadata, None)?
     } else {
-        verify_payload_with_instance_img(&metadata, &dice)?
+        verify_payload_with_instance_img(&metadata, &dice, &mut state)?
     };
 
     let payload_metadata = metadata.payload.ok_or_else(|| {
@@ -337,7 +341,6 @@
     // To minimize the exposure to untrusted data, derive dice profile as soon as possible.
     info!("DICE derivation for payload");
     let dice_artifacts = dice_derivation(dice, &instance_data, &payload_metadata)?;
-    let mut state = VmInstanceState::Unknown;
     let vm_secret = VmSecret::new(dice_artifacts, service, &mut state)
         .context("Failed to create VM secrets")?;
 
@@ -345,15 +348,7 @@
         VmInstanceState::NewlyCreated => true,
         VmInstanceState::PreviouslySeen => false,
         VmInstanceState::Unknown => {
-            // VmSecret instantiation was not able to determine the state. This should only happen
-            // for legacy secret mechanism (V1) - in which case fallback to legacy
-            // instance.img based determination of state.
-            ensure!(
-                !should_defer_rollback_protection(),
-                "VmInstanceState is Unknown whilst guest is expected to use V2 based secrets.
-                This should've never happened"
-            );
-            is_new_instance_legacy()
+            bail!("Vm instance state is still unknown, this should not have happened");
         }
     };
 
@@ -519,10 +514,6 @@
     Path::new(AVF_STRICT_BOOT).exists()
 }
 
-fn is_new_instance_legacy() -> bool {
-    Path::new(AVF_NEW_INSTANCE).exists()
-}
-
 fn is_verified_boot() -> bool {
     !Path::new(DEBUG_MICRODROID_NO_VERIFIED_BOOT).exists()
 }
diff --git a/guest/microdroid_manager/src/vm_secret.rs b/guest/microdroid_manager/src/vm_secret.rs
index f031859..674f010 100644
--- a/guest/microdroid_manager/src/vm_secret.rs
+++ b/guest/microdroid_manager/src/vm_secret.rs
@@ -35,6 +35,9 @@
     StoreSecretRequest, GetSecretResponse, GetSecretRequest};
 use secretkeeper_comm::data_types::error::SecretkeeperError;
 use std::fs;
+use std::thread;
+use rand::Rng;
+use std::time::Duration;
 use zeroize::Zeroizing;
 use std::sync::Mutex;
 use std::sync::Arc;
@@ -63,6 +66,8 @@
     0x55, 0xF8, 0x08, 0x23, 0x81, 0x5F, 0xF5, 0x16, 0x20, 0x3E, 0xBE, 0xBA, 0xB7, 0xA8, 0x43, 0x92,
 ];
 
+const BACKOFF_SK_ACCESS_MS: u64 = 100;
+
 pub enum VmSecret {
     // V2 secrets are derived from 2 independently secured secrets:
     //      1. Secretkeeper protected secrets (skp secret).
@@ -118,15 +123,19 @@
             .map_err(|e| anyhow!("Failed to build a sealing_policy: {e}"))?;
         let session = SkVmSession::new(vm_service, &explicit_dice, policy)?;
         let mut skp_secret = Zeroizing::new([0u8; SECRET_SIZE]);
-        if let Some(secret) = session.get_secret(id)? {
-            *skp_secret = secret;
-            *state = VmInstanceState::PreviouslySeen;
-        } else {
-            log::warn!("No entry found in Secretkeeper for this VM instance, creating new secret.");
-            *skp_secret = rand::random();
-            session.store_secret(id, skp_secret.clone())?;
-            *state = VmInstanceState::NewlyCreated;
-        }
+        get_or_create_sk_secret(&session, id, &mut skp_secret, state).or_else(|e| {
+            // TODO(b/399304956): Secretkeeper rejects requests when overloaded with
+            // connections from multiple clients. Backoff & retry again, hoping it is
+            // less busy then. Secretkeeper changes are required for more robust solutions.
+            log::info!(
+                "get_or_create_sk_secret failed with {e:?}. Refreshing connection & retrying!"
+            );
+            let mut rng = rand::thread_rng();
+            let backoff = rng.gen_range(BACKOFF_SK_ACCESS_MS..2 * BACKOFF_SK_ACCESS_MS);
+            thread::sleep(Duration::from_millis(backoff));
+            session.refresh()?;
+            get_or_create_sk_secret(&session, id, &mut skp_secret, state)
+        })?;
         Ok(Self::V2 {
             instance_id: id,
             dice_artifacts: explicit_dice,
@@ -283,8 +292,6 @@
     sealing_policy: Vec<u8>,
 }
 
-// TODO(b/378911776): This get_secret/store_secret fails on expired session.
-// Introduce retry after refreshing the session
 impl SkVmSession {
     fn new(
         vm_service: &Strong<dyn IVirtualMachineService>,
@@ -366,3 +373,21 @@
             ))
         })?)
 }
+
+fn get_or_create_sk_secret(
+    session: &SkVmSession,
+    id: [u8; ID_SIZE],
+    skp_secret: &mut Zeroizing<[u8; SECRET_SIZE]>,
+    state: &mut VmInstanceState,
+) -> Result<()> {
+    if let Some(secret) = session.get_secret(id)? {
+        **skp_secret = secret;
+        *state = VmInstanceState::PreviouslySeen;
+    } else {
+        log::warn!("No entry found in Secretkeeper for this VM instance, creating new secret.");
+        **skp_secret = rand::random();
+        session.store_secret(id, skp_secret.clone())?;
+        *state = VmInstanceState::NewlyCreated;
+    }
+    Ok(())
+}
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index 418a88e..9d08ed7 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -2166,11 +2166,6 @@
         assumeFalse(
                 "Cuttlefish/Goldfish doesn't support device tree under /proc/device-tree",
                 isCuttlefish() || isGoldfish());
-        if (!isUpdatableVmSupported()) {
-            // TODO(b/389611249): Non protected VMs using legacy secret mechanisms do not reliably
-            // implement `AVmPayload_isNewInstance`.
-            assumeProtectedVM();
-        }
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_a", config);
         TestResults testResults =
                 runVmTestService(