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::{Connection, OptionalExtension, params};
use x25519_dalek::{PublicKey as X25519Public, StaticSecret};

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 prekey_sk: Vec<u8>,
    pub prekey_pk: Vec<u8>,
}

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

pub struct FriendProfile {
    pub display_name: String,
    pub bio: String,
    pub xmr_addr: String,
    pub updated_at: i64,
}

#[derive(Debug)]
pub struct Message {
    pub dir: String,
    pub plaintext: Vec<u8>,
}

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

    pub fn home_dir() -> Result<PathBuf> {
        home_dir()
    }

    pub 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,
              prekey BLOB NOT NULL DEFAULT x''
            );
            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,
              prekey_sk BLOB NOT NULL,
              prekey_pk BLOB NOT NULL,
              hs_nickname TEXT NOT NULL DEFAULT 'ow0'
            );
            CREATE TABLE IF NOT EXISTS self_profile (
              id INTEGER PRIMARY KEY CHECK (id = 1),
              display_name TEXT NOT NULL DEFAULT '',
              bio TEXT NOT NULL DEFAULT '',
              xmr_addr TEXT NOT NULL DEFAULT '',
              updated_at INTEGER NOT NULL DEFAULT 0
            );
            CREATE TABLE IF NOT EXISTS friend_profiles (
              pubkey BLOB PRIMARY KEY,
              display_name TEXT NOT NULL,
              bio TEXT NOT NULL,
              xmr_addr TEXT NOT NULL DEFAULT '',
              updated_at INTEGER NOT NULL
            );
            INSERT OR IGNORE INTO self_profile (id) VALUES (1);
            ",
        )?;
        let store = Self { conn };
        store.migrate()?;
        store.ensure_self()?;
        Ok(store)
    }

    fn migrate(&self) -> Result<()> {
        self.add_column_if_missing(
            "friends",
            "last_connect_ok",
            "ALTER TABLE friends ADD COLUMN last_connect_ok INTEGER NOT NULL DEFAULT 1",
        )?;
        self.add_column_if_missing(
            "friends",
            "prekey",
            "ALTER TABLE friends ADD COLUMN prekey BLOB NOT NULL DEFAULT x''",
        )?;
        self.add_column_if_missing(
            "self",
            "prekey_sk",
            "ALTER TABLE self ADD COLUMN prekey_sk BLOB NOT NULL DEFAULT x''",
        )?;
        self.add_column_if_missing(
            "self",
            "prekey_pk",
            "ALTER TABLE self ADD COLUMN prekey_pk BLOB NOT NULL DEFAULT x''",
        )?;
        self.add_column_if_missing(
            "self",
            "hs_nickname",
            "ALTER TABLE self ADD COLUMN hs_nickname TEXT NOT NULL DEFAULT 'ow0'",
        )?;
        Ok(())
    }

    fn add_column_if_missing(&self, table: &str, column: &str, ddl: &str) -> Result<()> {
        let sql =
            format!("SELECT COUNT(*) FROM pragma_table_info('{table}') WHERE name = '{column}'");
        let has: i64 = self.conn.query_row(&sql, [], |row| row.get(0))?;
        if has == 0 {
            self.conn.execute(ddl, [])?;
        }
        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();
            let (psk, ppk) = gen_prekey();
            self.conn.execute(
                "INSERT INTO self (id, identity_sk, identity_pk, onion, onion_rotated_at, prekey_sk, prekey_pk, hs_nickname)
                 VALUES (1, ?1, ?2, '', 0, ?3, ?4, 'ow0')",
                params![sk, pk, psk, ppk],
            )?;
        } else {
            let empty: i64 = self.conn.query_row(
                "SELECT CASE WHEN length(prekey_sk) = 32 AND length(prekey_pk) = 32 THEN 0 ELSE 1 END FROM self WHERE id = 1",
                [],
                |row| row.get(0),
            )?;
            if empty != 0 {
                let (psk, ppk) = gen_prekey();
                self.conn.execute(
                    "UPDATE self SET prekey_sk = ?1, prekey_pk = ?2 WHERE id = 1",
                    params![psk, ppk],
                )?;
            }
        }
        Ok(())
    }

    pub fn self_identity(&self) -> Result<SelfIdentity> {
        self.conn
            .query_row(
                "SELECT identity_sk, identity_pk, onion, prekey_sk, prekey_pk FROM self WHERE id = 1",
                [],
                |row| {
                    Ok(SelfIdentity {
                        identity_sk: row.get(0)?,
                        identity_pk: row.get(1)?,
                        onion: row.get(2)?,
                        prekey_sk: row.get(3)?,
                        prekey_pk: row.get(4)?,
                    })
                },
            )
            .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, prekey FROM friends WHERE pubkey = ?1",
                params![pubkey],
                friend_from_row,
            )
            .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, prekey FROM friends
             ORDER BY COALESCE(petname, fingerprint) COLLATE NOCASE",
        )?;
        let rows = stmt.query_map([], friend_from_row)?;
        let mut out = Vec::new();
        for row in rows {
            out.push(row?);
        }
        Ok(out)
    }

    pub fn hs_nickname(&self) -> Result<String> {
        self.conn
            .query_row("SELECT hs_nickname FROM self WHERE id = 1", [], |row| {
                row.get(0)
            })
            .map_err(Into::into)
    }

    pub fn set_hs_nickname(&self, nickname: &str) -> Result<()> {
        self.conn.execute(
            "UPDATE self SET hs_nickname = ?1 WHERE id = 1",
            params![nickname],
        )?;
        Ok(())
    }

    pub fn set_onion(&self, onion: &str) -> Result<()> {
        let now = unix_now();
        self.conn.execute(
            "UPDATE self SET onion = ?1, onion_rotated_at = ?2 WHERE id = 1",
            params![onion, now],
        )?;
        Ok(())
    }

    /// Apply a signed locator update. Bad sig / stale ts / unknown friend → no change.
    pub fn apply_loc(&self, pubkey: &[u8], onion: &str, ts: i64, sig: &[u8]) -> Result<bool> {
        let loc = crate::loc::Loc {
            onion: onion.to_string(),
            ts,
            sig: sig.to_vec(),
        };
        if !crate::loc::verify(pubkey, &loc) {
            return Ok(false);
        }
        let Some(updated_at) = self
            .conn
            .query_row(
                "SELECT onion_updated_at FROM friends WHERE pubkey = ?1",
                params![pubkey],
                |row| row.get::<_, i64>(0),
            )
            .optional()?
        else {
            return Ok(false);
        };
        if ts <= updated_at {
            return Ok(false);
        }
        let n = self.conn.execute(
            "UPDATE friends SET onion = ?1, onion_updated_at = ?2 WHERE pubkey = ?3",
            params![onion, ts, pubkey],
        )?;
        Ok(n > 0)
    }

    pub fn set_self_profile(&self, display_name: &str, bio: &str, xmr_addr: &str) -> Result<()> {
        crate::profile::validate(display_name, bio, xmr_addr).map_err(|e| Error(e.to_string()))?;
        let prev: i64 = self.conn.query_row(
            "SELECT updated_at FROM self_profile WHERE id = 1",
            [],
            |row| row.get(0),
        )?;
        let now = unix_now().max(prev + 1);
        self.conn.execute(
            "UPDATE self_profile SET display_name = ?1, bio = ?2, xmr_addr = ?3, updated_at = ?4 WHERE id = 1",
            params![display_name, bio, xmr_addr, now],
        )?;
        Ok(())
    }

    pub fn self_profile(&self) -> Result<FriendProfile> {
        self.conn
            .query_row(
                "SELECT display_name, bio, xmr_addr, updated_at FROM self_profile WHERE id = 1",
                [],
                profile_from_row,
            )
            .map_err(Into::into)
    }

    /// Apply a signed profile. Unknown friend / bad sig / stale ts → no change, no insert.
    pub fn apply_profile(&self, pubkey: &[u8], prf: &crate::profile::Profile) -> Result<bool> {
        if !crate::profile::verify(pubkey, prf) {
            return Ok(false);
        }
        let exists: i64 = self.conn.query_row(
            "SELECT COUNT(*) FROM friends WHERE pubkey = ?1",
            params![pubkey],
            |row| row.get(0),
        )?;
        if exists == 0 {
            return Ok(false);
        }
        let prev: Option<i64> = self
            .conn
            .query_row(
                "SELECT updated_at FROM friend_profiles WHERE pubkey = ?1",
                params![pubkey],
                |row| row.get(0),
            )
            .optional()?;
        if prev.is_some_and(|t| prf.ts <= t) {
            return Ok(false);
        }
        self.conn.execute(
            "INSERT INTO friend_profiles (pubkey, display_name, bio, xmr_addr, updated_at)
             VALUES (?1, ?2, ?3, ?4, ?5)
             ON CONFLICT(pubkey) DO UPDATE SET
               display_name = excluded.display_name,
               bio = excluded.bio,
               xmr_addr = excluded.xmr_addr,
               updated_at = excluded.updated_at",
            params![pubkey, prf.display_name, prf.bio, prf.xmr_addr, prf.ts],
        )?;
        Ok(true)
    }

    pub fn friend_profile(&self, pubkey: &[u8]) -> Result<Option<FriendProfile>> {
        self.conn
            .query_row(
                "SELECT display_name, bio, xmr_addr, updated_at FROM friend_profiles WHERE pubkey = ?1",
                params![pubkey],
                profile_from_row,
            )
            .optional()
            .map_err(Into::into)
    }

    pub fn set_friend_prekey(&self, pubkey: &[u8], prekey: &[u8]) -> Result<()> {
        let n = self.conn.execute(
            "UPDATE friends SET prekey = ?1 WHERE pubkey = ?2",
            params![prekey, pubkey],
        )?;
        if n == 0 {
            return Err(Error("friend not found".into()));
        }
        Ok(())
    }

    pub fn get_friend_by_prekey(&self, prekey: &[u8]) -> Result<Option<Friend>> {
        self.conn
            .query_row(
                "SELECT pubkey, fingerprint, petname, onion, last_connect_ok, prekey FROM friends WHERE prekey = ?1",
                params![prekey],
                friend_from_row,
            )
            .optional()
            .map_err(Into::into)
    }

    pub fn append_message(&self, friend_pk: &[u8], dir: &str, plaintext: &[u8]) -> Result<()> {
        if dir != "in" && dir != "out" {
            return Err(Error("dir must be in or out".into()));
        }
        let friend_id: i64 = self
            .conn
            .query_row(
                "SELECT id FROM friends WHERE pubkey = ?1",
                params![friend_pk],
                |row| row.get(0),
            )
            .map_err(|_| Error("friend not found".into()))?;
        self.conn.execute(
            "INSERT INTO messages (friend_id, dir, plaintext, created_at) VALUES (?1, ?2, ?3, ?4)",
            params![friend_id, dir, plaintext, unix_now()],
        )?;
        Ok(())
    }

    pub fn list_messages(&self, friend_pk: &[u8]) -> Result<Vec<Message>> {
        let mut stmt = self.conn.prepare(
            "SELECT m.dir, m.plaintext FROM messages m
             JOIN friends f ON f.id = m.friend_id
             WHERE f.pubkey = ?1
             ORDER BY m.id",
        )?;
        let rows = stmt.query_map(params![friend_pk], |row| {
            Ok(Message {
                dir: row.get(0)?,
                plaintext: row.get(1)?,
            })
        })?;
        let mut out = Vec::new();
        for row in rows {
            out.push(row?);
        }
        Ok(out)
    }

    /// Overwrite message bodies, delete rows, VACUUM. Identity + friends stay.
    pub fn wipe_messages(&self) -> Result<()> {
        self.conn.execute(
            "UPDATE messages SET plaintext = zeroblob(length(plaintext))",
            [],
        )?;
        self.conn.execute("DELETE FROM messages", [])?;
        self.conn.execute_batch("VACUUM")?;
        Ok(())
    }

    /// Delete the data dir. Next `open_at` is a new identity. Caller must drop Store first.
    pub fn wipe_all(home: &Path) -> Result<()> {
        if home.exists() {
            fs::remove_dir_all(home)?;
        }
        Ok(())
    }
}

fn profile_from_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<FriendProfile> {
    Ok(FriendProfile {
        display_name: row.get(0)?,
        bio: row.get(1)?,
        xmr_addr: row.get(2)?,
        updated_at: row.get(3)?,
    })
}

fn friend_from_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<Friend> {
    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,
        prekey: row.get(5)?,
    })
}

fn gen_prekey() -> (Vec<u8>, Vec<u8>) {
    let sk = StaticSecret::random_from_rng(OsRng);
    let pk = X25519Public::from(&sk);
    (sk.to_bytes().to_vec(), pk.to_bytes().to_vec())
}

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()
}