onionwire/src/node.rs

//! Two-node onion session: host HS, Noise IK, persist plaintext locally.

use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use futures::StreamExt;
use futures::io::{AsyncRead, AsyncWrite};
use tor_cell::relaycell::msg::Connected;
use tor_hsservice::{RunningOnionService, handle_rend_requests};

use crate::dispatch::{self, Kind};
use crate::frame;
use crate::hs::{self, Client, HS_PORT};
use crate::loc;
use crate::pay;
use crate::profile;
use crate::qr;
use crate::ratelimit::TokenBucket;
use crate::session::{self, Keys};
use crate::store::{Friend, FriendProfile, Message, PaymentWrite, Store};
use crate::wallet::{self, Wallet};

pub struct RotateResult {
    pub notified: usize,
    pub friends: usize,
}

struct HsHandle {
    _svc: Arc<RunningOnionService>,
    rend: tokio::task::JoinHandle<()>,
}

impl Drop for HsHandle {
    fn drop(&mut self) {
        self.rend.abort();
    }
}

pub struct Node {
    home: PathBuf,
    store: Mutex<Store>,
    client: Client,
    hs: Mutex<Option<HsHandle>>,
    onion: Mutex<String>,
    keys: Mutex<Keys>,
    wallet: Wallet,
    incoming_limit: Mutex<TokenBucket>,
}

impl Node {
    pub async fn start(home: PathBuf) -> Result<Arc<Self>, String> {
        let store = Store::open_at(&home).map_err(|e| e.to_string())?;
        let me = store.self_identity().map_err(|e| e.to_string())?;
        let keys = Keys::from_self(&me).map_err(|e| e.to_string())?;
        let nickname = store.hs_nickname().map_err(|e| e.to_string())?;
        let state = home.join("arti");
        let cache = home.join("cache");
        let client = hs::bootstrapped(&state, &cache).await?;
        let launched = client
            .launch_onion_service(hs::hs_config(&nickname)?)
            .map_err(|e| format!("launch_onion_service: {e}"))?
            .ok_or_else(|| "onion service disabled in config — fail closed".to_string())?;
        let (svc, rend) = launched;
        let onion = hs::onion_string(&svc)?;
        hs::wait_until_published(&svc, &onion).await?;
        store.set_onion(&onion).map_err(|e| e.to_string())?;
        let node = Arc::new(Self {
            home,
            store: Mutex::new(store),
            client,
            hs: Mutex::new(None),
            onion: Mutex::new(onion),
            keys: Mutex::new(keys),
            wallet: Wallet::from_env(),
            incoming_limit: Mutex::new(TokenBucket::default()),
        });
        let rend = spawn_rend(Arc::clone(&node), rend);
        *node.hs.lock().map_err(|e| e.to_string())? = Some(HsHandle { _svc: svc, rend });
        Ok(node)
    }

    pub fn onion(&self) -> String {
        self.onion.lock().map(|g| g.clone()).unwrap_or_default()
    }

    fn keys(&self) -> Result<Keys, String> {
        self.keys
            .lock()
            .map(|g| g.clone())
            .map_err(|e| e.to_string())
    }

    pub fn identity_pk(&self) -> [u8; 32] {
        self.keys().map(|k| k.identity_pk).unwrap_or([0; 32])
    }

    pub fn arti_dir(&self) -> PathBuf {
        self.home.join("arti")
    }

    pub fn cache_dir(&self) -> PathBuf {
        self.home.join("cache")
    }

    pub fn qr_payload(&self) -> Result<String, String> {
        let k = self.keys()?;
        qr::encode(&k.identity_sk, &self.onion(), &k.prekey_pk).map_err(|e| e.to_string())
    }

    pub fn write_backup(&self, path: &str, passphrase: &str) -> Result<(), String> {
        use std::io::Write;
        use std::os::unix::fs::OpenOptionsExt;
        let k = self.keys()?;
        let blob = crate::backup::seal(
            passphrase,
            &crate::backup::BackupKeys {
                identity_sk: k.identity_sk,
                identity_pk: k.identity_pk,
                prekey_sk: k.prekey_sk,
                prekey_pk: k.prekey_pk,
            },
        )?;
        let mut f = std::fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .mode(0o600)
            .open(path)
            .map_err(|e| e.to_string())?;
        f.write_all(&blob).map_err(|e| e.to_string())
    }

