local-mini-kms/src/db.rs

use rusqlite::{params, Connection};
use rust_util::{opt_result, simple_error, XResult};

pub const DEFAULT_MASTER_KEY_VERIFICATION_KEY: &str = "__master_verification_key";

pub struct Key {
    pub name: String,
    pub encrypted_key: String,
    pub comment: Option<String>,
}

pub struct Keys {
    pub count: usize,
    pub keys: Vec<Key>,
}

pub fn make_value_key_name(name: &str) -> String {
    format!("value:{}", name)
}

pub fn make_data_key_name(name: &str) -> String {
    format!("data_key:{}", name)
}

pub fn open_db(db: &str) -> XResult<Connection> {
    let con = opt_result!(Connection::open(db), "Open sqlite db: {}, failed: {}", db);
    log::debug!("Db auto commit: {}", con.is_autocommit());
    Ok(con)
}

pub fn init_db(conn: &Connection) -> XResult<bool> {
    if let Ok(false) = check_table_keys(conn) {
        repair_table_keys(conn)?;
    }

    Ok(true)
}

fn repair_table_keys(conn: &Connection) -> XResult<()> {
    let field_names = list_table_fields(conn, "keys")?;
    let field_names = field_names.iter().map(|n| n.as_str()).collect::<Vec<_>>();

    if !field_names.contains(&"comment") {
        log::info!("Repair table keys, add column comment");
        let _ = conn.execute("ALTER TABLE keys ADD COLUMN comment TEXT", ())?;
    }

    Ok(())
}

fn list_table_fields(conn: &Connection, table: &str) -> XResult<Vec<String>> {
    let mut stmt_query_fields = conn.prepare(&format!("pragma table_info({})", table))?;
    let mut rows_query_field = stmt_query_fields.query(())?;

    let mut field_names = vec![];
    let mut next_query_field_opt = rows_query_field.next()?;
    while let Some(next_query_field) = next_query_field_opt {
        // cid|name|type|notnull|dflt_value|pk
        // ^   ^    ^    ^       ^          ^
        // |   |    |    |       |          [5] - Is column PK
        // |   |    |    |       [4] Column default value
        // |   |    |    [3] Is column not null
        // |   |    [2] Column type
        // |   [1] Column name
        // [0] Column index
        let field_name: String = next_query_field.get(1)?;
        field_names.push(field_name.to_lowercase());
        next_query_field_opt = rows_query_field.next()?;
    }
    log::trace!("Table {} fields: {:?}", table, field_names);
    Ok(field_names)
}

fn check_table_keys(conn: &Connection) -> XResult<bool> {
    let mut stmt =
        conn.prepare("SELECT name FROM sqlite_master WHERE type='table' AND name='keys'")?;
    let mut rows = stmt.query(())?;
    if rows.next()?.is_none() {
        log::info!("Create table keys");
        create_table_keys(conn)?;
        Ok(true)
    } else {
        Ok(false)
    }
}

fn create_table_keys(conn: &Connection) -> XResult<()> {
    let _ = conn.execute(
        r##"
        CREATE TABLE keys (
            id INTEGER PRIMARY KEY,
            name  TEXT NOT NULL,
            value  TEXT,
            comment TEXT
        )"##,
        (),
    )?;
    Ok(())
}

pub fn insert_key(conn: &Connection, key: &Key) -> XResult<()> {
    let default_comment = "".to_string();
    log::debug!("insert key name={}", &key.name);
    let _ = conn.execute(
        "INSERT INTO keys (name, value, comment) VALUES (?1, ?2, ?3)",
        (
            &key.name,
            &key.encrypted_key,
            key.comment.as_ref().unwrap_or(&default_comment),
        ),
    )?;
    Ok(())
}

pub fn update_key(conn: &Connection, key: &Key) -> XResult<()> {
    log::debug!("update key name={}", &key.name);
    if let Some(comment) = &key.comment {
        let _ = conn.execute(
            "UPDATE keys SET value = ?1, comment = ?2 WHERE name = ?3",
            (&key.encrypted_key, comment, &key.name),
        )?;
    } else {
        let _ = conn.execute(
            "UPDATE keys SET value = ?1 WHERE name = ?2",
            (&key.encrypted_key, &key.name),
        )?;
    }
    Ok(())
}

pub fn list_keys(conn: &Connection, ty: &str, search: &str, limit: usize) -> XResult<Keys> {
    let name = format!("{}:%{}%", ty, search);
    let mut count_stmt = conn.prepare("SELECT count(*) FROM keys WHERE name like ?1")?;
    let mut count_iter = count_stmt.query_map(params![name], |row| {
        let count: usize = row.get(0)?;
        Ok(count)
    })?;
    let count = match count_iter.next() {
        None => 0,
        Some(Ok(count)) => count,
        Some(Err(e)) => return simple_error!("List keys failed: {}", e),
    };
    log::debug!("found {} keys via: {}, limit: {}", count, name, limit);
    let mut keys = vec![];
    if count > 0 {
        let mut list_stmt =
            conn.prepare("SELECT id, name, value, comment FROM keys WHERE name like ?1 LIMIT ?2")?;
        let mut list_iter = list_stmt.query_map(params![name, limit], |row| {
            Ok(Key {
                name: row.get(1)?,
                encrypted_key: row.get(2)?,
                comment: row.get(3)?,
            })
        })?;

        loop {
            match list_iter.next() {
                None => {
                    break;
                }
                Some(Ok(r)) => {
                    log::debug!("found key name={}", r.name);
                    keys.push(r);
                }
                Some(Err(e)) => return simple_error!("List keys failed: {}", e),
            }
        }
    }

    Ok(Keys { count, keys })
}

pub fn find_key(conn: &Connection, name: &str) -> XResult<Option<Key>> {
    let mut stmt = conn.prepare("SELECT id, name, value, comment FROM keys WHERE name = ?1")?;
    let mut key_iter = stmt.query_map(params![name], |row| {
        Ok(Key {
            name: row.get(1)?,
            encrypted_key: row.get(2)?,
            comment: row.get(3)?,
        })
    })?;
    match key_iter.next() {
        None => {
            log::debug!("key name={} not exists", name);
            Ok(None)
        }
        Some(Ok(r)) => {
            log::debug!("found key name={}", name);
            Ok(Some(r))
        }
        Some(Err(e)) => simple_error!("Find key failed: {}", e),
    }
}