onionwire/src/store.rs

use std::fs;
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};

use ed25519_dalek::SigningKey;
use rand::rngs::OsRng;
use rusqlite::{params, Connection, OptionalExtension};

pub type Result<T> = std::result::Result<T, Error>;

#[derive(Debug)]
pub struct Error(String);

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.fmt(f)
    }
}

impl std::error::Error for Error {}

impl From<rusqlite::Error> for Error {
    fn from(e: rusqlite::Error) -> Self {
        Self(e.to_string())
    }
}

impl From<std::io::Error> for Error {
    fn from(e: std::io::Error) -> Self {
        Self(e.to_string())
    }
}

pub struct Store {
    conn: Connection,
}

pub struct SelfIdentity {
    pub identity_sk: Vec<u8>,
    pub identity_pk: Vec<u8>,
    pub onion: String,
}

pub struct Friend {
    pub pubkey: Vec<u8>,
    pub fingerprint: String,
    pub petname: Option<String>,
    pub onion: String,
    pub last_connect_ok: bool,
}

impl Store {
    pub fn open() -> Result<Self> {
        Self::open_at(&home_dir()?)
    }

    fn open_at(home: &Path) -> Result<Self> {
        mkdir_700(home)?;
        mkdir_700(&home.join("arti"))?;
        let db_path = home.join("onionwire.db");
        let conn = Connection::open(&db_path)?;
        conn.execute_batch(
            "
            PRAGMA foreign_keys = ON;
            CREATE TABLE IF NOT EXISTS friends (
              id INTEGER PRIMARY KEY,
              pubkey BLOB NOT NULL UNIQUE,
              fingerprint TEXT NOT NULL,
              petname TEXT,
              onion TEXT NOT NULL,
              onion_updated_at INTEGER NOT NULL,
              added_at INTEGER NOT NULL,
              last_connect_ok INTEGER NOT NULL DEFAULT 1
            );
            CREATE TABLE IF NOT EXISTS messages (
              id INTEGER PRIMARY KEY,
              friend_id INTEGER NOT NULL REFERENCES friends(id),
              dir TEXT NOT NULL CHECK(dir IN ('in','out')),
              plaintext BLOB NOT NULL,
              created_at INTEGER NOT NULL
            );
            CREATE TABLE IF NOT EXISTS self (
              id INTEGER PRIMARY KEY CHECK (id = 1),
              identity_sk BLOB NOT NULL,
              identity_pk BLOB NOT NULL,
              onion TEXT NOT NULL,
              onion_rotated_at INTEGER NOT NULL
            );
            ",
        )?;
        let store = Self { conn };
        store.migrate()?;
        store.ensure_self()?;
        Ok(store)
    }

    fn migrate(&self) -> Result<()> {
        let has: i64 = self.conn.query_row(
            "SELECT COUNT(*) FROM pragma_table_info('friends') WHERE name = 'last_connect_ok'",
            [],
            |row| row.get(0),
        )?;
        if has == 0 {
            self.conn.execute(
                "ALTER TABLE friends ADD COLUMN last_connect_ok INTEGER NOT NULL DEFAULT 1",
                [],
            )?;
        }
        Ok(())
    }

    fn ensure_self(&self) -> Result<()> {
        let exists: i64 =
            self.conn
                .query_row("SELECT COUNT(*) FROM self WHERE id = 1", [], |row| {
                    row.get(0)
                })?;
        if exists == 0 {
            let signing = SigningKey::generate(&mut OsRng);
            let sk = signing.to_bytes().to_vec();
            let pk = signing.verifying_key().to_bytes().to_vec();
            self.conn.execute(
                "INSERT INTO self (id, identity_sk, identity_pk, onion, onion_rotated_at)
                 VALUES (1, ?1, ?2, '', 0)",
                params![sk, pk],
            )?;
        }
        Ok(())
    }

    pub fn self_identity(&self) -> Result<SelfIdentity> {
        self.conn
            .query_row(
                "SELECT identity_sk, identity_pk, onion FROM self WHERE id = 1",
                [],
                |row| {
                    Ok(SelfIdentity {
                        identity_sk: row.get(0)?,
                        identity_pk: row.get(1)?,
                        onion: row.get(2)?,
                    })
                },
            )
            .map_err(Into::into)
    }

    pub fn friend_count(&self) -> Result<i64> {
        self.conn
            .query_row("SELECT COUNT(*) FROM friends", [], |row| row.get(0))
            .map_err(Into::into)
    }

    pub fn upsert_friend(&self, pubkey: &[u8], onion: &str, petname: Option<&str>) -> Result<()> {
        let now = unix_now();
        let fp = fingerprint(pubkey);
        self.conn.execute(
            "INSERT INTO friends (pubkey, fingerprint, petname, onion, onion_updated_at, added_at)
             VALUES (?1, ?2, ?3, ?4, ?5, ?5)
             ON CONFLICT(pubkey) DO UPDATE SET
               onion = excluded.onion,
               onion_updated_at = excluded.onion_updated_at,
               petname = COALESCE(excluded.petname, friends.petname)",
            params![pubkey, fp, petname, onion, now],
        )?;
        Ok(())
    }

    pub fn get_friend(&self, pubkey: &[u8]) -> Result<Option<Friend>> {
        self.conn
            .query_row(
                "SELECT pubkey, fingerprint, petname, onion, last_connect_ok FROM friends WHERE pubkey = ?1",
                params![pubkey],
                |row| {
                    Ok(Friend {
                        pubkey: row.get(0)?,
                        fingerprint: row.get(1)?,
                        petname: row.get(2)?,
                        onion: row.get(3)?,
                        last_connect_ok: row.get::<_, i64>(4)? != 0,
                    })
                },
            )
            .optional()
            .map_err(Into::into)
    }

    pub fn list_friends(&self) -> Result<Vec<Friend>> {
        let mut stmt = self.conn.prepare(
            "SELECT pubkey, fingerprint, petname, onion, last_connect_ok FROM friends
             ORDER BY COALESCE(petname, fingerprint) COLLATE NOCASE",
        )?;
        let rows = stmt.query_map([], |row| {
            Ok(Friend {
                pubkey: row.get(0)?,
                fingerprint: row.get(1)?,
                petname: row.get(2)?,
                onion: row.get(3)?,
                last_connect_ok: row.get::<_, i64>(4)? != 0,
            })
        })?;
        let mut out = Vec::new();
        for row in rows {
            out.push(row?);
        }
        Ok(out)
    }
}

fn home_dir() -> Result<PathBuf> {
    if let Some(p) = std::env::var_os("ONIONWIRE_HOME") {
        return Ok(PathBuf::from(p));
    }
    let home = std::env::var_os("HOME").ok_or_else(|| Error("HOME not set".into()))?;
    Ok(PathBuf::from(home).join(".local/share/onionwire"))
}

fn mkdir_700(path: &Path) -> Result<()> {
    fs::create_dir_all(path)?;
    let mut perms = fs::metadata(path)?.permissions();
    perms.set_mode(0o700);
    fs::set_permissions(path, perms)?;
    Ok(())
}

fn unix_now() -> i64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs() as i64)
        .unwrap_or(0)
}

fn fingerprint(pubkey: &[u8]) -> String {
    pubkey.iter().map(|b| format!("{b:02x}")).collect()
}