    pub fn restore_backup(&self, path: &str, passphrase: &str) -> Result<(), String> {
        let blob = std::fs::read(path).map_err(|e| e.to_string())?;
        let k = crate::backup::open(passphrase, &blob)?;
        self.store
            .lock()
            .map_err(|e| e.to_string())?
            .replace_identity_keys(&k.identity_sk, &k.identity_pk, &k.prekey_sk, &k.prekey_pk)
            .map_err(|e| e.to_string())?;
        *self.keys.lock().map_err(|e| e.to_string())? = Keys {
            identity_sk: k.identity_sk,
            identity_pk: k.identity_pk,
            prekey_sk: k.prekey_sk,
            prekey_pk: k.prekey_pk,
        };
        Ok(())
    }

    pub fn add_friend_from_qr(&self, raw: &str) -> Result<(), String> {
        let p = qr::decode(raw).map_err(|e| e.to_string())?;
        self.add_friend_payload(&p)
    }

    pub fn add_friend_payload(&self, p: &qr::QrPayload) -> Result<(), String> {
        let store = self.store.lock().map_err(|e| e.to_string())?;
        store
            .upsert_friend(&p.pubkey, &p.onion, None)
            .map_err(|e| e.to_string())?;
        store
            .set_friend_prekey(&p.pubkey, &p.signed_prekey)
            .map_err(|e| e.to_string())?;
        Ok(())
    }

    pub fn get_friend(&self, pubkey: &[u8]) -> Result<Option<Friend>, String> {
        self.store
            .lock()
            .map_err(|e| e.to_string())?
            .get_friend(pubkey)
            .map_err(|e| e.to_string())
    }

    pub fn list_friends(&self) -> Result<Vec<Friend>, String> {
        self.store
            .lock()
            .map_err(|e| e.to_string())?
            .list_friends()
            .map_err(|e| e.to_string())
    }

    pub fn friend(&self, pubkey: &[u8]) -> Result<Friend, String> {
        self.store
            .lock()
            .map_err(|e| e.to_string())?
            .get_friend(pubkey)
            .map_err(|e| e.to_string())?
            .ok_or_else(|| "friend not found".into())
    }

    pub fn list_messages(&self, friend_pk: &[u8]) -> Result<Vec<Message>, String> {
        self.store
            .lock()
            .map_err(|e| e.to_string())?
            .list_messages(friend_pk)
            .map_err(|e| e.to_string())
    }

    pub fn wipe_messages(&self) -> Result<(), String> {
        self.store
            .lock()
            .map_err(|e| e.to_string())?
            .wipe_messages()
            .map_err(|e| e.to_string())
    }

    pub fn self_profile(&self) -> Result<FriendProfile, String> {
        self.store
            .lock()
            .map_err(|e| e.to_string())?
            .self_profile()
            .map_err(|e| e.to_string())
    }

    pub fn friend_profile(&self, pubkey: &[u8]) -> Result<Option<FriendProfile>, String> {
        self.store
            .lock()
            .map_err(|e| e.to_string())?
            .friend_profile(pubkey)
            .map_err(|e| e.to_string())
    }

    pub fn set_self_profile(
        &self,
        display_name: &str,
        bio: &str,
        xmr_addr: &str,
    ) -> Result<(), String> {
        self.store
            .lock()
            .map_err(|e| e.to_string())?
            .set_self_profile(display_name, bio, xmr_addr)
            .map_err(|e| e.to_string())
    }

