Merge "[apkverify] Refactor verify_signed_data to use SignatureAlgorithmID"
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 828ac9f..ae84c08 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -45,11 +45,15 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
+import java.lang.ref.WeakReference;
 import java.nio.file.FileAlreadyExistsException;
 import java.nio.file.Files;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
+import java.util.WeakHashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -65,6 +69,11 @@
  * @hide
  */
 public class VirtualMachine {
+    private static final Map<Context, Map<String, WeakReference<VirtualMachine>>> sInstances =
+            new WeakHashMap<>();
+
+    private static final Object sInstancesLock = new Object();
+
     /** Name of the directory under the files directory where all VMs created for the app exist. */
     private static final String VM_DIR = "vm";
 
@@ -159,6 +168,8 @@
 
     private final ExecutorService mExecutorService = Executors.newCachedThreadPool();
 
+    @NonNull private final Context mContext;
+
     static {
         System.loadLibrary("virtualmachine_jni");
     }
@@ -166,6 +177,7 @@
     private VirtualMachine(
             @NonNull Context context, @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
+        mContext = context;
         mPackageName = context.getPackageName();
         mName = name;
         mConfig = config;
@@ -231,6 +243,18 @@
             throw new VirtualMachineException("failed to create instance partition", e);
         }
 
+        synchronized (sInstancesLock) {
+            Map<String, WeakReference<VirtualMachine>> instancesMap;
+            if (sInstances.containsKey(context)) {
+                instancesMap = sInstances.get(context);
+            } else {
+                instancesMap = new HashMap<>();
+                sInstances.put(context, instancesMap);
+            }
+
+            instancesMap.put(name, new WeakReference<>(vm));
+        }
+
         return vm;
     }
 
@@ -249,7 +273,23 @@
             throw new VirtualMachineException(e);
         }
 
-        VirtualMachine vm = new VirtualMachine(context, name, config);
+        VirtualMachine vm;
+        synchronized (sInstancesLock) {
+            Map<String, WeakReference<VirtualMachine>> instancesMap;
+            if (sInstances.containsKey(context)) {
+                instancesMap = sInstances.get(context);
+            } else {
+                instancesMap = new HashMap<>();
+                sInstances.put(context, instancesMap);
+            }
+
+            if (instancesMap.containsKey(name)) {
+                vm = instancesMap.get(name).get();
+            } else {
+                vm = new VirtualMachine(context, name, config);
+                instancesMap.put(name, new WeakReference<>(vm));
+            }
+        }
 
         // If config file exists, but the instance image file doesn't, it means that the VM is
         // corrupted. That's different from the case that the VM doesn't exist. Throw an exception
@@ -544,6 +584,11 @@
         mInstanceFilePath.delete();
         mIdsigFilePath.delete();
         vmRootDir.delete();
