1 // Copyright 2024, The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 //! Database of VM IDs.
16 
17 use anyhow::{anyhow, Context, Result};
18 use log::{debug, error, info, warn};
19 use rusqlite::{params, params_from_iter, Connection, OpenFlags, Rows};
20 use std::path::PathBuf;
21 
22 /// Subdirectory to hold the database.
23 const DB_DIR: &str = "vmdb";
24 
25 /// Name of the file that holds the database.
26 const DB_FILENAME: &str = "vmids.sqlite";
27 
28 /// Maximum number of host parameters in a single SQL statement.
29 /// (Default value of `SQLITE_LIMIT_VARIABLE_NUMBER` for <= 3.32.0)
30 const MAX_VARIABLES: usize = 999;
31 
32 /// Return the current time as milliseconds since epoch.
db_now() -> u6433 fn db_now() -> u64 {
34     let now = std::time::SystemTime::now()
35         .duration_since(std::time::UNIX_EPOCH)
36         .unwrap_or(std::time::Duration::ZERO)
37         .as_millis();
38     now.try_into().unwrap_or(u64::MAX)
39 }
40 
41 /// Identifier for a VM and its corresponding secret.
42 pub type VmId = [u8; 64];
43 
44 /// Representation of an on-disk database of VM IDs.
45 pub struct VmIdDb {
46     conn: Connection,
47 }
48 
49 struct RetryOnFailure(bool);
50 
51 impl VmIdDb {
52     /// Connect to the VM ID database file held in the given directory, creating it if necessary.
53     /// The second return value indicates whether a new database file was created.
54     ///
55     /// This function assumes no other threads/processes are attempting to connect concurrently.
new(db_dir: &str) -> Result<(Self, bool)>56     pub fn new(db_dir: &str) -> Result<(Self, bool)> {
57         let mut db_path = PathBuf::from(db_dir);
58         db_path.push(DB_DIR);
59         if !db_path.exists() {
60             std::fs::create_dir(&db_path).context("failed to create {db_path:?}")?;
61             info!("created persistent db dir {db_path:?}");
62         }
63         db_path.push(DB_FILENAME);
64         Self::new_at_path(db_path, RetryOnFailure(true))
65     }
66 
new_at_path(db_path: PathBuf, retry: RetryOnFailure) -> Result<(Self, bool)>67     fn new_at_path(db_path: PathBuf, retry: RetryOnFailure) -> Result<(Self, bool)> {
68         let (flags, created) = if db_path.exists() {
69             debug!("connecting to existing database {db_path:?}");
70             (
71                 OpenFlags::SQLITE_OPEN_READ_WRITE
72                     | OpenFlags::SQLITE_OPEN_URI
73                     | OpenFlags::SQLITE_OPEN_NO_MUTEX,
74                 false,
75             )
76         } else {
77             info!("creating fresh database {db_path:?}");
78             (
79                 OpenFlags::SQLITE_OPEN_READ_WRITE
80                     | OpenFlags::SQLITE_OPEN_CREATE
81                     | OpenFlags::SQLITE_OPEN_URI
82                     | OpenFlags::SQLITE_OPEN_NO_MUTEX,
83                 true,
84             )
85         };
86         let mut db = Self {
87             conn: Connection::open_with_flags(&db_path, flags)
88                 .context(format!("failed to open/create DB with {flags:?}"))?,
89         };
90 
91         if created {
92             db.init_tables().context("failed to create tables")?;
93         } else {
94             // An existing .sqlite file may have an earlier schema.
95             match db.schema_version() {
96                 Err(e) => {
97                     // Couldn't determine a schema version, so wipe and try again.
98                     error!("failed to determine VM DB schema: {e:?}");
99                     if retry.0 {
100                         // This is the first attempt, so wipe and retry.
101                         error!("resetting database file {db_path:?}");
102                         let _ = std::fs::remove_file(&db_path);
103                         return Self::new_at_path(db_path, RetryOnFailure(false));
104                     } else {
105                         // An earlier attempt at wiping/retrying has failed, so give up.
106                         return Err(anyhow!("failed to reset database file {db_path:?}"));
107                     }
108                 }
109                 Ok(0) => db.upgrade_tables_v0_v1().context("failed to upgrade schema v0 -> v1")?,
110                 Ok(1) => {
111                     // Current version, no action needed.
112                 }
113                 Ok(version) => {
114                     // If the database looks like it's from a future version, leave it alone and
115                     // fail to connect to it.
116                     error!("database from the future (v{version})");
117                     return Err(anyhow!("database from the future (v{version})"));
118                 }
119             }
120         }
121         Ok((db, created))
122     }
123 
124     /// Delete the associated database file.
delete_db_file(self, db_dir: &str)125     pub fn delete_db_file(self, db_dir: &str) {
126         let mut db_path = PathBuf::from(db_dir);
127         db_path.push(DB_DIR);
128         db_path.push(DB_FILENAME);
129 
130         // Drop the connection before removing the backing file.
131         drop(self);
132         warn!("removing database file {db_path:?}");
133         if let Err(e) = std::fs::remove_file(&db_path) {
134             error!("failed to remove database file {db_path:?}: {e:?}");
135         }
136     }
137 
schema_version(&mut self) -> Result<i32>138     fn schema_version(&mut self) -> Result<i32> {
139         let version: i32 = self
140             .conn
141             .query_row("PRAGMA main.user_version", (), |row| row.get(0))
142             .context("failed to read pragma")?;
143         Ok(version)
144     }
145 
146     /// Create the database table and indices using the current schema.
init_tables(&mut self) -> Result<()>147     fn init_tables(&mut self) -> Result<()> {
148         self.init_tables_v1()
149     }
150 
151     /// Create the database table and indices using the v1 schema.
init_tables_v1(&mut self) -> Result<()>152     fn init_tables_v1(&mut self) -> Result<()> {
153         info!("creating v1 database schema");
154         self.conn
155             .execute(
156                 "CREATE TABLE IF NOT EXISTS main.vmids (
157                      vm_id BLOB PRIMARY KEY,
158                      user_id INTEGER,
159                      app_id INTEGER,
160                      created INTEGER
161                  ) WITHOUT ROWID;",
162                 (),
163             )
164             .context("failed to create table")?;
165         self.conn
166             .execute("CREATE INDEX IF NOT EXISTS main.vmids_user_index ON vmids(user_id);", [])
167             .context("Failed to create user index")?;
168         self.conn
169             .execute(
170                 "CREATE INDEX IF NOT EXISTS main.vmids_app_index ON vmids(user_id, app_id);",
171                 [],
172             )
173             .context("Failed to create app index")?;
174         self.conn
175             .execute("PRAGMA main.user_version = 1;", ())
176             .context("failed to declare version")?;
177         Ok(())
178     }
179 
upgrade_tables_v0_v1(&mut self) -> Result<()>180     fn upgrade_tables_v0_v1(&mut self) -> Result<()> {
181         let _rows = self
182             .conn
183             .execute("ALTER TABLE main.vmids ADD COLUMN created INTEGER;", ())
184             .context("failed to alter table v0->v1")?;
185         self.conn
186             .execute("PRAGMA main.user_version = 1;", ())
187             .context("failed to set schema version")?;
188         Ok(())
189     }
190 
191     /// Create the database table and indices using the v0 schema.
192     #[cfg(test)]
init_tables_v0(&mut self) -> Result<()>193     fn init_tables_v0(&mut self) -> Result<()> {
194         info!("creating v0 database schema");
195         self.conn
196             .execute(
197                 "CREATE TABLE IF NOT EXISTS main.vmids (
198                      vm_id BLOB PRIMARY KEY,
199                      user_id INTEGER,
200                      app_id INTEGER
201                  ) WITHOUT ROWID;",
202                 (),
203             )
204             .context("failed to create table")?;
205         self.conn
206             .execute("CREATE INDEX IF NOT EXISTS main.vmids_user_index ON vmids(user_id);", [])
207             .context("Failed to create user index")?;
208         self.conn
209             .execute(
210                 "CREATE INDEX IF NOT EXISTS main.vmids_app_index ON vmids(user_id, app_id);",
211                 [],
212             )
213             .context("Failed to create app index")?;
214         Ok(())
215     }
216 
217     /// Add the given VM ID into the database.
add_vm_id(&mut self, vm_id: &VmId, user_id: i32, app_id: i32) -> Result<()>218     pub fn add_vm_id(&mut self, vm_id: &VmId, user_id: i32, app_id: i32) -> Result<()> {
219         let now = db_now();
220         let _rows = self
221             .conn
222             .execute(
223                 "REPLACE INTO main.vmids (vm_id, user_id, app_id, created) VALUES (?1, ?2, ?3, ?4);",
224                 params![vm_id, &user_id, &app_id, &now],
225             )
226             .context("failed to add VM ID")?;
227         Ok(())
228     }
229 
230     /// Remove the given VM IDs from the database.  The collection of IDs is assumed to be smaller
231     /// than the maximum number of SQLite parameters.
delete_vm_ids(&mut self, vm_ids: &[VmId]) -> Result<()>232     pub fn delete_vm_ids(&mut self, vm_ids: &[VmId]) -> Result<()> {
233         assert!(vm_ids.len() < MAX_VARIABLES);
234         let mut vars = "?,".repeat(vm_ids.len());
235         vars.pop(); // remove trailing comma
236         let sql = format!("DELETE FROM main.vmids WHERE vm_id IN ({});", vars);
237         let mut stmt = self.conn.prepare(&sql).context("failed to prepare DELETE stmt")?;
238         let _rows = stmt.execute(params_from_iter(vm_ids)).context("failed to delete VM IDs")?;
239         Ok(())
240     }
241 
242     /// Return the VM IDs associated with Android user ID `user_id`.
vm_ids_for_user(&mut self, user_id: i32) -> Result<Vec<VmId>>243     pub fn vm_ids_for_user(&mut self, user_id: i32) -> Result<Vec<VmId>> {
244         let mut stmt = self
245             .conn
246             .prepare("SELECT vm_id FROM main.vmids WHERE user_id = ?;")
247             .context("failed to prepare SELECT stmt")?;
248         let rows = stmt.query(params![user_id]).context("query failed")?;
249         Self::vm_ids_from_rows(rows)
250     }
251 
252     /// Return the VM IDs associated with `(user_id, app_id)`.
vm_ids_for_app(&mut self, user_id: i32, app_id: i32) -> Result<Vec<VmId>>253     pub fn vm_ids_for_app(&mut self, user_id: i32, app_id: i32) -> Result<Vec<VmId>> {
254         let mut stmt = self
255             .conn
256             .prepare("SELECT vm_id FROM main.vmids WHERE user_id = ? AND app_id = ?;")
257             .context("failed to prepare SELECT stmt")?;
258         let rows = stmt.query(params![user_id, app_id]).context("query failed")?;
259         Self::vm_ids_from_rows(rows)
260     }
261 
262     /// Retrieve a collection of VM IDs from database rows.
vm_ids_from_rows(mut rows: Rows) -> Result<Vec<VmId>>263     fn vm_ids_from_rows(mut rows: Rows) -> Result<Vec<VmId>> {
264         let mut vm_ids: Vec<VmId> = Vec::new();
265         while let Some(row) = rows.next().context("failed row unpack")? {
266             match row.get(0) {
267                 Ok(vm_id) => vm_ids.push(vm_id),
268                 Err(e) => error!("failed to parse row: {e:?}"),
269             }
270         }
271 
272         Ok(vm_ids)
273     }
274 
275     /// Determine whether the specified VM ID is associated with `(user_id, app_id)`. Returns false
276     /// if there is no such VM ID, or it exists but is not associated.
is_vm_id_for_app(&mut self, vm_id: &VmId, user_id: u32, app_id: u32) -> Result<bool>277     pub fn is_vm_id_for_app(&mut self, vm_id: &VmId, user_id: u32, app_id: u32) -> Result<bool> {
278         let mut stmt = self
279             .conn
280             .prepare(
281                 "SELECT COUNT(*) FROM main.vmids \
282                         WHERE vm_id = ? AND user_id = ? AND app_id = ?;",
283             )
284             .context("failed to prepare SELECT stmt")?;
285         stmt.query_row(params![vm_id, user_id, app_id], |row| row.get(0))
286             .context("query failed")
287             .map(|n: usize| n != 0)
288     }
289 
290     /// Determine the number of VM IDs associated with `(user_id, app_id)`.
count_vm_ids_for_app(&mut self, user_id: i32, app_id: i32) -> Result<usize>291     pub fn count_vm_ids_for_app(&mut self, user_id: i32, app_id: i32) -> Result<usize> {
292         let mut stmt = self
293             .conn
294             .prepare("SELECT COUNT(vm_id) FROM main.vmids WHERE user_id = ? AND app_id = ?;")
295             .context("failed to prepare SELECT stmt")?;
296         stmt.query_row(params![user_id, app_id], |row| row.get(0)).context("query failed")
297     }
298 
299     /// Return the `count` oldest VM IDs associated with `(user_id, app_id)`.
oldest_vm_ids_for_app( &mut self, user_id: i32, app_id: i32, count: usize, ) -> Result<Vec<VmId>>300     pub fn oldest_vm_ids_for_app(
301         &mut self,
302         user_id: i32,
303         app_id: i32,
304         count: usize,
305     ) -> Result<Vec<VmId>> {
306         // SQLite considers NULL columns to be smaller than values, so rows left over from a v0
307         // database will be listed first.
308         let mut stmt = self
309             .conn
310             .prepare(
311                 "SELECT vm_id FROM main.vmids WHERE user_id = ? AND app_id = ? ORDER BY created LIMIT ?;",
312             )
313             .context("failed to prepare SELECT stmt")?;
314         let rows = stmt.query(params![user_id, app_id, count]).context("query failed")?;
315         Self::vm_ids_from_rows(rows)
316     }
317 
318     /// Return all of the `(user_id, app_id)` pairs present in the database.
get_all_owners(&mut self) -> Result<Vec<(i32, i32)>>319     pub fn get_all_owners(&mut self) -> Result<Vec<(i32, i32)>> {
320         let mut stmt = self
321             .conn
322             .prepare("SELECT DISTINCT user_id, app_id FROM main.vmids;")
323             .context("failed to prepare SELECT stmt")?;
324         let mut rows = stmt.query(()).context("query failed")?;
325         let mut owners: Vec<(i32, i32)> = Vec::new();
326         while let Some(row) = rows.next().context("failed row unpack")? {
327             let user_id = match row.get(0) {
328                 Ok(v) => v,
329                 Err(e) => {
330                     error!("failed to parse row: {e:?}");
331                     continue;
332                 }
333             };
334             let app_id = match row.get(1) {
335                 Ok(v) => v,
336                 Err(e) => {
337                     error!("failed to parse row: {e:?}");
338                     continue;
339                 }
340             };
341             owners.push((user_id, app_id));
342         }
343 
344         Ok(owners)
345     }
346 }
347 
348 /// Current schema version.
349 #[cfg(test)]
350 const SCHEMA_VERSION: usize = 1;
351 
352 /// Create a new in-memory database for testing.
353 #[cfg(test)]
new_test_db() -> VmIdDb354 pub fn new_test_db() -> VmIdDb {
355     tests::new_test_db_version(SCHEMA_VERSION)
356 }
357 
358 #[cfg(test)]
359 mod tests {
360     use super::*;
361     use std::io::Write;
362 
363     const VM_ID1: VmId = [1u8; 64];
364     const VM_ID2: VmId = [2u8; 64];
365     const VM_ID3: VmId = [3u8; 64];
366     const VM_ID4: VmId = [4u8; 64];
367     const VM_ID5: VmId = [5u8; 64];
368     const VM_ID_UNKNOWN: VmId = [6u8; 64];
369     const USER1: i32 = 1;
370     const USER2: i32 = 2;
371     const USER3: i32 = 3;
372     const USER_UNKNOWN: i32 = 4;
373     const APP_A: i32 = 50;
374     const APP_B: i32 = 60;
375     const APP_C: i32 = 70;
376     const APP_UNKNOWN: i32 = 99;
377 
new_test_db_version(version: usize) -> VmIdDb378     pub fn new_test_db_version(version: usize) -> VmIdDb {
379         let mut db = VmIdDb { conn: Connection::open_in_memory().unwrap() };
380         match version {
381             0 => db.init_tables_v0().unwrap(),
382             1 => db.init_tables_v1().unwrap(),
383             _ => panic!("unexpected version {version}"),
384         }
385         db
386     }
387 
show_contents(db: &VmIdDb)388     fn show_contents(db: &VmIdDb) {
389         let mut stmt = db.conn.prepare("SELECT * FROM main.vmids;").unwrap();
390         let mut rows = stmt.query(()).unwrap();
391         println!("DB contents:");
392         while let Some(row) = rows.next().unwrap() {
393             println!("  {row:?}");
394         }
395     }
396 
show_contents_for_app(db: &VmIdDb, user_id: i32, app_id: i32, count: usize)397     fn show_contents_for_app(db: &VmIdDb, user_id: i32, app_id: i32, count: usize) {
398         let mut stmt = db
399             .conn
400             .prepare("SELECT vm_id, created FROM main.vmids WHERE user_id = ? AND app_id = ? ORDER BY created LIMIT ?;")
401             .unwrap();
402         let mut rows = stmt.query(params![user_id, app_id, count]).unwrap();
403         println!("First (by created) {count} rows for app_id={app_id}");
404         while let Some(row) = rows.next().unwrap() {
405             println!("  {row:?}");
406         }
407     }
408 
409     #[test]
test_schema_version0()410     fn test_schema_version0() {
411         let mut db0 = VmIdDb { conn: Connection::open_in_memory().unwrap() };
412         db0.init_tables_v0().unwrap();
413         let version = db0.schema_version().unwrap();
414         assert_eq!(0, version);
415     }
416 
417     #[test]
test_schema_version1()418     fn test_schema_version1() {
419         let mut db1 = VmIdDb { conn: Connection::open_in_memory().unwrap() };
420         db1.init_tables_v1().unwrap();
421         let version = db1.schema_version().unwrap();
422         assert_eq!(1, version);
423     }
424 
425     #[test]
test_schema_upgrade_v0_v1()426     fn test_schema_upgrade_v0_v1() {
427         let mut db = new_test_db_version(0);
428         let version = db.schema_version().unwrap();
429         assert_eq!(0, version);
430 
431         // Manually insert a row before upgrade.
432         db.conn
433             .execute(
434                 "REPLACE INTO main.vmids (vm_id, user_id, app_id) VALUES (?1, ?2, ?3);",
435                 params![&VM_ID1, &USER1, APP_A],
436             )
437             .unwrap();
438 
439         db.upgrade_tables_v0_v1().unwrap();
440         let version = db.schema_version().unwrap();
441         assert_eq!(1, version);
442 
443         assert_eq!(vec![VM_ID1], db.vm_ids_for_user(USER1).unwrap());
444         show_contents(&db);
445     }
446 
447     #[test]
test_corrupt_database_file()448     fn test_corrupt_database_file() {
449         let db_dir = tempfile::Builder::new().prefix("vmdb-test-").tempdir().unwrap();
450         let mut db_path = db_dir.path().to_owned();
451         db_path.push(DB_FILENAME);
452         {
453             let mut file = std::fs::File::create(db_path).unwrap();
454             let _ = file.write_all(b"This is not an SQLite file!");
455         }
456 
457         // Non-DB file should be wiped and start over.
458         let (mut db, created) =
459             VmIdDb::new(&db_dir.path().to_string_lossy()).expect("failed to replace bogus DB");
460         assert!(created);
461         db.add_vm_id(&VM_ID1, USER1, APP_A).unwrap();
462         assert_eq!(vec![VM_ID1], db.vm_ids_for_user(USER1).unwrap());
463     }
464 
465     #[test]
test_non_upgradable_database_file()466     fn test_non_upgradable_database_file() {
467         let db_dir = tempfile::Builder::new().prefix("vmdb-test-").tempdir().unwrap();
468         let mut db_path = db_dir.path().to_owned();
469         db_path.push(DB_FILENAME);
470         {
471             // Create an unrelated database that happens to apparently have a schema version of 0.
472             let (db, created) = VmIdDb::new(&db_dir.path().to_string_lossy()).unwrap();
473             assert!(created);
474             db.conn.execute("DROP TABLE main.vmids", ()).unwrap();
475             db.conn.execute("PRAGMA main.user_version = 0;", ()).unwrap();
476         }
477 
478         // Should fail to open a database because the upgrade fails.
479         let result = VmIdDb::new(&db_dir.path().to_string_lossy());
480         assert!(result.is_err());
481     }
482 
483     #[test]
test_database_from_the_future()484     fn test_database_from_the_future() {
485         let db_dir = tempfile::Builder::new().prefix("vmdb-test-").tempdir().unwrap();
486         {
487             let (mut db, created) = VmIdDb::new(&db_dir.path().to_string_lossy()).unwrap();
488             assert!(created);
489             db.add_vm_id(&VM_ID1, USER1, APP_A).unwrap();
490             // Make the database look like it's from a future version.
491             db.conn.execute("PRAGMA main.user_version = 99;", ()).unwrap();
492         }
493         // Should fail to open a database from the future.
494         let result = VmIdDb::new(&db_dir.path().to_string_lossy());
495         assert!(result.is_err());
496     }
497 
498     #[test]
test_add_remove()499     fn test_add_remove() {
500         let mut db = new_test_db();
501         db.add_vm_id(&VM_ID1, USER1, APP_A).unwrap();
502         db.add_vm_id(&VM_ID2, USER1, APP_A).unwrap();
503         db.add_vm_id(&VM_ID3, USER1, APP_A).unwrap();
504         db.add_vm_id(&VM_ID4, USER2, APP_B).unwrap();
505         db.add_vm_id(&VM_ID5, USER3, APP_A).unwrap();
506         db.add_vm_id(&VM_ID5, USER3, APP_C).unwrap(); // Overwrites APP_A
507 
508         assert_eq!(
509             vec![(USER1, APP_A), (USER2, APP_B), (USER3, APP_C)],
510             db.get_all_owners().unwrap()
511         );
512 
513         let empty: Vec<VmId> = Vec::new();
514 
515         assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_user(USER1).unwrap());
516         assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_app(USER1, APP_A).unwrap());
517         assert_eq!(3, db.count_vm_ids_for_app(USER1, APP_A).unwrap());
518         assert_eq!(vec![VM_ID4], db.vm_ids_for_app(USER2, APP_B).unwrap());
519         assert_eq!(1, db.count_vm_ids_for_app(USER2, APP_B).unwrap());
520         assert_eq!(vec![VM_ID5], db.vm_ids_for_user(USER3).unwrap());
521         assert_eq!(empty, db.vm_ids_for_user(USER_UNKNOWN).unwrap());
522         assert_eq!(empty, db.vm_ids_for_app(USER1, APP_UNKNOWN).unwrap());
523         assert_eq!(0, db.count_vm_ids_for_app(USER1, APP_UNKNOWN).unwrap());
524 
525         assert!(db.is_vm_id_for_app(&VM_ID1, USER1 as u32, APP_A as u32).unwrap());
526         assert!(!db.is_vm_id_for_app(&VM_ID1, USER2 as u32, APP_A as u32).unwrap());
527         assert!(!db.is_vm_id_for_app(&VM_ID1, USER1 as u32, APP_B as u32).unwrap());
528         assert!(!db.is_vm_id_for_app(&VM_ID_UNKNOWN, USER1 as u32, APP_A as u32).unwrap());
529         assert!(!db.is_vm_id_for_app(&VM_ID5, USER3 as u32, APP_A as u32).unwrap());
530         assert!(db.is_vm_id_for_app(&VM_ID5, USER3 as u32, APP_C as u32).unwrap());
531 
532         db.delete_vm_ids(&[VM_ID2, VM_ID3]).unwrap();
533 
534         assert_eq!(vec![VM_ID1], db.vm_ids_for_user(USER1).unwrap());
535         assert_eq!(vec![VM_ID1], db.vm_ids_for_app(USER1, APP_A).unwrap());
536         assert_eq!(1, db.count_vm_ids_for_app(USER1, APP_A).unwrap());
537 
538         // OK to delete things that don't exist.
539         db.delete_vm_ids(&[VM_ID2, VM_ID3]).unwrap();
540 
541         assert_eq!(vec![VM_ID1], db.vm_ids_for_user(USER1).unwrap());
542         assert_eq!(vec![VM_ID1], db.vm_ids_for_app(USER1, APP_A).unwrap());
543         assert_eq!(1, db.count_vm_ids_for_app(USER1, APP_A).unwrap());
544 
545         db.add_vm_id(&VM_ID2, USER1, APP_A).unwrap();
546         db.add_vm_id(&VM_ID3, USER1, APP_A).unwrap();
547 
548         assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_user(USER1).unwrap());
549         assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_app(USER1, APP_A).unwrap());
550         assert_eq!(3, db.count_vm_ids_for_app(USER1, APP_A).unwrap());
551         assert_eq!(vec![VM_ID4], db.vm_ids_for_app(USER2, APP_B).unwrap());
552         assert_eq!(1, db.count_vm_ids_for_app(USER2, APP_B).unwrap());
553         assert_eq!(vec![VM_ID5], db.vm_ids_for_user(USER3).unwrap());
554         assert_eq!(empty, db.vm_ids_for_user(USER_UNKNOWN).unwrap());
555         assert_eq!(empty, db.vm_ids_for_app(USER1, APP_UNKNOWN).unwrap());
556         assert_eq!(0, db.count_vm_ids_for_app(USER1, APP_UNKNOWN).unwrap());
557 
558         assert_eq!(
559             vec![(USER1, APP_A), (USER2, APP_B), (USER3, APP_C)],
560             db.get_all_owners().unwrap()
561         );
562 
563         show_contents(&db);
564     }
565 
566     #[test]
test_invalid_vm_id()567     fn test_invalid_vm_id() {
568         let mut db = new_test_db();
569         db.add_vm_id(&VM_ID3, USER1, APP_A).unwrap();
570         db.add_vm_id(&VM_ID2, USER1, APP_A).unwrap();
571         db.add_vm_id(&VM_ID1, USER1, APP_A).unwrap();
572 
573         // Note that results are returned in `vm_id` order, because the table is `WITHOUT ROWID`.
574         assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_user(USER1).unwrap());
575 
576         // Manually insert a row with a VM ID that's the wrong size.
577         db.conn
578             .execute(
579                 "REPLACE INTO main.vmids (vm_id, user_id, app_id, created) VALUES (?1, ?2, ?3, ?4);",
580                 params![&[99u8; 60], &USER1, APP_A, &db_now()],
581             )
582             .unwrap();
583 
584         // Invalid row is skipped and remainder returned.
585         assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_user(USER1).unwrap());
586         show_contents(&db);
587     }
588 
589     #[test]
test_remove_oldest_with_upgrade()590     fn test_remove_oldest_with_upgrade() {
591         let mut db = new_test_db_version(0);
592         let version = db.schema_version().unwrap();
593         assert_eq!(0, version);
594 
595         let remove_count = 10;
596         let mut want = vec![];
597 
598         // Manually insert rows before upgrade.
599         const V0_COUNT: usize = 5;
600         for idx in 0..V0_COUNT {
601             let mut vm_id = [0u8; 64];
602             vm_id[0..8].copy_from_slice(&(idx as u64).to_be_bytes());
603             if want.len() < remove_count {
604                 want.push(vm_id);
605             }
606             db.conn
607                 .execute(
608                     "REPLACE INTO main.vmids (vm_id, user_id, app_id) VALUES (?1, ?2, ?3);",
609                     params![&vm_id, &USER1, APP_A],
610                 )
611                 .unwrap();
612         }
613 
614         // Now move to v1.
615         db.upgrade_tables_v0_v1().unwrap();
616         let version = db.schema_version().unwrap();
617         assert_eq!(1, version);
618 
619         for idx in V0_COUNT..40 {
620             let mut vm_id = [0u8; 64];
621             vm_id[0..8].copy_from_slice(&(idx as u64).to_be_bytes());
622             if want.len() < remove_count {
623                 want.push(vm_id);
624             }
625             db.add_vm_id(&vm_id, USER1, APP_A).unwrap();
626         }
627         show_contents_for_app(&db, USER1, APP_A, 10);
628         let got = db.oldest_vm_ids_for_app(USER1, APP_A, 10).unwrap();
629         assert_eq!(got, want);
630     }
631 }
632