    /// One-shot signed `prf` to a friend. Fail closed; no retry, no outbox.
    pub async fn push_self_profile(&self, friend_pk: &[u8]) -> Result<(), String> {
        let me = self.self_profile()?;
        let k = self.keys()?;
        let prf = profile::sign(
            &k.identity_sk,
            &me.display_name,
            &me.bio,
            &me.xmr_addr,
            me.updated_at,
        )
        .map_err(|e| e.to_string())?;
        let pt = profile::encode(&prf);
        let (onion, prekey) = {
            let f = self.friend(friend_pk)?;
            if f.prekey.len() != 32 {
                return Err("friend missing prekey".into());
            }
            (f.onion, f.prekey)
        };
        self.try_send(&onion, friend_pk, &prekey, &pt).await
    }

    /// Invoice to receive. Wallet subaddress if up, else self_profile.xmr_addr.
    pub async fn pay_invoice(
        &self,
        friend_pk: &[u8],
        atomic: &str,
        memo: &str,
    ) -> Result<(), String> {
        let address = self.invoice_address().await?;
        let k = self.keys()?;
        let inv = pay::sign_invoice(&k.identity_sk, atomic, &address, memo, unix_now())
            .map_err(|e| e.to_string())?;
        let pt = pay::encode_invoice(&inv);
        self.send_once(friend_pk, &pt).await?;
        let store = self.store.lock().map_err(|e| e.to_string())?;
        store
            .insert_payment(
                friend_pk,
                PaymentWrite {
                    dir: "out",
                    kind: "invoice",
                    amount_atomic: atomic,
                    address: &address,
                    memo,
                    txid: None,
                    verified: false,
                },
            )
            .map_err(|e| e.to_string())?;
        store
            .append_message(
                friend_pk,
                "out",
                pay::invoice_chat_line(atomic, memo).as_bytes(),
            )
            .map_err(|e| e.to_string())?;
        Ok(())
    }

    /// Pay selected friend via wallet RPC, then signed receipt.
    pub async fn tip(&self, friend_pk: &[u8], atomic: &str, _memo: &str) -> Result<(), String> {
        if !self.wallet.configured() {
            return Err("wallet not connected — set ONIONWIRE_WALLET_RPC".into());
        }
        let address = self
            .friend_profile(friend_pk)?
            .map(|p| p.xmr_addr)
            .unwrap_or_default();
        if address.is_empty() {
            return Err("friend has no Monero address — they must /profile".into());
        }
        pay::check_address(&address).map_err(|e| e.to_string())?;
        let amount: u64 = atomic
            .parse()
            .map_err(|_| "amount too large for wallet RPC".to_string())?;
        let txid = self
            .wallet
            .transfer(&address, amount)
            .await
            .map_err(|e| e.to_string())?;
        let k = self.keys()?;
        let rcp = pay::sign_receipt(&k.identity_sk, &txid, atomic, &address, unix_now())
            .map_err(|e| e.to_string())?;
        let pt = pay::encode_receipt(&rcp);
        self.send_once(friend_pk, &pt).await?;
        let store = self.store.lock().map_err(|e| e.to_string())?;
        store
            .insert_payment(
                friend_pk,
                PaymentWrite {
                    dir: "out",
                    kind: "receipt",
                    amount_atomic: atomic,
                    address: &address,
                    memo: "",
                    txid: Some(&txid),
                    verified: true,
                },
            )
            .map_err(|e| e.to_string())?;
        store
            .append_message(
                friend_pk,
                "out",
                pay::receipt_chat_line(atomic, true).as_bytes(),
            )
            .map_err(|e| e.to_string())?;
        Ok(())
    }

    async fn invoice_address(&self) -> Result<String, String> {
        if self.wallet.configured()
            && let Ok(addr) = self.wallet.create_address().await
            && pay::check_address(&addr).is_ok()
        {
            return Ok(addr);
        }
        let profile_addr = self.self_profile()?.xmr_addr;
        if profile_addr.is_empty() {
            return Err("no receive address — set profile xmr or ONIONWIRE_WALLET_RPC".into());
        }
        pay::check_address(&profile_addr).map_err(|e| e.to_string())?;
        Ok(profile_addr)
    }

