use internal address HashMap registry for the TOFU validation

src/profile/tofu.rs

@@ -7,16 +7,15 @@ use database::Database;
 use gtk::{gio::TlsCertificate, glib::Uri};
 use r2d2::Pool;
 use r2d2_sqlite::SqliteConnectionManager;
-use sourceview::prelude::TlsCertificateExt;
 use sqlite::Transaction;
-use std::cell::RefCell;
+use std::{cell::RefCell, collections::HashMap};
 
 /// TOFU wrapper for the Gemini protocol
 ///
 /// https://geminiprotocol.net/docs/protocol-specification.gmi#tls-server-certificate-validation
 pub struct Tofu {
     database: Database,
-    memory: RefCell<Vec<Certificate>>,
+    memory: RefCell<HashMap<String, Certificate>>,
 }
 
 impl Tofu {
@@ -26,13 +25,17 @@ impl Tofu {
         let database = Database::init(database_pool, profile_id);
 
         let records = database.records()?;
-        let memory = RefCell::new(Vec::with_capacity(records.len()));
+        let memory = RefCell::new(HashMap::with_capacity(records.len()));
 
         {
             // build in-memory index...
             let mut m = memory.borrow_mut();
             for r in records {
-                m.push(Certificate::from_db(Some(r.id), &r.pem, r.time)?)
+                if m.insert(r.address, Certificate::from_db(Some(r.id), &r.pem, r.time)?)
+                    .is_some()
+                {
+                    panic!() // expect unique address
+                }
             }
         }
 
@@ -41,47 +44,37 @@ impl Tofu {
 
     // Actions
 
-    pub fn add(&self, tls_certificate: TlsCertificate) -> Result<()> {
-        self.memory
-            .borrow_mut()
-            .push(Certificate::from_tls_certificate(tls_certificate)?);
-        Ok(())
+    pub fn add(
+        &self,
+        uri: &Uri,
+        default_port: i32,
+        tls_certificate: TlsCertificate,
+    ) -> Result<bool> {
+        match address(uri, default_port) {
+            Some(k) => Ok(self
+                .memory
+                .borrow_mut()
+                .insert(k, Certificate::from_tls_certificate(tls_certificate)?)
+                .is_none()),
+            None => Ok(false),
+        }
     }
 
-    pub fn server_certificates(&self, uri: &Uri) -> Option<Vec<TlsCertificate>> {
-        fn f(subject_name: &str) -> String {
-            subject_name
-                .trim_start_matches("CN=")
-                .trim_start_matches('*')
-                .trim_matches('.')
-                .to_lowercase()
-        }
-        if let Some(h) = uri.host() {
-            let k = f(&h);
-            let m = self.memory.borrow();
-            let b: Vec<TlsCertificate> = m
-                .iter()
-                .filter_map(|certificate| {
-                    let tls_certificate = certificate.tls_certificate();
-                    if k.ends_with(&f(&tls_certificate.subject_name().unwrap())) {
-                        Some(tls_certificate.clone())
-                    } else {
-                        None
-                    }
-                })
-                .collect();
-            if !b.is_empty() {
-                return Some(b);
-            }
-        }
-        None
+    pub fn server_certificate(&self, uri: &Uri, default_port: i32) -> Option<TlsCertificate> {
+        address(uri, default_port).and_then(|k| {
+            self.memory
+                .borrow()
+                .get(&k)
+                .map(|c| c.tls_certificate().clone())
+        })
     }
 
     /// Save in-memory index to the permanent database (on app close)
     pub fn save(&self) -> Result<()> {
-        for c in self.memory.borrow().iter() {
-            if c.id().is_none() {
-                self.database.add(c.time(), &c.pem())?;
+        for (address, certificate) in self.memory.borrow_mut().drain() {
+            if certificate.id().is_none() {
+                self.database
+                    .add(address, certificate.time(), &certificate.pem())?;
             }
         }
         Ok(())
@@ -100,3 +93,18 @@ pub fn migrate(tx: &Transaction) -> Result<()> {
     // Success
     Ok(())
 }
+
+fn address(uri: &Uri, default_port: i32) -> Option<String> {
+    uri.host().map(|host| {
+        let port = uri.port();
+        format!(
+            "{}:{}",
+            host,
+            if port.is_positive() {
+                port
+            } else {
+                default_port
+            }
+        )
+    })
+}

src/profile/tofu/database.rs

@@ -34,10 +34,10 @@ impl Database {
 
     /// Create new record in database
     /// * return last insert ID on success
-    pub fn add(&self, time: &DateTime, pem: &str) -> Result<i64> {
+    pub fn add(&self, address: String, time: &DateTime, pem: &str) -> Result<i64> {
         let mut connection = self.pool.get()?;
         let tx = connection.transaction()?;
-        let id = insert(&tx, self.profile_id, time, pem)?;
+        let id = insert(&tx, self.profile_id, address, time, pem)?;
         tx.commit()?;
         Ok(id)
     }
@@ -52,37 +52,49 @@ pub fn init(tx: &Transaction) -> Result<usize> {
             `id`         INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
             `profile_id` INTEGER NOT NULL,
             `time`       INTEGER NOT NULL,
+            `address`    VARCHAR(255) NOT NULL,
             `pem`        TEXT NOT NULL,
 
-            FOREIGN KEY (`profile_id`) REFERENCES `profile` (`id`)
+            FOREIGN KEY (`profile_id`) REFERENCES `profile` (`id`),
+            UNIQUE (`address`)
         )",
         [],
     )?)
 }
 
-pub fn insert(tx: &Transaction, profile_id: i64, time: &DateTime, pem: &str) -> Result<i64> {
+pub fn insert(
+    tx: &Transaction,
+    profile_id: i64,
+    address: String,
+    time: &DateTime,
+    pem: &str,
+) -> Result<i64> {
     tx.execute(
         "INSERT INTO `profile_tofu` (
             `profile_id`,
             `time`,
+            `address`,
             `pem`
-        ) VALUES (?, ?, ?)",
-        (profile_id, time.to_unix(), pem),
+        ) VALUES (?, ?, ?, ?) ON CONFLICT (`address`)
+                              DO UPDATE SET `time` = `excluded`.`time`,
+                                            `pem`  = `excluded`.`pem`",
+        (profile_id, time.to_unix(), address, pem),
     )?;
     Ok(tx.last_insert_rowid())
 }
 
 pub fn select(tx: &Transaction, profile_id: i64) -> Result<Vec<Row>> {
     let mut stmt = tx.prepare(
-        "SELECT `id`, `profile_id`, `time`, `pem` FROM `profile_tofu` WHERE `profile_id` = ?",
+        "SELECT `id`, `profile_id`, `address`, `time`, `pem` FROM `profile_tofu` WHERE `profile_id` = ?",
     )?;
 
     let result = stmt.query_map([profile_id], |row| {
         Ok(Row {
             id: row.get(0)?,
             //profile_id: row.get(1)?,
-            time: DateTime::from_unix_local(row.get(2)?).unwrap(),
-            pem: row.get(3)?,
+            address: row.get(2)?,
+            time: DateTime::from_unix_local(row.get(3)?).unwrap(),
+            pem: row.get(4)?,
         })
     })?;
 

src/profile/tofu/database/row.rs

@@ -1,4 +1,5 @@
 pub struct Row {
+    pub address: String,
     pub id: i64,
     pub pem: String,
     pub time: gtk::glib::DateTime,