Merge "forwarder_host terminates when terminating VM" into main
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/InstallerService.java b/android/TerminalApp/java/com/android/virtualization/terminal/InstallerService.java
index 5967b6f..82df5c5 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/InstallerService.java
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/InstallerService.java
@@ -25,6 +25,7 @@
 import android.content.pm.ServiceInfo;
 import android.os.Build;
 import android.os.IBinder;
+import android.os.SELinux;
 import android.util.Log;
 
 import androidx.annotation.Nullable;
@@ -37,6 +38,7 @@
 import org.apache.commons.compress.compressors.gzip.GzipCompressorInputStream;
 
 import java.io.BufferedInputStream;
+import java.io.File;
 import java.io.IOException;
 import java.lang.ref.WeakReference;
 import java.net.URL;
@@ -45,6 +47,7 @@
 import java.nio.file.Path;
 import java.nio.file.StandardCopyOption;
 import java.util.Arrays;
+import java.util.Objects;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 
@@ -60,6 +63,9 @@
                     ? "https://github.com/ikicha/debian_ci/releases/download/release_x86_64/images.tar.gz"
                     : "https://github.com/ikicha/debian_ci/releases/download/release_aarch64/images.tar.gz";
 
+    private static final String SELINUX_FILE_CONTEXT =
+            "u:object_r:virtualizationservice_data_file:";
+
     private final Object mLock = new Object();
 
     private Notification mNotification;
@@ -144,7 +150,9 @@
                 () -> {
                     // TODO(b/374015561): Provide progress update
                     boolean success = downloadFromSdcard() || downloadFromUrl();
-
+                    if (success) {
+                        reLabelImagesSELinuxContext();
+                    }
                     stopForeground(STOP_FOREGROUND_REMOVE);
 
                     synchronized (mLock) {
@@ -156,6 +164,24 @@
                 });
     }
 
+    private void reLabelImagesSELinuxContext() {
+        File payloadFolder = InstallUtils.getInternalStorageDir(this);
+
+        // The context should be u:object_r:privapp_data_file:s0:c35,c257,c512,c768
+        // and we want to get s0:c35,c257,c512,c768 part
+        String level = SELinux.getFileContext(payloadFolder.toString()).split(":", 4)[3];
+        String targetContext = SELINUX_FILE_CONTEXT + level;
+
+        File[] files = payloadFolder.listFiles();
+        for (File file : files) {
+            if (file.isFile() &&
+                    !Objects.equals(SELinux.getFileContext(file.toString()),
+                            targetContext)) {
+                SELinux.setFileContext(file.toString(), targetContext);
+            }
+        }
+    }
+
     private boolean downloadFromSdcard() {
         // Installing from sdcard is preferred, but only supported only in debuggable build.
         if (Build.isDebuggable()) {
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.java b/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.java
index 7575d05..4e32de7 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.java
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.java
@@ -89,6 +89,7 @@
     private ConditionVariable mBootCompleted = new ConditionVariable();
     private static final int POST_NOTIFICATIONS_PERMISSION_REQUEST_CODE = 101;
     private ActivityResultLauncher<Intent> manageExternalStorageActivityResultLauncher;
+    private static int diskSizeStep;
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
@@ -104,6 +105,8 @@
         }
 
         setContentView(R.layout.activity_headless);
+        diskSizeStep = getResources().getInteger(
+                R.integer.disk_size_round_up_step_size_in_mb) << 20;
 
         MaterialToolbar toolbar = (MaterialToolbar) findViewById(R.id.toolbar);
         setSupportActionBar(toolbar);
@@ -293,9 +296,9 @@
         }
     }
 