    async fn send_once(&self, friend_pk: &[u8], plaintext: &[u8]) -> Result<(), String> {
        let (onion, prekey) = {
            let f = self.friend(friend_pk)?;
            if f.prekey.len() != 32 {
                return Err("friend missing prekey".into());
            }
            (f.onion, f.prekey)
        };
        self.try_send(&onion, friend_pk, &prekey, plaintext).await
    }

    pub async fn connect_onion(&self, onion: &str) -> Result<(), String> {
        tokio::time::timeout(
            Duration::from_secs(20),
            self.client.connect((onion, HS_PORT)),
        )
        .await
        .map_err(|_| format!("connect {onion}:{HS_PORT} timed out"))?
        .map_err(|e| format!("connect {onion}:{HS_PORT}: {e}"))?;
        Ok(())
    }

    pub async fn rotate(self: &Arc<Self>) -> Result<RotateResult, String> {
        let old_nick = self
            .store
            .lock()
            .map_err(|e| e.to_string())?
            .hs_nickname()
            .map_err(|e| e.to_string())?;
        let new_nick = next_hs_nickname(&old_nick);
        let launched = self
            .client
            .launch_onion_service(hs::hs_config(&new_nick)?)
            .map_err(|e| format!("launch_onion_service: {e}"))?
            .ok_or_else(|| "onion service disabled in config — fail closed".to_string())?;
        let (svc, rend) = launched;
        let onion = hs::onion_string(&svc)?;
        // ponytail: hard-cut old HS before waiting; keeping both stalled ow1 at Bootstrapping.
        *self.hs.lock().map_err(|e| e.to_string())? = None;
        hs::wait_until_published(&svc, &onion).await?;
        let rend = spawn_rend(Arc::clone(self), rend);
        *self.hs.lock().map_err(|e| e.to_string())? = Some(HsHandle { _svc: svc, rend });
        {
            let store = self.store.lock().map_err(|e| e.to_string())?;
            store.set_onion(&onion).map_err(|e| e.to_string())?;
            store
                .set_hs_nickname(&new_nick)
                .map_err(|e| e.to_string())?;
        }
        *self.onion.lock().map_err(|e| e.to_string())? = onion.clone();

        let ts = unix_now();
        let k = self.keys()?;
        let loc = loc::sign(&k.identity_sk, &onion, ts).map_err(|e| e.to_string())?;
        let loc_pt = loc::encode(&loc);
        let friends = self
            .store
            .lock()
            .map_err(|e| e.to_string())?
            .list_friends()
            .map_err(|e| e.to_string())?;
        let mut notified = 0;
        for f in &friends {
            if self
                .push_loc(&f.pubkey, &f.onion, &f.prekey, &loc_pt)
                .await
                .is_ok()
            {
                notified += 1;
            }
        }

        Ok(RotateResult {
            notified,
            friends: friends.len(),
        })
    }

    pub async fn send(&self, friend_pk: &[u8], plaintext: &[u8]) -> Result<(), String> {
        let (onion, prekey) = {
            let store = self.store.lock().map_err(|e| e.to_string())?;
            let f = store
                .get_friend(friend_pk)
                .map_err(|e| e.to_string())?
                .ok_or_else(|| "unknown friend".to_string())?;
            if f.prekey.len() != 32 {
                return Err("friend missing prekey".into());
            }
            (f.onion, f.prekey)
        };
        let deadline = Instant::now() + Duration::from_secs(180);
        let mut last = None::<String>;
        while Instant::now() < deadline {
            match self.try_send(&onion, friend_pk, &prekey, plaintext).await {
                Ok(()) => {
                    self.store
                        .lock()
                        .map_err(|e| e.to_string())?
                        .append_message(friend_pk, "out", plaintext)
                        .map_err(|e| e.to_string())?;
                    return Ok(());
                }
                Err(e) if e.contains("fingerprint mismatch") => return Err(e),
                Err(e) => last = Some(e),
            }
            tokio::time::sleep(Duration::from_secs(3)).await;
        }
        Err(last.unwrap_or_else(|| "send timed out".into()))
    }