+
+        synchronized (sInstancesLock) {
+            Map<String, WeakReference<VirtualMachine>> instancesMap = sInstances.get(mContext);
+            if (instancesMap != null) instancesMap.remove(mName);
+        }
     }
 
     /**
diff --git a/libs/apkverify/src/lib.rs b/libs/apkverify/src/lib.rs
index f6c1a21..040c304 100644
--- a/libs/apkverify/src/lib.rs
+++ b/libs/apkverify/src/lib.rs
@@ -24,5 +24,5 @@
 mod v3;
 mod ziputil;
 
-// TODO(jooyung) fallback to v2 when v3 not found
+// TODO(b/197052981) fallback to v2 when v3 not found
 pub use v3::{get_public_key_der, pick_v4_apk_digest, verify};
diff --git a/libs/apkverify/src/v3.rs b/libs/apkverify/src/v3.rs
index 2f13837..557abcd 100644
--- a/libs/apkverify/src/v3.rs
+++ b/libs/apkverify/src/v3.rs
@@ -34,7 +34,7 @@
 
 pub const APK_SIGNATURE_SCHEME_V3_BLOCK_ID: u32 = 0xf05368c0;
 
-// TODO(jooyung): get "ro.build.version.sdk"
+// TODO(b/190343842): get "ro.build.version.sdk"
 const SDK_INT: u32 = 31;
 
 type Signers = LengthPrefixed<Vec<LengthPrefixed<Signer>>>;
@@ -215,7 +215,9 @@
             bail!("Public key mismatch between certificate and signature record");
         }
 
-        // TODO(jooyung) 8. If the proof-of-rotation attribute exists for the signer verify that the struct is valid and this signer is the last certificate in the list.
+        // TODO(b/245914104)
+        // 8. If the proof-of-rotation attribute exists for the signer verify that the
+        // struct is valid and this signer is the last certificate in the list.
         Ok(self.public_key.to_vec().into_boxed_slice())
     }
 }
@@ -235,7 +237,7 @@
 }
 
 // ReadFromBytes implementations
-// TODO(jooyung): add derive macro: #[derive(ReadFromBytes)]
+// TODO(b/190343842): add derive macro: #[derive(ReadFromBytes)]
 
 impl ReadFromBytes for Signer {
     fn read_from_bytes(buf: &mut Bytes) -> Result<Self> {
diff --git a/libs/apkverify/tests/apkverify_test.rs b/libs/apkverify/tests/apkverify_test.rs
index 0b33208..7e1438f 100644
--- a/libs/apkverify/tests/apkverify_test.rs
+++ b/libs/apkverify/tests/apkverify_test.rs
@@ -25,7 +25,7 @@
 fn test_verify_truncated_cd() {
     use zip::result::ZipError;
     let res = verify("tests/data/v2-only-truncated-cd.apk");
-    // TODO(jooyung): consider making a helper for err assertion
+    // TODO(b/190343842): consider making a helper for err assertion
     assert!(matches!(
         res.unwrap_err().root_cause().downcast_ref::<ZipError>().unwrap(),
         ZipError::InvalidArchive(_),
@@ -37,7 +37,6 @@
     assert!(verify("tests/data/test.apex").is_ok());
 }
 
-// TODO(b/190343842)
 #[test]
 fn test_verify_v3_dsa_sha256() {
     for key_name in KEY_NAMES_DSA.iter() {
@@ -54,7 +53,6 @@
     }
 }
 
-// TODO(b/190343842)
 #[test]
 fn test_verify_v3_ecdsa_sha512() {
     for key_name in KEY_NAMES_ECDSA.iter() {
@@ -82,7 +80,6 @@
     }
 }
 
-// TODO(b/190343842)
 #[test]
 fn test_verify_v3_sig_does_not_verify() {
     let path_list = [
@@ -100,7 +97,6 @@
     }
 }
 
-// TODO(b/190343842)
 #[test]
 fn test_verify_v3_digest_mismatch() {
     let path_list = [
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 99afe98..c2060cb 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -504,4 +504,21 @@
         assertThat(bootResult.deathReason).isEqualTo(
                 VirtualMachineCallback.DEATH_REASON_MICRODROID_INVALID_PAYLOAD_CONFIG);
     }
+
+    @Test
+    public void sameInstancesShareTheSameVmObject()
+            throws VirtualMachineException, InterruptedException, IOException {
+        VirtualMachineConfig.Builder builder =
+                mInner.newVmConfigBuilder("assets/vm_config.json");
+        VirtualMachineConfig normalConfig = builder.debugLevel(DebugLevel.NONE).build();
+        VirtualMachine vm = mInner.forceCreateNewVirtualMachine("test_vm", normalConfig);
+        VirtualMachine vm2 = mInner.getVirtualMachineManager().get("test_vm");
+        assertThat(vm).isEqualTo(vm2);
+
+        VirtualMachine newVm = mInner.forceCreateNewVirtualMachine("test_vm", normalConfig);
+        VirtualMachine newVm2 = mInner.getVirtualMachineManager().get("test_vm");
+        assertThat(newVm).isEqualTo(newVm2);
+
+        assertThat(vm).isNotEqualTo(newVm);
+    }
 }