Add a limit to how many VMs one app can own

Running `atest MicrodroidTests` creates ~260 VMs, so choose a bigger
limit.

Bug: 294177871
Test: virtualizationservice_test
Change-Id: Idb1b7a54d1ff8fb79f7024d325221b3912560b0a
diff --git a/virtualizationservice/src/maintenance.rs b/virtualizationservice/src/maintenance.rs
index f950db9..8efc58d 100644
--- a/virtualizationservice/src/maintenance.rs
+++ b/virtualizationservice/src/maintenance.rs
@@ -40,6 +40,9 @@
 /// parcel fits within max AIDL message size.
 const DELETE_MAX_BATCH_SIZE: usize = 100;
 
+/// Maximum number of VM IDs that a single app can have.
+const MAX_VM_IDS_PER_APP: usize = 400;
+
 /// State related to VM secrets.
 pub struct State {
     sk: binder::Strong<dyn ISecretkeeper>,
@@ -101,6 +104,24 @@
     pub fn add_id(&mut self, vm_id: &VmId, user_id: u32, app_id: u32) -> Result<()> {
         let user_id: i32 = user_id.try_into().context(format!("user_id {user_id} out of range"))?;
         let app_id: i32 = app_id.try_into().context(format!("app_id {app_id} out of range"))?;
+
+        // To prevent unbounded growth of VM IDs (and the associated state) for an app, limit the
+        // number of VM IDs per app.
+        let count = self
+            .vm_id_db
+            .count_vm_ids_for_app(user_id, app_id)
+            .context("failed to determine VM count")?;
+        if count >= MAX_VM_IDS_PER_APP {
+            // The owner has too many VM IDs, so delete the oldest IDs so that the new VM ID
+            // creation can progress/succeed.
+            let purge = 1 + count - MAX_VM_IDS_PER_APP;
+            let old_vm_ids = self
+                .vm_id_db
+                .oldest_vm_ids_for_app(user_id, app_id, purge)
+                .context("failed to find oldest VM IDs")?;
+            error!("Deleting {purge} of {count} VM IDs for user_id={user_id}, app_id={app_id}");
+            self.delete_ids(&old_vm_ids);
+        }
         self.vm_id_db.add_vm_id(vm_id, user_id, app_id)
     }
 
@@ -396,6 +417,39 @@
         assert_eq!(vec![VM_ID5], sk_state.vm_id_db.vm_ids_for_user(USER3).unwrap());
     }
 
+    #[test]
+    fn test_sk_state_too_many_vms() {
+        let history = Arc::new(Mutex::new(Vec::new()));
+        let mut sk_state = new_test_state(history.clone(), 20);
+
+        // Every VM ID added up to the limit is kept.
+        for idx in 0..MAX_VM_IDS_PER_APP {
+            let mut vm_id = [0u8; 64];
+            vm_id[0..8].copy_from_slice(&(idx as u64).to_be_bytes());
+            sk_state.add_id(&vm_id, USER1 as u32, APP_A as u32).unwrap();
+            assert_eq!(idx + 1, sk_state.vm_id_db.count_vm_ids_for_app(USER1, APP_A).unwrap());
+        }
+        assert_eq!(
+            MAX_VM_IDS_PER_APP,
+            sk_state.vm_id_db.count_vm_ids_for_app(USER1, APP_A).unwrap()
+        );
+
+        // Beyond the limit it's one in, one out.
+        for idx in MAX_VM_IDS_PER_APP..MAX_VM_IDS_PER_APP + 10 {
+            let mut vm_id = [0u8; 64];
+            vm_id[0..8].copy_from_slice(&(idx as u64).to_be_bytes());
+            sk_state.add_id(&vm_id, USER1 as u32, APP_A as u32).unwrap();
+            assert_eq!(
+                MAX_VM_IDS_PER_APP,
+                sk_state.vm_id_db.count_vm_ids_for_app(USER1, APP_A).unwrap()
+            );
+        }
+        assert_eq!(
+            MAX_VM_IDS_PER_APP,
+            sk_state.vm_id_db.count_vm_ids_for_app(USER1, APP_A).unwrap()
+        );
+    }
+
     struct Irreconcilable;
 
     impl IVirtualizationReconciliationCallback for Irreconcilable {
diff --git a/virtualizationservice/src/maintenance/vmdb.rs b/virtualizationservice/src/maintenance/vmdb.rs
index 47704bc..273f340 100644
--- a/virtualizationservice/src/maintenance/vmdb.rs
+++ b/virtualizationservice/src/maintenance/vmdb.rs
@@ -272,6 +272,34 @@
         Ok(vm_ids)
     }
 