    async fn push_loc(
        &self,
        friend_pk: &[u8],
        onion: &str,
        prekey: &[u8],
        loc_pt: &[u8],
    ) -> Result<(), String> {
        if prekey.len() != 32 {
            return Err("friend missing prekey".into());
        }
        let deadline = Instant::now() + Duration::from_secs(180);
        let mut last = None::<String>;
        while Instant::now() < deadline {
            match self.try_send(onion, friend_pk, prekey, loc_pt).await {
                Ok(()) => return Ok(()),
                Err(e) if e.contains("fingerprint mismatch") => return Err(e),
                Err(e) => last = Some(e),
            }
            tokio::time::sleep(Duration::from_secs(3)).await;
        }
        Err(last.unwrap_or_else(|| "loc push timed out".into()))
    }

    async fn try_send(
        &self,
        onion: &str,
        pinned_id: &[u8],
        remote_prekey: &[u8],
        plaintext: &[u8],
    ) -> Result<(), String> {
        let mut stream = self
            .client
            .connect((onion, HS_PORT))
            .await
            .map_err(|e| format!("connect {onion}:{HS_PORT}: {e}"))?;
        let keys = self.keys()?;
        let mut sess = session::handshake_initiator(&mut stream, &keys, pinned_id, remote_prekey)
            .await
            .map_err(session_err)?;
        let ct = sess.encrypt(plaintext).map_err(session_err)?;
        frame::write_frame(&mut stream, &ct)
            .await
            .map_err(|e| e.to_string())?;
        Ok(())
    }

    async fn handle_incoming<S>(&self, stream: &mut S) -> Result<(), String>
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        let store = &self.store;
        let keys = self.keys()?;
        let mut sess = session::handshake_responder(stream, &keys, |spk| {
            let Ok(g) = store.lock() else {
                return None;
            };
            g.get_friend_by_prekey(spk).ok().flatten().map(|f| f.pubkey)
        })
        .await
        .map_err(session_err)?;
        let ct = frame::read_frame(stream).await.map_err(|e| e.to_string())?;
        let pt = sess.decrypt(&ct).map_err(session_err)?;
        match dispatch::classify(&pt) {
            Kind::Loc => {
                if let Some(loc) = loc::decode(&pt) {
                    let applied = store.lock().map_err(|e| e.to_string())?.apply_loc(
                        &sess.peer_identity,
                        &loc.onion,
                        loc.ts,
                        &loc.sig,
                    );
                    match applied {
                        Ok(true) => {}
                        Ok(false) => {
                            eprintln!("loc dropped (bad sig, stale ts, or unknown friend)")
                        }
                        Err(e) => return Err(e.to_string()),
                    }
                }
                Ok(())
            }
            Kind::Profile => {
                if let Some(prf) = profile::decode(&pt) {
                    let applied = store
                        .lock()
                        .map_err(|e| e.to_string())?
                        .apply_profile(&sess.peer_identity, &prf);
                    match applied {
                        Ok(true) => {}
                        Ok(false) => {
                            eprintln!("prf dropped (bad sig, stale ts, or unknown friend)")
                        }
                        Err(e) => return Err(e.to_string()),
                    }
                }
                Ok(())
            }
            Kind::Chat => {
                store
                    .lock()
                    .map_err(|e| e.to_string())?
                    .append_message(&sess.peer_identity, "in", &pt)
                    .map_err(|e| e.to_string())?;
                Ok(())
            }
            Kind::Invoice => self.ingest_invoice(&sess.peer_identity, &pt).await,
            Kind::Receipt => self.ingest_receipt(&sess.peer_identity, &pt).await,
            Kind::Ping | Kind::Drop => Ok(()),
        }
    }

    async fn ingest_invoice(&self, peer: &[u8], pt: &[u8]) -> Result<(), String> {
        let Some(inv) = pay::decode_invoice(pt) else {
            return Ok(());
        };
        if !pay::verify_invoice(peer, &inv) {
            eprintln!("inv dropped (bad sig)");
            return Ok(());
        }
        let store = self.store.lock().map_err(|e| e.to_string())?;
        store
            .insert_payment(
                peer,
                PaymentWrite {
                    dir: "in",
                    kind: "invoice",
                    amount_atomic: &inv.amount_atomic,
                    address: &inv.address,
                    memo: &inv.memo,
                    txid: None,
                    verified: false,
                },
            )
            .map_err(|e| e.to_string())?;
        store
            .append_message(
                peer,
                "in",
                pay::invoice_chat_line(&inv.amount_atomic, &inv.memo).as_bytes(),
            )
            .map_err(|e| e.to_string())?;
        Ok(())
    }

    async fn ingest_receipt(&self, peer: &[u8], pt: &[u8]) -> Result<(), String> {
        let Some(rcp) = pay::decode_receipt(pt) else {
            return Ok(());
        };
        if !pay::verify_receipt(peer, &rcp) {
            eprintln!("rcp dropped (bad sig)");
            return Ok(());
        }
        let id = {
            let store = self.store.lock().map_err(|e| e.to_string())?;
            store
                .insert_payment(
                    peer,
                    PaymentWrite {
                        dir: "in",
                        kind: "receipt",
                        amount_atomic: &rcp.amount_atomic,
                        address: &rcp.address,
                        memo: "",
                        txid: Some(&rcp.txid),
                        verified: false,
                    },
                )
                .map_err(|e| e.to_string())?
        };
        let mut verified = false;
        if self.wallet.configured() {
            match self.wallet.get_transfers().await {
                Ok(rows)
                    if wallet::transfers_match(
                        &rows,
                        &rcp.txid,
                        &rcp.amount_atomic,
                        &rcp.address,
                    ) =>
                {
                    self.store
                        .lock()
                        .map_err(|e| e.to_string())?
                        .mark_verified(id)
                        .map_err(|e| e.to_string())?;
                    verified = true;
                }
                Ok(_) => {}
                Err(e) => eprintln!("receipt unverified: {e}"),
            }
        }
        if !verified {
            eprintln!("receipt unverified");
        }
        self.store
            .lock()
            .map_err(|e| e.to_string())?
            .append_message(
                peer,
                "in",
                pay::receipt_chat_line(&rcp.amount_atomic, verified).as_bytes(),
            )
            .map_err(|e| e.to_string())?;
        Ok(())
    }
}

