virtualizationservice/src/maintenance/vmdb.rs

// Copyright 2024, The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Database of VM IDs.

use anyhow::{anyhow, Context, Result};
use log::{debug, error, info, warn};
use rusqlite::{params, params_from_iter, Connection, OpenFlags, Rows};
use std::path::PathBuf;

/// Subdirectory to hold the database.
const DB_DIR: &str = "vmdb";

/// Name of the file that holds the database.
const DB_FILENAME: &str = "vmids.sqlite";

/// Maximum number of host parameters in a single SQL statement.
/// (Default value of `SQLITE_LIMIT_VARIABLE_NUMBER` for <= 3.32.0)
const MAX_VARIABLES: usize = 999;

/// Return the current time as milliseconds since epoch.
fn db_now() -> u64 {
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or(std::time::Duration::ZERO)
        .as_millis();
    now.try_into().unwrap_or(u64::MAX)
}

/// Identifier for a VM and its corresponding secret.
pub type VmId = [u8; 64];

/// Representation of an on-disk database of VM IDs.
pub struct VmIdDb {
    conn: Connection,
}

struct RetryOnFailure(bool);

impl VmIdDb {
    /// Connect to the VM ID database file held in the given directory, creating it if necessary.
    /// The second return value indicates whether a new database file was created.
    ///
    /// This function assumes no other threads/processes are attempting to connect concurrently.
    pub fn new(db_dir: &str) -> Result<(Self, bool)> {
        let mut db_path = PathBuf::from(db_dir);
        db_path.push(DB_DIR);
        if !db_path.exists() {
            std::fs::create_dir(&db_path).context("failed to create {db_path:?}")?;
            info!("created persistent db dir {db_path:?}");
        }
        db_path.push(DB_FILENAME);
        Self::new_at_path(db_path, RetryOnFailure(true))
    }

    fn new_at_path(db_path: PathBuf, retry: RetryOnFailure) -> Result<(Self, bool)> {
        let (flags, created) = if db_path.exists() {
            debug!("connecting to existing database {db_path:?}");
            (
                OpenFlags::SQLITE_OPEN_READ_WRITE
                    | OpenFlags::SQLITE_OPEN_URI
                    | OpenFlags::SQLITE_OPEN_NO_MUTEX,
                false,
            )
        } else {
            info!("creating fresh database {db_path:?}");
            (
                OpenFlags::SQLITE_OPEN_READ_WRITE
                    | OpenFlags::SQLITE_OPEN_CREATE
                    | OpenFlags::SQLITE_OPEN_URI
                    | OpenFlags::SQLITE_OPEN_NO_MUTEX,
                true,
            )
        };
        let mut db = Self {
            conn: Connection::open_with_flags(&db_path, flags)
                .context(format!("failed to open/create DB with {flags:?}"))?,
        };

        if created {
            db.init_tables().context("failed to create tables")?;
        } else {
            // An existing .sqlite file may have an earlier schema.
            match db.schema_version() {
                Err(e) => {
                    // Couldn't determine a schema version, so wipe and try again.
                    error!("failed to determine VM DB schema: {e:?}");
                    if retry.0 {
                        // This is the first attempt, so wipe and retry.
                        error!("resetting database file {db_path:?}");
                        let _ = std::fs::remove_file(&db_path);
                        return Self::new_at_path(db_path, RetryOnFailure(false));
                    } else {
                        // An earlier attempt at wiping/retrying has failed, so give up.
                        return Err(anyhow!("failed to reset database file {db_path:?}"));
                    }
                }
                Ok(0) => db.upgrade_tables_v0_v1().context("failed to upgrade schema v0 -> v1")?,
                Ok(1) => {
                    // Current version, no action needed.
                }
                Ok(version) => {
                    // If the database looks like it's from a future version, leave it alone and
                    // fail to connect to it.
                    error!("database from the future (v{version})");
                    return Err(anyhow!("database from the future (v{version})"));
                }
            }
        }
        Ok((db, created))
    }

    /// Delete the associated database file.
    pub fn delete_db_file(self, db_dir: &str) {
        let mut db_path = PathBuf::from(db_dir);
        db_path.push(DB_DIR);
        db_path.push(DB_FILENAME);

        // Drop the connection before removing the backing file.
        drop(self);
        warn!("removing database file {db_path:?}");
        if let Err(e) = std::fs::remove_file(&db_path) {
            error!("failed to remove database file {db_path:?}: {e:?}");
        }
    }