+    /// Determine the number of VM IDs associated with `(user_id, app_id)`.
+    pub fn count_vm_ids_for_app(&mut self, user_id: i32, app_id: i32) -> Result<usize> {
+        let mut stmt = self
+            .conn
+            .prepare("SELECT COUNT(vm_id) FROM main.vmids WHERE user_id = ? AND app_id = ?;")
+            .context("failed to prepare SELECT stmt")?;
+        stmt.query_row(params![user_id, app_id], |row| row.get(0)).context("query failed")
+    }
+
+    /// Return the `count` oldest VM IDs associated with `(user_id, app_id)`.
+    pub fn oldest_vm_ids_for_app(
+        &mut self,
+        user_id: i32,
+        app_id: i32,
+        count: usize,
+    ) -> Result<Vec<VmId>> {
+        // SQLite considers NULL columns to be smaller than values, so rows left over from a v0
+        // database will be listed first.
+        let mut stmt = self
+            .conn
+            .prepare(
+                "SELECT vm_id FROM main.vmids WHERE user_id = ? AND app_id = ? ORDER BY created LIMIT ?;",
+            )
+            .context("failed to prepare SELECT stmt")?;
+        let rows = stmt.query(params![user_id, app_id, count]).context("query failed")?;
+        Self::vm_ids_from_rows(rows)
+    }
+
     /// Return all of the `(user_id, app_id)` pairs present in the database.
     pub fn get_all_owners(&mut self) -> Result<Vec<(i32, i32)>> {
         let mut stmt = self
@@ -344,6 +372,19 @@
     fn show_contents(db: &VmIdDb) {
         let mut stmt = db.conn.prepare("SELECT * FROM main.vmids;").unwrap();
         let mut rows = stmt.query(()).unwrap();
+        println!("DB contents:");
+        while let Some(row) = rows.next().unwrap() {
+            println!("  {row:?}");
+        }
+    }
+
+    fn show_contents_for_app(db: &VmIdDb, user_id: i32, app_id: i32, count: usize) {
+        let mut stmt = db
+            .conn
+            .prepare("SELECT vm_id, created FROM main.vmids WHERE user_id = ? AND app_id = ? ORDER BY created LIMIT ?;")
+            .unwrap();
+        let mut rows = stmt.query(params![user_id, app_id, count]).unwrap();
+        println!("First (by created) {count} rows for app_id={app_id}");
         while let Some(row) = rows.next().unwrap() {
             println!("  {row:?}");
         }
@@ -457,31 +498,39 @@
 
         assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_user(USER1).unwrap());
         assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_app(USER1, APP_A).unwrap());
+        assert_eq!(3, db.count_vm_ids_for_app(USER1, APP_A).unwrap());
         assert_eq!(vec![VM_ID4], db.vm_ids_for_app(USER2, APP_B).unwrap());
+        assert_eq!(1, db.count_vm_ids_for_app(USER2, APP_B).unwrap());
         assert_eq!(vec![VM_ID5], db.vm_ids_for_user(USER3).unwrap());
         assert_eq!(empty, db.vm_ids_for_user(USER_UNKNOWN).unwrap());
         assert_eq!(empty, db.vm_ids_for_app(USER1, APP_UNKNOWN).unwrap());