-    private static File getPartitionFile(Context context, String fileName)
+    public static File getPartitionFile(Context context, String fileName)
             throws FileNotFoundException {
-        File file = new File(context.getFilesDir(), fileName);
+        File file = new File(InstallUtils.getInternalStorageDir(context), fileName);
         if (!file.exists()) {
             Log.d(TAG, file.getAbsolutePath() + " - file not found");
             throw new FileNotFoundException("File not found: " + fileName);
@@ -516,15 +519,14 @@
         return mBootCompleted.block(timeoutMillis);
     }
 
-    private long roundUpDiskSize(long diskSize) {
-        // Round up every disk_size_round_up_step_size_in_mb MB
-        int disk_size_step = getResources().getInteger(
-                R.integer.disk_size_round_up_step_size_in_mb) * 1024 * 1024;
-        return (long) Math.ceil(((double) diskSize) / disk_size_step) * disk_size_step;
+    private static long roundUpDiskSize(long diskSize) {
+        // Round up every diskSizeStep MB
+        return (long) Math.ceil(((double) diskSize) / diskSizeStep) * diskSizeStep;
     }
 
-    private long getMinFilesystemSize(File file) throws IOException, NumberFormatException {
+    public static long getMinFilesystemSize(File file) throws IOException, NumberFormatException {
         try {
+            runE2fsck(file.getAbsolutePath());
             String result = runCommand("/system/bin/resize2fs", "-P", file.getAbsolutePath());
             // The return value is the number of 4k block
             long minSize = Long.parseLong(
@@ -547,12 +549,10 @@
                     getString(R.string.preference_file_key), Context.MODE_PRIVATE);
             SharedPreferences.Editor editor = sharedPref.edit();
 
-            long minDiskSize = getMinFilesystemSize(file);
-            editor.putLong(getString(R.string.preference_min_disk_size_key), minDiskSize);
-
             long currentDiskSize = getFilesystemSize(file);
+            // The default partition size is 6G
             long newSizeInBytes = sharedPref.getLong(getString(R.string.preference_disk_size_key),
-                    roundUpDiskSize(currentDiskSize));
+                    6L << 30);
             editor.putLong(getString(R.string.preference_disk_size_key), newSizeInBytes);
             editor.apply();
 
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/SettingsDiskResizeActivity.kt b/android/TerminalApp/java/com/android/virtualization/terminal/SettingsDiskResizeActivity.kt
index 54e8ab2..58be98d 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/SettingsDiskResizeActivity.kt
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/SettingsDiskResizeActivity.kt
@@ -55,13 +55,10 @@
                     0
                 )
             ).toFloat();
+        val partition = MainActivity.getPartitionFile(this, "root_part")
         val minDiskSizeMb =
-            bytesToMb(
-                    sharedPref.getLong(
-                    getString(R.string.preference_min_disk_size_key),
-                    0
-                )
-            ).toFloat();
+            bytesToMb(MainActivity.getMinFilesystemSize(partition)).toFloat()
+                .coerceAtMost(diskSizeMb)
 
         val diskSizeText = findViewById<TextView>(R.id.settings_disk_resize_resize_gb_assigned)
         val diskMaxSizeText = findViewById<TextView>(R.id.settings_disk_resize_resize_gb_max)
diff --git a/build/debian/build.sh b/build/debian/build.sh
index b4436c1..7fc9035 100755
--- a/build/debian/build.sh
+++ b/build/debian/build.sh
@@ -169,6 +169,19 @@
 	mv "${debian_cloud_image}/image_bookworm_nocloud_${debian_arch}.raw" "${out}"
 }
 
+extract_partitions() {
+	root_partition_num=1
+	efi_partition_num=15
+
+	loop=$(losetup -f --show --partscan image.raw)
+	dd if=${loop}p$root_partition_num of=root_part
+	dd if=${loop}p$efi_partition_num of=efi_part
+	losetup -d ${loop}
+
+	sed -i "s/{root_part_guid}/$(sfdisk --part-uuid image.raw $root_partition_num)/g" vm_config.json
+	sed -i "s/{efi_part_guid}/$(sfdisk --part-uuid image.raw $efi_partition_num)/g" vm_config.json
+}
+
 clean_up() {
 	rm -rf "${workdir}"
 }
@@ -191,18 +204,29 @@
 copy_android_config
 run_fai
 fdisk -l image.raw
-images=(image.raw)
+images=()
+
+cp $(dirname $0)/vm_config.json.${arch} vm_config.json
+
+if [[ "$arch" == "aarch64" ]]; then
+	extract_partitions
+	images+=(
+		root_part
+		efi_part
+	)
+fi
+
 # TODO(b/365955006): remove these lines when uboot supports x86_64 EFI application
 if [[ "$arch" == "x86_64" ]]; then
 	virt-get-kernel -a image.raw
 	mv vmlinuz* vmlinuz
 	mv initrd.img* initrd.img
 	images+=(
+		image.raw
 		vmlinuz
 		initrd.img
 	)
 fi
 
-cp $(dirname $0)/vm_config.json.${arch} vm_config.json
 # --sparse option isn't supported in apache-commons-compress
-tar czv -f images.tar.gz ${images[@]} vm_config.json
\ No newline at end of file
+tar czv -f images.tar.gz ${images[@]} vm_config.json
diff --git a/build/debian/vm_config.json.aarch64 b/build/debian/vm_config.json.aarch64
index 2df0a05..bbe590f 100644
--- a/build/debian/vm_config.json.aarch64
+++ b/build/debian/vm_config.json.aarch64
@@ -2,8 +2,20 @@
     "name": "debian",
     "disks": [
         {
-            "image": "$PAYLOAD_DIR/image.raw",
-            "partitions": [],
+            "partitions": [
+                {
+                    "label": "ROOT",
+                    "path": "$PAYLOAD_DIR/root_part",
+                    "writable": true,
+                    "guid": "{root_part_guid}"
+                },
+                {
+                    "label": "EFI",
+                    "path": "$PAYLOAD_DIR/efi_part",
+                    "writable": false,
+                    "guid": "{efi_part_guid}"
+                }
+            ],
             "writable": true
         }
     ],
diff --git a/guest/pvmfw/Android.bp b/guest/pvmfw/Android.bp
index 4586cca..a5b7494 100644
--- a/guest/pvmfw/Android.bp
+++ b/guest/pvmfw/Android.bp
@@ -16,6 +16,7 @@
         "libcbor_util_nostd",
         "libciborium_nostd",
         "libciborium_io_nostd",
+        "libcoset_nostd",
         "libcstr",
         "libdiced_open_dice_nostd",
         "liblibfdt_nostd",
diff --git a/guest/pvmfw/src/bcc.rs b/guest/pvmfw/src/bcc.rs
index 5317ce9..9260d7f 100644
--- a/guest/pvmfw/src/bcc.rs
+++ b/guest/pvmfw/src/bcc.rs
@@ -21,6 +21,7 @@
 use ciborium::value::Value;
 use core::fmt;
 use core::mem::size_of;
+use coset::{iana, Algorithm, CborSerializable, CoseKey};
 use diced_open_dice::{BccHandover, Cdi, DiceArtifacts, DiceMode};
 use log::trace;
 
@@ -29,16 +30,24 @@
 pub enum BccError {
     CborDecodeError,
     CborEncodeError,
+    CosetError(coset::CoseError),
     DiceError(diced_open_dice::DiceError),
     MalformedBcc(&'static str),
     MissingBcc,
 }
 
+impl From<coset::CoseError> for BccError {
+    fn from(e: coset::CoseError) -> Self {
+        Self::CosetError(e)
+    }
+}
+
 impl fmt::Display for BccError {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             Self::CborDecodeError => write!(f, "Error parsing BCC CBOR"),
             Self::CborEncodeError => write!(f, "Error encoding BCC CBOR"),
+            Self::CosetError(e) => write!(f, "Encountered an error with coset: {e}"),
             Self::DiceError(e) => write!(f, "Dice error: {e:?}"),
             Self::MalformedBcc(s) => {
                 write!(f, "BCC does not have the expected CBOR structure: {s}")
@@ -84,6 +93,7 @@
 /// Represents a (partially) decoded BCC DICE chain.
 pub struct Bcc {
     is_debug_mode: bool,
+    leaf_subject_pubkey: PublicKey,
 }
 
 impl Bcc {
@@ -117,12 +127,18 @@
             .collect::<Result<Vec<_>>>()?;
 
         let is_debug_mode = is_any_payload_debug_mode(&payloads)?;
-        Ok(Self { is_debug_mode })
+        // Safe to unwrap because we checked the length above.
+        let leaf_subject_pubkey = payloads.last().unwrap().subject_public_key()?;
+        Ok(Self { is_debug_mode, leaf_subject_pubkey })
     }
 
     pub fn is_debug_mode(&self) -> bool {
         self.is_debug_mode
     }
+
+    pub fn leaf_subject_pubkey(&self) -> &PublicKey {
+        &self.leaf_subject_pubkey
+    }
 }
 
 fn is_any_payload_debug_mode(payloads: &[BccPayload]) -> Result<bool> {
@@ -144,6 +160,13 @@
 #[repr(transparent)]
 struct BccPayload(Value);
 
+#[derive(Debug, Clone)]
+pub struct PublicKey {
+    /// The COSE key algorithm for the public key, representing the value of the `alg`
+    /// field in the COSE key format of the public key. See RFC 8152, section 7 for details.
+    pub cose_alg: iana::Algorithm,
+}
+
 impl BccEntry {
     pub fn new(entry: Value) -> Self {
         Self(entry)
@@ -178,6 +201,7 @@
 
 const KEY_MODE: i32 = -4670551;
 const MODE_DEBUG: u8 = DiceMode::kDiceModeDebug as u8;
+const SUBJECT_PUBLIC_KEY: i32 = -4670552;
 
 impl BccPayload {
     pub fn is_debug_mode(&self) -> Result<bool> {
@@ -204,6 +228,21 @@
         Ok(mode == MODE_DEBUG.into())
     }
 
+    fn subject_public_key(&self) -> Result<PublicKey> {
+        // BccPayload = {                             ; CWT [RFC8392]
+        // ...
+        //   -4670552 : bstr .cbor PubKeyEd25519 /
+        //              bstr .cbor PubKeyECDSA256 /
+        //              bstr .cbor PubKeyECDSA384,    ; Subject Public Key
+        // ...
+        // }
+        self.value_from_key(SUBJECT_PUBLIC_KEY)
+            .ok_or(BccError::MalformedBcc("Subject public key missing"))?
+            .as_bytes()
+            .ok_or(BccError::MalformedBcc("Subject public key is not a byte string"))
+            .and_then(|v| PublicKey::from_slice(v))
+    }
+
     fn value_from_key(&self, key: i32) -> Option<&Value> {
         // BccPayload is just a map; we only use integral keys, but in general it's legitimate
         // for other things to be present, or for the key we care about not to be present.
@@ -218,3 +257,13 @@
         None
     }
 }
+
+impl PublicKey {
+    fn from_slice(slice: &[u8]) -> Result<Self> {
+        let key = CoseKey::from_slice(slice)?;
+        let Some(Algorithm::Assigned(cose_alg)) = key.alg else {
+            return Err(BccError::MalformedBcc("Invalid algorithm in public key"));
+        };
+        Ok(Self { cose_alg })
+    }
+}
diff --git a/guest/pvmfw/src/main.rs b/guest/pvmfw/src/main.rs
index 1e88c4b..aeced51 100644
--- a/guest/pvmfw/src/main.rs
+++ b/guest/pvmfw/src/main.rs
@@ -200,6 +200,8 @@
         Cow::Owned(truncated_bcc_handover)
     };
 
+    trace!("BCC leaf subject public key algorithm: {:?}", bcc.leaf_subject_pubkey().cose_alg);
+
     dice_inputs
         .write_next_bcc(
             new_bcc_handover.as_ref(),
diff --git a/guest/trusty/security_vm/launcher/src/main.rs b/guest/trusty/security_vm/launcher/src/main.rs
index c5cc6b4..bdb4ed8 100644
--- a/guest/trusty/security_vm/launcher/src/main.rs
+++ b/guest/trusty/security_vm/launcher/src/main.rs
@@ -34,6 +34,14 @@
     /// Whether the VM is protected or not.
     #[arg(long)]
     protected: bool,
+
+    /// Name of the VM. Used to pull correct config from early_vms.xml
+    #[arg(long, default_value = "trusty_security_vm_launcher")]
+    name: String,
+
+    /// Memory size of the VM in MiB
+    #[arg(long, default_value_t = 128)]
+    memory_size_mib: i32,
 }
 
 fn get_service() -> Result<Strong<dyn IVirtualizationService>> {
@@ -51,10 +59,10 @@
         File::open(&args.kernel).with_context(|| format!("Failed to open {:?}", &args.kernel))?;
 
     let vm_config = VirtualMachineConfig::RawConfig(VirtualMachineRawConfig {
-        name: "trusty_security_vm_launcher".to_owned(),
+        name: args.name.to_owned(),
         kernel: Some(ParcelFileDescriptor::new(kernel)),
         protectedVm: args.protected,
-        memoryMib: 128,
+        memoryMib: args.memory_size_mib,
         platformVersion: "~1.0".to_owned(),
         // TODO: add instanceId
         ..Default::default()