    fn schema_version(&mut self) -> Result<i32> {
        let version: i32 = self
            .conn
            .query_row("PRAGMA main.user_version", (), |row| row.get(0))
            .context("failed to read pragma")?;
        Ok(version)
    }

    /// Create the database table and indices using the current schema.
    fn init_tables(&mut self) -> Result<()> {
        self.init_tables_v1()
    }

    /// Create the database table and indices using the v1 schema.
    fn init_tables_v1(&mut self) -> Result<()> {
        info!("creating v1 database schema");
        self.conn
            .execute(
                "CREATE TABLE IF NOT EXISTS main.vmids (
                     vm_id BLOB PRIMARY KEY,
                     user_id INTEGER,
                     app_id INTEGER,
                     created INTEGER
                 ) WITHOUT ROWID;",
                (),
            )
            .context("failed to create table")?;
        self.conn
            .execute("CREATE INDEX IF NOT EXISTS main.vmids_user_index ON vmids(user_id);", [])
            .context("Failed to create user index")?;
        self.conn
            .execute(
                "CREATE INDEX IF NOT EXISTS main.vmids_app_index ON vmids(user_id, app_id);",
                [],
            )
            .context("Failed to create app index")?;
        self.conn
            .execute("PRAGMA main.user_version = 1;", ())
            .context("failed to declare version")?;
        Ok(())
    }

    fn upgrade_tables_v0_v1(&mut self) -> Result<()> {
        let _rows = self
            .conn
            .execute("ALTER TABLE main.vmids ADD COLUMN created INTEGER;", ())
            .context("failed to alter table v0->v1")?;
        self.conn
            .execute("PRAGMA main.user_version = 1;", ())
            .context("failed to set schema version")?;
        Ok(())
    }

    /// Create the database table and indices using the v0 schema.
    #[cfg(test)]
    fn init_tables_v0(&mut self) -> Result<()> {
        info!("creating v0 database schema");
        self.conn
            .execute(
                "CREATE TABLE IF NOT EXISTS main.vmids (
                     vm_id BLOB PRIMARY KEY,
                     user_id INTEGER,
                     app_id INTEGER
                 ) WITHOUT ROWID;",
                (),
            )
            .context("failed to create table")?;
        self.conn
            .execute("CREATE INDEX IF NOT EXISTS main.vmids_user_index ON vmids(user_id);", [])
            .context("Failed to create user index")?;
        self.conn
            .execute(
                "CREATE INDEX IF NOT EXISTS main.vmids_app_index ON vmids(user_id, app_id);",
                [],
            )
            .context("Failed to create app index")?;
        Ok(())
    }

    /// Add the given VM ID into the database.
    pub fn add_vm_id(&mut self, vm_id: &VmId, user_id: i32, app_id: i32) -> Result<()> {
        let now = db_now();
        let _rows = self
            .conn
            .execute(
                "REPLACE INTO main.vmids (vm_id, user_id, app_id, created) VALUES (?1, ?2, ?3, ?4);",
                params![vm_id, &user_id, &app_id, &now],
            )
            .context("failed to add VM ID")?;
        Ok(())
    }

    /// Remove the given VM IDs from the database.  The collection of IDs is assumed to be smaller
    /// than the maximum number of SQLite parameters.
    pub fn delete_vm_ids(&mut self, vm_ids: &[VmId]) -> Result<()> {
        assert!(vm_ids.len() < MAX_VARIABLES);
        let mut vars = "?,".repeat(vm_ids.len());
        vars.pop(); // remove trailing comma
        let sql = format!("DELETE FROM main.vmids WHERE vm_id IN ({});", vars);
        let mut stmt = self.conn.prepare(&sql).context("failed to prepare DELETE stmt")?;
        let _rows = stmt.execute(params_from_iter(vm_ids)).context("failed to delete VM IDs")?;
        Ok(())
    }

    /// Return the VM IDs associated with Android user ID `user_id`.
    pub fn vm_ids_for_user(&mut self, user_id: i32) -> Result<Vec<VmId>> {
        let mut stmt = self
            .conn
            .prepare("SELECT vm_id FROM main.vmids WHERE user_id = ?;")
            .context("failed to prepare SELECT stmt")?;
        let rows = stmt.query(params![user_id]).context("query failed")?;
        Self::vm_ids_from_rows(rows)
    }

    /// Return the VM IDs associated with `(user_id, app_id)`.
    pub fn vm_ids_for_app(&mut self, user_id: i32, app_id: i32) -> Result<Vec<VmId>> {
        let mut stmt = self
            .conn
            .prepare("SELECT vm_id FROM main.vmids WHERE user_id = ? AND app_id = ?;")
            .context("failed to prepare SELECT stmt")?;
        let rows = stmt.query(params![user_id, app_id]).context("query failed")?;
        Self::vm_ids_from_rows(rows)
    }

    /// Retrieve a collection of VM IDs from database rows.
    fn vm_ids_from_rows(mut rows: Rows) -> Result<Vec<VmId>> {
        let mut vm_ids: Vec<VmId> = Vec::new();
        while let Some(row) = rows.next().context("failed row unpack")? {
            match row.get(0) {
                Ok(vm_id) => vm_ids.push(vm_id),
                Err(e) => error!("failed to parse row: {e:?}"),
            }
        }

        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
            .conn
            .prepare("SELECT DISTINCT user_id, app_id FROM main.vmids;")
            .context("failed to prepare SELECT stmt")?;
        let mut rows = stmt.query(()).context("query failed")?;
        let mut owners: Vec<(i32, i32)> = Vec::new();
        while let Some(row) = rows.next().context("failed row unpack")? {
            let user_id = match row.get(0) {
                Ok(v) => v,
                Err(e) => {
                    error!("failed to parse row: {e:?}");
                    continue;
                }
            };
            let app_id = match row.get(1) {
                Ok(v) => v,
                Err(e) => {
                    error!("failed to parse row: {e:?}");
                    continue;
                }
            };
            owners.push((user_id, app_id));
        }

        Ok(owners)
    }
}

