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>,
}

#[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 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
            );
            ",
        )?;
        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''",
        )?;
        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)
                 VALUES (1, ?1, ?2, '', 0, ?3, ?4)",
                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 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(())
    }

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

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