+        assert_eq!(0, db.count_vm_ids_for_app(USER1, APP_UNKNOWN).unwrap());
 
         db.delete_vm_ids(&[VM_ID2, VM_ID3]).unwrap();
 
         assert_eq!(vec![VM_ID1], db.vm_ids_for_user(USER1).unwrap());
         assert_eq!(vec![VM_ID1], db.vm_ids_for_app(USER1, APP_A).unwrap());
+        assert_eq!(1, db.count_vm_ids_for_app(USER1, APP_A).unwrap());
 
         // OK to delete things that don't exist.
         db.delete_vm_ids(&[VM_ID2, VM_ID3]).unwrap();
 
         assert_eq!(vec![VM_ID1], db.vm_ids_for_user(USER1).unwrap());
         assert_eq!(vec![VM_ID1], db.vm_ids_for_app(USER1, APP_A).unwrap());
+        assert_eq!(1, db.count_vm_ids_for_app(USER1, APP_A).unwrap());
 
         db.add_vm_id(&VM_ID2, USER1, APP_A).unwrap();
         db.add_vm_id(&VM_ID3, USER1, APP_A).unwrap();
 
         assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_user(USER1).unwrap());
         assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_app(USER1, APP_A).unwrap());
+        assert_eq!(3, db.count_vm_ids_for_app(USER1, APP_A).unwrap());
         assert_eq!(vec![VM_ID4], db.vm_ids_for_app(USER2, APP_B).unwrap());
+        assert_eq!(1, db.count_vm_ids_for_app(USER2, APP_B).unwrap());
         assert_eq!(vec![VM_ID5], db.vm_ids_for_user(USER3).unwrap());
         assert_eq!(empty, db.vm_ids_for_user(USER_UNKNOWN).unwrap());
         assert_eq!(empty, db.vm_ids_for_app(USER1, APP_UNKNOWN).unwrap());
+        assert_eq!(0, db.count_vm_ids_for_app(USER1, APP_UNKNOWN).unwrap());
 
         assert_eq!(
             vec![(USER1, APP_A), (USER2, APP_B), (USER3, APP_C)],
@@ -513,4 +562,47 @@
         assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_user(USER1).unwrap());
         show_contents(&db);
     }
+
+    #[test]
+    fn test_remove_oldest_with_upgrade() {
+        let mut db = new_test_db_version(0);
+        let version = db.schema_version().unwrap();
+        assert_eq!(0, version);
+
+        let remove_count = 10;
+        let mut want = vec![];
+
+        // Manually insert rows before upgrade.
+        const V0_COUNT: usize = 5;
+        for idx in 0..V0_COUNT {
+            let mut vm_id = [0u8; 64];
+            vm_id[0..8].copy_from_slice(&(idx as u64).to_be_bytes());
+            if want.len() < remove_count {
+                want.push(vm_id);
+            }
+            db.conn
+                .execute(
+                    "REPLACE INTO main.vmids (vm_id, user_id, app_id) VALUES (?1, ?2, ?3);",
+                    params![&vm_id, &USER1, APP_A],
+                )
+                .unwrap();
+        }
+
+        // Now move to v1.
+        db.upgrade_tables_v0_v1().unwrap();
+        let version = db.schema_version().unwrap();
+        assert_eq!(1, version);
+
+        for idx in V0_COUNT..40 {
+            let mut vm_id = [0u8; 64];
+            vm_id[0..8].copy_from_slice(&(idx as u64).to_be_bytes());
+            if want.len() < remove_count {
+                want.push(vm_id);
+            }
+            db.add_vm_id(&vm_id, USER1, APP_A).unwrap();
+        }
+        show_contents_for_app(&db, USER1, APP_A, 10);
+        let got = db.oldest_vm_ids_for_app(USER1, APP_A, 10).unwrap();
+        assert_eq!(got, want);
+    }
 }