1 files changed, 41 insertions, 5 deletions

diff --git a/src/bin/client/database.rs b/src/bin/client/database.rs
index a274f31..55051fb 100644
--- a/src/bin/client/database.rs
+++ b/src/bin/client/database.rs
@@ -1,15 +1,17 @@
 use rusqlite::Connection;
 
-pub fn init_db() -> Connection {
-    let conn = Connection::open("client.db").unwrap();
+pub fn init_db(conn: &Connection) {
     conn.execute(
         "CREATE TABLE IF NOT EXISTS hosts (
             host TEXT PRIMARY KEY,
-            last_used DATETIME DEFAULT CURRENT_TIMESTAMP
+            last_used DATETIME DEFAULT CURRENT_TIMESTAMP,
+            fingerprint TEXT
         )",
         [],
     )
     .unwrap();
+    // Migration for databases created before TOFU trust.
+    let _ = conn.execute("ALTER TABLE hosts ADD COLUMN fingerprint TEXT", []);
     conn.execute(
         "CREATE TABLE IF NOT EXISTS credentials (
             host TEXT PRIMARY KEY,
@@ -19,7 +21,6 @@ pub fn init_db() -> Connection {
         [],
     )
     .unwrap();
-    conn
 }
 
 pub fn get_last_host(conn: &Connection) -> Option<String> {
@@ -36,7 +37,8 @@ pub fn get_last_host(conn: &Connection) -> Option<String> {
 
 pub fn save_host(conn: &Connection, host: &str) {
     conn.execute(
-        "INSERT OR REPLACE INTO hosts (host, last_used) VALUES (?1, CURRENT_TIMESTAMP)",
+        "INSERT INTO hosts (host) VALUES (?1)
+         ON CONFLICT(host) DO UPDATE SET last_used = CURRENT_TIMESTAMP",
         [host],
     )
     .unwrap();
@@ -68,3 +70,37 @@ pub fn get_credentials(conn: &Connection, host: &str) -> Option<(String, String)
         None
     }
 }
+
+/// SHA-256 SPKI fingerprint of a host's certificate, if trusted.
+pub fn get_fingerprint(conn: &Connection, host: &str) -> Option<String> {
+    let mut stmt = conn
+        .prepare("SELECT fingerprint FROM hosts WHERE host = ?1")
+        .ok()?;
+    let mut rows = stmt.query([host]).ok()?;
+    match rows.next() {
+        Ok(Some(row)) => row.get(0).ok(),
+        _ => None,
+    }
+}
+
+pub fn save_fingerprint(
+    conn: &Connection,
+    host: &str,
+    fingerprint: &str,
+) -> Result<(), rusqlite::Error> {
+    // Hosts may not exist yet (fingerprint can be saved before first
+    // successful connect completes save_host).
+    conn.execute("INSERT OR IGNORE INTO hosts (host) VALUES (?1)", [host])?;
+    conn.execute(
+        "UPDATE hosts SET fingerprint = ?2 WHERE host = ?1",
+        (host, fingerprint),
+    )?;
+    Ok(())
+}
+
+/// Remove all trust and credentials for a host ("forget this host").
+pub fn forget_host(conn: &Connection, host: &str) -> Result<(), rusqlite::Error> {
+    conn.execute("DELETE FROM hosts WHERE host = ?1", [host])?;
+    conn.execute("DELETE FROM credentials WHERE host = ?1", [host])?;
+    Ok(())
+}