diff --git a/Cargo.lock b/Cargo.lock index fb0ab95..792d55d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -997,7 +997,7 @@ checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89" [[package]] name = "local-mini-kms" -version = "1.0.4" +version = "1.0.5" dependencies = [ "aes-gcm-stream", "aes-kw", diff --git a/Cargo.toml b/Cargo.toml index 32aff43..32081de 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "local-mini-kms" -version = "1.0.4" +version = "1.0.5" edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html diff --git a/src/db.rs b/src/db.rs index 081595e..4db43c2 100644 --- a/src/db.rs +++ b/src/db.rs @@ -24,24 +24,73 @@ pub fn open_db(db: &str) -> XResult { } pub fn init_db(conn: &Connection) -> XResult { - 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_some() { - log::info!("Table keys exists, skip init"); - return Ok(false); + if let Ok(false) = check_table_keys(conn) { + repair_table_keys(conn)?; } - let _ = conn.execute(r##" + 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::>(); + + 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> { + 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 { + 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 - ) - "##, ())?; - log::info!("Table keys created"); - Ok(true) + )"##, + (), + )?; + Ok(()) } pub fn insert_key(conn: &Connection, key: &Key) -> XResult<()> { @@ -49,7 +98,11 @@ pub fn insert_key(conn: &Connection, key: &Key) -> XResult<()> { 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)), + ( + &key.name, + &key.encrypted_key, + key.comment.as_ref().unwrap_or(&default_comment), + ), )?; Ok(()) } @@ -90,4 +143,4 @@ pub fn find_key(conn: &Connection, name: &str) -> XResult> { } Some(Err(e)) => simple_error!("Find key failed: {}", e), } -} \ No newline at end of file +}