fn spawn_rend(
    node: Arc<Node>,
    rend: impl futures::Stream<Item = tor_hsservice::RendRequest> + Send + 'static,
) -> tokio::task::JoinHandle<()> {
    tokio::spawn(async move {
        let mut requests = std::pin::pin!(handle_rend_requests(rend));
        while let Some(req) = requests.next().await {
            let allow = node
                .incoming_limit
                .lock()
                .map(|mut b| b.try_acquire())
                .unwrap_or(false);
            if !allow {
                eprintln!("rate-limit drop");
                drop(req);
                continue;
            }
            let serve = Arc::clone(&node);
            tokio::spawn(async move {
                let Ok(mut stream) = req.accept(Connected::new_empty()).await else {
                    return;
                };
                if let Err(e) = serve.handle_incoming(&mut stream).await {
                    eprintln!("incoming: {e}");
                }
            });
        }
    })
}

fn next_hs_nickname(cur: &str) -> String {
    let n = cur
        .strip_prefix("ow")
        .and_then(|s| s.parse::<u32>().ok())
        .unwrap_or(0);
    format!("ow{}", n.saturating_add(1))
}

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

fn session_err(e: session::Error) -> String {
    if e.is_fingerprint_mismatch() {
        crate::tui::fingerprint_mismatch_banner().to_string()
    } else {
        e.to_string()
    }
}

pub fn dir_contains_bytes(dir: &Path, needle: &[u8]) -> bool {
    let mut stack = vec![dir.to_path_buf()];
    while let Some(p) = stack.pop() {
        let Ok(rd) = std::fs::read_dir(&p) else {
            continue;
        };
        for ent in rd.flatten() {
            let path = ent.path();
            if path.is_dir() {
                stack.push(path);
            } else if let Ok(bytes) = std::fs::read(&path)
                && bytes.windows(needle.len()).any(|w| w == needle)
            {
                return true;
            }
        }
    }
    false
}