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