geth/crates/geth-store/src/lib.rs

use rusqlite::{Connection, params};
use std::path::Path;

#[derive(Debug, thiserror::Error)]
pub enum StoreError {
    #[error("sqlite error: {0}")]
    Sqlite(#[from] rusqlite::Error),
    #[error("json error: {0}")]
    Json(#[from] serde_json::Error),
}

pub struct Store {
    conn: Connection,
}

impl Store {
    pub fn open(path: &Path) -> Result<Self, StoreError> {
        let conn = Connection::open(path)?;
        let store = Self { conn };
        store.migrate()?;
        Ok(store)
    }

    pub fn open_memory() -> Result<Self, StoreError> {
        let conn = Connection::open_in_memory()?;
        let store = Self { conn };
        store.migrate()?;
        Ok(store)
    }

    pub fn migrate(&self) -> Result<(), StoreError> {
        self.conn.execute_batch(
            r#"
            PRAGMA foreign_keys = ON;
            CREATE TABLE IF NOT EXISTS meta (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL
            );
            CREATE TABLE IF NOT EXISTS agents (
                agent_id TEXT PRIMARY KEY,
                public_key TEXT NOT NULL,
                created_at_ms INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS nodes (
                node_id TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                agent_id TEXT,
                created_at_ms INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS node_endpoints (
                endpoint_id TEXT PRIMARY KEY,
                node_id TEXT NOT NULL,
                agent_id TEXT NOT NULL,
                transport TEXT NOT NULL,
                created_at_ms INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS resources (
                resource_id TEXT PRIMARY KEY,
                kind TEXT NOT NULL,
                name TEXT NOT NULL,
                authority_ref TEXT NOT NULL,
                local_role TEXT NOT NULL,
                replication_policy TEXT NOT NULL,
                retention_policy TEXT NOT NULL,
                status TEXT NOT NULL,
                created_at_ms INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS keychain_ops (
                op_id TEXT PRIMARY KEY,
                op_json TEXT NOT NULL,
                created_at_ms INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS auth_ops (
                op_id TEXT PRIMARY KEY,
                resource_id TEXT NOT NULL,
                op_json TEXT NOT NULL,
                created_at_ms INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS grants (
                grant_id TEXT PRIMARY KEY,
                resource_id TEXT NOT NULL,
                principal_id TEXT NOT NULL,
                capability TEXT NOT NULL,
                revoked INTEGER NOT NULL DEFAULT 0
            );
            CREATE TABLE IF NOT EXISTS resource_secrets (
                secret_id TEXT PRIMARY KEY,
                resource_id TEXT NOT NULL,
                epoch INTEGER NOT NULL,
                status TEXT NOT NULL,
                created_at_ms INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS cas_objects (
                hash TEXT PRIMARY KEY,
                size_bytes INTEGER NOT NULL,
                path TEXT NOT NULL,
                created_at_ms INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS cas_pins (
                hash TEXT PRIMARY KEY,
                pinned_at_ms INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS kv_stores (
                kv_id TEXT PRIMARY KEY,
                resource_id TEXT NOT NULL,
                name TEXT NOT NULL
            );
            CREATE TABLE IF NOT EXISTS db_resources (
                db_id TEXT PRIMARY KEY,
                resource_id TEXT NOT NULL,
                path TEXT NOT NULL
            );
            CREATE TABLE IF NOT EXISTS document_resources (
                document_id TEXT PRIMARY KEY,
                resource_id TEXT NOT NULL,
                name TEXT NOT NULL
            );
            CREATE TABLE IF NOT EXISTS peer_cards (
                peer_id TEXT PRIMARY KEY,
                card_json TEXT NOT NULL,
                updated_at_ms INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS module_state (
                module TEXT PRIMARY KEY,
                state_json TEXT NOT NULL,
                updated_at_ms INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS ssh_cert_requests (
                request_id TEXT PRIMARY KEY,
                requester_node TEXT NOT NULL,
                public_key TEXT NOT NULL,
                public_key_fingerprint TEXT NOT NULL,
                cert_kind TEXT NOT NULL,
                principals_json TEXT NOT NULL,
                requested_validity TEXT,
                renewal_of TEXT,
                reason TEXT,
                status TEXT NOT NULL,
                created_at_ms INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS ssh_certificates (
                cert_id TEXT PRIMARY KEY,
                request_id TEXT NOT NULL,
                certificate TEXT NOT NULL,
                certificate_fingerprint TEXT NOT NULL,
                imported_at_ms INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS ssh_revocations (
                revocation_id TEXT PRIMARY KEY,
                kind TEXT NOT NULL,
                target TEXT NOT NULL,
                reason TEXT,
                created_at_ms INTEGER NOT NULL,
                published INTEGER NOT NULL DEFAULT 0
            );
            INSERT OR IGNORE INTO meta(key, value) VALUES ('schema_version', '1');
            "#,
        )?;
        Ok(())
    }

    pub fn upsert_agent(&self, agent_id: &str, public_key: &str) -> Result<(), StoreError> {
        self.conn.execute(
            "INSERT OR IGNORE INTO agents(agent_id, public_key, created_at_ms) VALUES (?1, ?2, ?3)",
            params![agent_id, public_key, now_ms()],
        )?;
        Ok(())
    }

    pub fn upsert_node(&self, node_id: &str, name: &str, agent_id: &str) -> Result<(), StoreError> {
        self.conn.execute(
            "INSERT OR IGNORE INTO nodes(node_id, name, agent_id, created_at_ms) VALUES (?1, ?2, ?3, ?4)",
            params![node_id, name, agent_id, now_ms()],
        )?;
        Ok(())
    }

    pub fn upsert_node_endpoint(
        &self,
        endpoint_id: &str,
        node_id: &str,
        agent_id: &str,
        transport: &str,
    ) -> Result<(), StoreError> {
        self.conn.execute(
            "INSERT OR IGNORE INTO node_endpoints(endpoint_id, node_id, agent_id, transport, created_at_ms) VALUES (?1, ?2, ?3, ?4, ?5)",
            params![endpoint_id, node_id, agent_id, transport, now_ms()],
        )?;
        Ok(())
    }

    pub fn list_resources(&self) -> Result<Vec<StoredResource>, StoreError> {
        let mut stmt = self.conn.prepare(
            "SELECT resource_id, kind, name, status FROM resources ORDER BY kind, name, resource_id",
        )?;
        let rows = stmt.query_map([], |row| {
            Ok(StoredResource {
                resource_id: row.get(0)?,
                kind: row.get(1)?,
                name: row.get(2)?,
                status: row.get(3)?,
            })
        })?;
        rows.collect::<Result<Vec<_>, _>>()
            .map_err(StoreError::from)
    }

    pub fn insert_resource(&self, resource: &StoredResource) -> Result<(), StoreError> {
        self.conn.execute(
            r#"INSERT OR IGNORE INTO resources(
                resource_id, kind, name, authority_ref, local_role, replication_policy,
                retention_policy, status, created_at_ms
            ) VALUES (?1, ?2, ?3, 'local', 'owner', 'local-only', 'keep', ?4, ?5)"#,
            params![
                resource.resource_id,
                resource.kind,
                resource.name,
                resource.status,
                now_ms()
            ],
        )?;
        Ok(())
    }

    pub fn record_cas_object(
        &self,
        hash: &str,
        size_bytes: u64,
        path: &str,
    ) -> Result<(), StoreError> {
        self.conn.execute(
            "INSERT OR REPLACE INTO cas_objects(hash, size_bytes, path, created_at_ms) VALUES (?1, ?2, ?3, ?4)",
            params![hash, size_bytes as i64, path, now_ms()],
        )?;
        Ok(())
    }

    pub fn list_cas_objects(&self) -> Result<Vec<CasObject>, StoreError> {
        let mut stmt = self.conn.prepare(
            "SELECT hash, size_bytes, path FROM cas_objects ORDER BY created_at_ms, hash",
        )?;
        let rows = stmt.query_map([], |row| {
            Ok(CasObject {
                hash: row.get(0)?,
                size_bytes: row.get::<_, i64>(1)? as u64,
                path: row.get(2)?,
            })
        })?;
        rows.collect::<Result<Vec<_>, _>>()
            .map_err(StoreError::from)
    }

    pub fn insert_ssh_cert_request(
        &self,
        request: &StoredSshCertRequest,
    ) -> Result<(), StoreError> {
        self.conn.execute(
            r#"INSERT OR REPLACE INTO ssh_cert_requests(
                request_id, requester_node, public_key, public_key_fingerprint, cert_kind,
                principals_json, requested_validity, renewal_of, reason, status, created_at_ms
            ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)"#,
            params![
                request.request_id,
                request.requester_node,
                request.public_key,
                request.public_key_fingerprint,
                request.cert_kind,
                serde_json::to_string(&request.principals)?,
                request.requested_validity,
                request.renewal_of,
                request.reason,
                request.status,
                request.created_at_ms
            ],
        )?;
        Ok(())
    }

    pub fn get_ssh_cert_request(
        &self,
        request_id: &str,
    ) -> Result<Option<StoredSshCertRequest>, StoreError> {
        let mut stmt = self.conn.prepare(
            r#"SELECT request_id, requester_node, public_key, public_key_fingerprint, cert_kind,
               principals_json, requested_validity, renewal_of, reason, status, created_at_ms
               FROM ssh_cert_requests WHERE request_id = ?1"#,
        )?;
        let mut rows = stmt.query(params![request_id])?;
        if let Some(row) = rows.next()? {
            Ok(Some(stored_ssh_cert_request_from_row(row)?))
        } else {
            Ok(None)
        }
    }

    pub fn update_ssh_cert_request_status(
        &self,
        request_id: &str,
        status: &str,
    ) -> Result<(), StoreError> {
        self.conn.execute(
            "UPDATE ssh_cert_requests SET status = ?2 WHERE request_id = ?1",
            params![request_id, status],
        )?;
        Ok(())
    }

    pub fn list_ssh_cert_requests(&self) -> Result<Vec<StoredSshCertRequest>, StoreError> {
        let mut stmt = self.conn.prepare(
            r#"SELECT request_id, requester_node, public_key, public_key_fingerprint, cert_kind,
               principals_json, requested_validity, renewal_of, reason, status, created_at_ms
               FROM ssh_cert_requests ORDER BY created_at_ms, request_id"#,
        )?;
        let rows = stmt.query_map([], stored_ssh_cert_request_from_row)?;
        rows.collect::<Result<Vec<_>, _>>()
            .map_err(StoreError::from)
    }

    pub fn insert_ssh_certificate(
        &self,
        certificate: &StoredSshCertificate,
    ) -> Result<(), StoreError> {
        self.conn.execute(
            r#"INSERT OR REPLACE INTO ssh_certificates(
                cert_id, request_id, certificate, certificate_fingerprint, imported_at_ms
            ) VALUES (?1, ?2, ?3, ?4, ?5)"#,
            params![
                certificate.cert_id,
                certificate.request_id,
                certificate.certificate,
                certificate.certificate_fingerprint,
                certificate.imported_at_ms
            ],
        )?;
        Ok(())
    }

    pub fn list_ssh_certificates(&self) -> Result<Vec<StoredSshCertificate>, StoreError> {
        let mut stmt = self.conn.prepare(
            r#"SELECT cert_id, request_id, certificate, certificate_fingerprint, imported_at_ms
               FROM ssh_certificates ORDER BY imported_at_ms, cert_id"#,
        )?;
        let rows = stmt.query_map([], |row| {
            Ok(StoredSshCertificate {
                cert_id: row.get(0)?,
                request_id: row.get(1)?,
                certificate: row.get(2)?,
                certificate_fingerprint: row.get(3)?,
                imported_at_ms: row.get(4)?,
            })
        })?;
        rows.collect::<Result<Vec<_>, _>>()
            .map_err(StoreError::from)
    }

    pub fn insert_ssh_revocation(
        &self,
        revocation: &StoredSshRevocation,
    ) -> Result<(), StoreError> {
        self.conn.execute(
            r#"INSERT OR REPLACE INTO ssh_revocations(
                revocation_id, kind, target, reason, created_at_ms, published
            ) VALUES (?1, ?2, ?3, ?4, ?5, ?6)"#,
            params![
                revocation.revocation_id,
                revocation.kind,
                revocation.target,
                revocation.reason,
                revocation.created_at_ms,
                if revocation.published { 1_i64 } else { 0_i64 }
            ],
        )?;
        Ok(())
    }

    pub fn list_ssh_revocations(&self) -> Result<Vec<StoredSshRevocation>, StoreError> {
        let mut stmt = self.conn.prepare(
            r#"SELECT revocation_id, kind, target, reason, created_at_ms, published
               FROM ssh_revocations ORDER BY created_at_ms, revocation_id"#,
        )?;
        let rows = stmt.query_map([], |row| {
            Ok(StoredSshRevocation {
                revocation_id: row.get(0)?,
                kind: row.get(1)?,
                target: row.get(2)?,
                reason: row.get(3)?,
                created_at_ms: row.get(4)?,
                published: row.get::<_, i64>(5)? != 0,
            })
        })?;
        rows.collect::<Result<Vec<_>, _>>()
            .map_err(StoreError::from)
    }
}

fn stored_ssh_cert_request_from_row(
    row: &rusqlite::Row<'_>,
) -> Result<StoredSshCertRequest, rusqlite::Error> {
    let principals_json: String = row.get(5)?;
    let principals = serde_json::from_str(&principals_json).map_err(|error| {
        rusqlite::Error::FromSqlConversionFailure(5, rusqlite::types::Type::Text, Box::new(error))
    })?;
    Ok(StoredSshCertRequest {
        request_id: row.get(0)?,
        requester_node: row.get(1)?,
        public_key: row.get(2)?,
        public_key_fingerprint: row.get(3)?,
        cert_kind: row.get(4)?,
        principals,
        requested_validity: row.get(6)?,
        renewal_of: row.get(7)?,
        reason: row.get(8)?,
        status: row.get(9)?,
        created_at_ms: row.get(10)?,
    })
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StoredResource {
    pub resource_id: String,
    pub kind: String,
    pub name: String,
    pub status: String,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CasObject {
    pub hash: String,
    pub size_bytes: u64,
    pub path: String,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StoredSshCertRequest {
    pub request_id: String,
    pub requester_node: String,
    pub public_key: String,
    pub public_key_fingerprint: String,
    pub cert_kind: String,
    pub principals: Vec<String>,
    pub requested_validity: Option<String>,
    pub renewal_of: Option<String>,
    pub reason: Option<String>,
    pub status: String,
    pub created_at_ms: i64,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StoredSshCertificate {
    pub cert_id: String,
    pub request_id: String,
    pub certificate: String,
    pub certificate_fingerprint: String,
    pub imported_at_ms: i64,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StoredSshRevocation {
    pub revocation_id: String,
    pub kind: String,
    pub target: String,
    pub reason: Option<String>,
    pub created_at_ms: i64,
    pub published: bool,
}

#[must_use]
pub fn now_ms() -> i64 {
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default();
    i64::try_from(now.as_millis()).unwrap_or(i64::MAX)
}

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

    #[test]
    fn migrations_are_idempotent() {
        let store = Store::open_memory().expect("open");
        store.migrate().expect("migrate again");
        store.migrate().expect("migrate third time");
        let resources = store.list_resources().expect("resources");
        assert!(resources.is_empty());
    }

    #[test]
    fn ssh_cert_request_and_revocation_roundtrip() {
        let store = Store::open_memory().expect("open");
        let request = StoredSshCertRequest {
            request_id: "ssh-cert-request:1".to_owned(),
            requester_node: "node:laptop".to_owned(),
            public_key: "ssh-ed25519 AAAA test".to_owned(),
            public_key_fingerprint: "ssh:blake3:test".to_owned(),
            cert_kind: "user".to_owned(),
            principals: vec!["eric".to_owned()],
            requested_validity: Some("+52w".to_owned()),
            renewal_of: None,
            reason: Some("renewal".to_owned()),
            status: "pending".to_owned(),
            created_at_ms: 1,
        };
        store
            .insert_ssh_cert_request(&request)
            .expect("insert request");
        assert_eq!(
            store
                .get_ssh_cert_request("ssh-cert-request:1")
                .expect("get request"),
            Some(request.clone())
        );

        let revocation = StoredSshRevocation {
            revocation_id: "ssh-revocation:1".to_owned(),
            kind: "public-key".to_owned(),
            target: "ssh:blake3:test".to_owned(),
            reason: Some("lost key".to_owned()),
            created_at_ms: 2,
            published: true,
        };
        store
            .insert_ssh_revocation(&revocation)
            .expect("insert revocation");
        assert_eq!(
            store.list_ssh_revocations().expect("list revocations"),
            vec![revocation]
        );
    }
}