/// Current schema version.
#[cfg(test)]
const SCHEMA_VERSION: usize = 1;

/// Create a new in-memory database for testing.
#[cfg(test)]
pub fn new_test_db() -> VmIdDb {
    tests::new_test_db_version(SCHEMA_VERSION)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;

    const VM_ID1: VmId = [1u8; 64];
    const VM_ID2: VmId = [2u8; 64];
    const VM_ID3: VmId = [3u8; 64];
    const VM_ID4: VmId = [4u8; 64];
    const VM_ID5: VmId = [5u8; 64];
    const USER1: i32 = 1;
    const USER2: i32 = 2;
    const USER3: i32 = 3;
    const USER_UNKNOWN: i32 = 4;
    const APP_A: i32 = 50;
    const APP_B: i32 = 60;
    const APP_C: i32 = 70;
    const APP_UNKNOWN: i32 = 99;

    pub fn new_test_db_version(version: usize) -> VmIdDb {
        let mut db = VmIdDb { conn: Connection::open_in_memory().unwrap() };
        match version {
            0 => db.init_tables_v0().unwrap(),
            1 => db.init_tables_v1().unwrap(),
            _ => panic!("unexpected version {version}"),
        }
        db
    }

    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:?}");
        }
    }

    #[test]
    fn test_schema_version0() {
        let mut db0 = VmIdDb { conn: Connection::open_in_memory().unwrap() };
        db0.init_tables_v0().unwrap();
        let version = db0.schema_version().unwrap();
        assert_eq!(0, version);
    }

    #[test]
    fn test_schema_version1() {
        let mut db1 = VmIdDb { conn: Connection::open_in_memory().unwrap() };
        db1.init_tables_v1().unwrap();
        let version = db1.schema_version().unwrap();
        assert_eq!(1, version);
    }

    #[test]
    fn test_schema_upgrade_v0_v1() {
        let mut db = new_test_db_version(0);
        let version = db.schema_version().unwrap();
        assert_eq!(0, version);

        // Manually insert a row before upgrade.
        db.conn
            .execute(
                "REPLACE INTO main.vmids (vm_id, user_id, app_id) VALUES (?1, ?2, ?3);",
                params![&VM_ID1, &USER1, APP_A],
            )
            .unwrap();

        db.upgrade_tables_v0_v1().unwrap();
        let version = db.schema_version().unwrap();
        assert_eq!(1, version);

        assert_eq!(vec![VM_ID1], db.vm_ids_for_user(USER1).unwrap());
        show_contents(&db);
    }

    #[test]
    fn test_corrupt_database_file() {
        let db_dir = tempfile::Builder::new().prefix("vmdb-test-").tempdir().unwrap();
        let mut db_path = db_dir.path().to_owned();
        db_path.push(DB_FILENAME);
        {
            let mut file = std::fs::File::create(db_path).unwrap();
            let _ = file.write_all(b"This is not an SQLite file!");
        }

        // Non-DB file should be wiped and start over.
        let (mut db, created) =
            VmIdDb::new(&db_dir.path().to_string_lossy()).expect("failed to replace bogus DB");
        assert!(created);
        db.add_vm_id(&VM_ID1, USER1, APP_A).unwrap();
        assert_eq!(vec![VM_ID1], db.vm_ids_for_user(USER1).unwrap());
    }

    #[test]
    fn test_non_upgradable_database_file() {
        let db_dir = tempfile::Builder::new().prefix("vmdb-test-").tempdir().unwrap();
        let mut db_path = db_dir.path().to_owned();
        db_path.push(DB_FILENAME);
        {
            // Create an unrelated database that happens to apparently have a schema version of 0.
            let (db, created) = VmIdDb::new(&db_dir.path().to_string_lossy()).unwrap();
            assert!(created);
            db.conn.execute("DROP TABLE main.vmids", ()).unwrap();
            db.conn.execute("PRAGMA main.user_version = 0;", ()).unwrap();
        }

        // Should fail to open a database because the upgrade fails.
        let result = VmIdDb::new(&db_dir.path().to_string_lossy());
        assert!(result.is_err());
    }

    #[test]
    fn test_database_from_the_future() {
        let db_dir = tempfile::Builder::new().prefix("vmdb-test-").tempdir().unwrap();
        {
            let (mut db, created) = VmIdDb::new(&db_dir.path().to_string_lossy()).unwrap();
            assert!(created);
            db.add_vm_id(&VM_ID1, USER1, APP_A).unwrap();
            // Make the database look like it's from a future version.
            db.conn.execute("PRAGMA main.user_version = 99;", ()).unwrap();
        }
        // Should fail to open a database from the future.
        let result = VmIdDb::new(&db_dir.path().to_string_lossy());
        assert!(result.is_err());
    }

    #[test]
    fn test_add_remove() {
        let mut db = new_test_db();
        db.add_vm_id(&VM_ID1, USER1, APP_A).unwrap();
        db.add_vm_id(&VM_ID2, USER1, APP_A).unwrap();
        db.add_vm_id(&VM_ID3, USER1, APP_A).unwrap();
        db.add_vm_id(&VM_ID4, USER2, APP_B).unwrap();
        db.add_vm_id(&VM_ID5, USER3, APP_A).unwrap();
        db.add_vm_id(&VM_ID5, USER3, APP_C).unwrap(); // Overwrites APP_A

        assert_eq!(
            vec![(USER1, APP_A), (USER2, APP_B), (USER3, APP_C)],
            db.get_all_owners().unwrap()
        );

        let empty: Vec<VmId> = Vec::new();

        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)],
            db.get_all_owners().unwrap()
        );

        show_contents(&db);
    }

    #[test]
    fn test_invalid_vm_id() {
        let mut db = new_test_db();
        db.add_vm_id(&VM_ID3, USER1, APP_A).unwrap();
        db.add_vm_id(&VM_ID2, USER1, APP_A).unwrap();
        db.add_vm_id(&VM_ID1, USER1, APP_A).unwrap();

        // Note that results are returned in `vm_id` order, because the table is `WITHOUT ROWID`.
        assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_user(USER1).unwrap());

        // Manually insert a row with a VM ID that's the wrong size.
        db.conn
            .execute(
                "REPLACE INTO main.vmids (vm_id, user_id, app_id, created) VALUES (?1, ?2, ?3, ?4);",
                params![&[99u8; 60], &USER1, APP_A, &db_now()],
            )
            .unwrap();

        // Invalid row is skipped and remainder returned.
        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);
    }
}