onionwire/src/wallet.rs

use std::time::Duration;

use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

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

#[derive(Debug, Clone)]
struct Endpoint {
    host: String,
    port: u16,
    path: String,
}

#[derive(Debug, Clone)]
pub struct Wallet {
    endpoint: Option<Endpoint>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TransferRow {
    pub txid: String,
    pub amount: String,
    pub address: String,
}

const RPC_TIMEOUT: Duration = Duration::from_secs(5);

impl Wallet {
    pub fn disabled() -> Self {
        Self { endpoint: None }
    }

    pub fn from_env() -> Self {
        match std::env::var("ONIONWIRE_WALLET_RPC") {
            Ok(s) if !s.trim().is_empty() => match Self::from_url(s.trim()) {
                Ok(w) => w,
                Err(e) => {
                    eprintln!("ONIONWIRE_WALLET_RPC: {e}");
                    Self::disabled()
                }
            },
            _ => Self::disabled(),
        }
    }

    pub fn from_url(url: &str) -> Result<Self> {
        Ok(Self {
            endpoint: Some(parse_http_url(url)?),
        })
    }

    pub fn configured(&self) -> bool {
        self.endpoint.is_some()
    }

    pub async fn get_address(&self) -> Result<String> {
        let v = self
            .rpc("get_address", serde_json::json!({"account_index": 0}))
            .await?;
        json_str(&v, "address")
    }

    pub async fn create_address(&self) -> Result<String> {
        let v = self
            .rpc("create_address", serde_json::json!({"account_index": 0}))
            .await?;
        json_str(&v, "address")
    }

    pub async fn transfer(&self, address: &str, amount: u64) -> Result<String> {
        let v = self
            .rpc(
                "transfer",
                serde_json::json!({
                    "destinations": [{"amount": amount, "address": address}],
                    "account_index": 0
                }),
            )
            .await?;
        json_str(&v, "tx_hash")
    }

    pub async fn get_transfers(&self) -> Result<Vec<TransferRow>> {
        let v = self
            .rpc(
                "get_transfers",
                serde_json::json!({"in": true, "pending": true}),
            )
            .await?;
        let mut out = Vec::new();
        for key in ["in", "pending"] {
            if let Some(arr) = v.get(key).and_then(|x| x.as_array()) {
                for item in arr {
                    out.push(TransferRow {
                        txid: json_str(item, "txid").unwrap_or_default(),
                        amount: json_amount(item),
                        address: json_str(item, "address").unwrap_or_default(),
                    });
                }
            }
        }
        Ok(out)
    }

    async fn rpc(&self, method: &str, params: serde_json::Value) -> Result<serde_json::Value> {
        let ep = self
            .endpoint
            .as_ref()
            .ok_or_else(|| Error("not configured".into()))?;
        let body = serde_json::json!({
            "jsonrpc": "2.0",
            "id": "0",
            "method": method,
            "params": params,
        })
        .to_string();
        let host_hdr = host_header(&ep.host, ep.port);
        let req = format!(
            "POST {} HTTP/1.1\r\nHost: {host_hdr}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
            ep.path,
            body.len()
        );
        let raw = tokio::time::timeout(RPC_TIMEOUT, http_post(ep, req.as_bytes()))
            .await
            .map_err(|_| Error("wallet RPC timed out".into()))??;
        parse_json_rpc(&raw)
    }
}

pub fn transfers_match(rows: &[TransferRow], txid: &str, amount: &str, address: &str) -> bool {
    rows.iter().any(|r| {
        r.txid == txid
            || (!address.is_empty()
                && r.address == address
                && (amount.is_empty() || r.amount == amount))
    })
}

fn parse_http_url(url: &str) -> Result<Endpoint> {
    let rest = url
        .strip_prefix("http://")
        .ok_or_else(|| Error("wallet RPC must be http:// (no TLS)".into()))?;
    if rest.contains("://") {
        return Err(Error("wallet RPC must be http:// (no TLS)".into()));
    }
    let (hostport, path) = match rest.split_once('/') {
        Some((hp, p)) => (hp, format!("/{p}")),
        None => (rest, "/json_rpc".into()),
    };
    let path = if path == "/" {
        "/json_rpc".into()
    } else {
        path
    };
    let (host, port) = parse_hostport(hostport)?;
    if !allowed_host(&host) {
        return Err(Error("wallet RPC host must be loopback or .onion".into()));
    }
    Ok(Endpoint { host, port, path })
}

fn parse_hostport(hostport: &str) -> Result<(String, u16)> {
    if let Some(rest) = hostport.strip_prefix('[') {
        let (host, rest) = rest
            .split_once(']')
            .ok_or_else(|| Error("invalid IPv6 host".into()))?;
        let port = match rest.strip_prefix(':') {
            Some(p) if !p.is_empty() => parse_port(p)?,
            _ => 18083,
        };
        if host.is_empty() {
            return Err(Error("empty host".into()));
        }
        return Ok((host.to_string(), port));
    }
    match hostport.rsplit_once(':') {
        Some((h, p)) if !h.is_empty() && !p.is_empty() && p.bytes().all(|b| b.is_ascii_digit()) => {
            Ok((h.to_string(), parse_port(p)?))
        }
        _ if !hostport.is_empty() => Ok((hostport.to_string(), 18083)),
        _ => Err(Error("empty host".into())),
    }
}

fn parse_port(p: &str) -> Result<u16> {
    p.parse()
        .map_err(|_| Error("invalid wallet RPC port".into()))
}

fn allowed_host(host: &str) -> bool {
    let h = host.trim();
    if h.to_ascii_lowercase().ends_with(".onion") {
        return true;
    }
    if h.eq_ignore_ascii_case("localhost") {
        return true;
    }
    h.parse::<std::net::IpAddr>()
        .map(|ip| ip.is_loopback())
        .unwrap_or(false)
}

fn host_header(host: &str, port: u16) -> String {
    if host.contains(':') {
        format!("[{host}]:{port}")
    } else {
        format!("{host}:{port}")
    }
}

async fn http_post(ep: &Endpoint, req: &[u8]) -> Result<Vec<u8>> {
    let mut stream = TcpStream::connect((ep.host.as_str(), ep.port))
        .await
        .map_err(|e| Error(format!("wallet connect: {e}")))?;
    stream
        .write_all(req)
        .await
        .map_err(|e| Error(format!("wallet write: {e}")))?;
    let mut buf = Vec::new();
    stream
        .read_to_end(&mut buf)
        .await
        .map_err(|e| Error(format!("wallet read: {e}")))?;
    Ok(buf)
}

fn parse_json_rpc(raw: &[u8]) -> Result<serde_json::Value> {
    let text = std::str::from_utf8(raw).map_err(|_| Error("wallet RPC not UTF-8".into()))?;
    let (head, body) = text
        .split_once("\r\n\r\n")
        .or_else(|| text.split_once("\n\n"))
        .ok_or_else(|| Error("wallet RPC truncated".into()))?;
    let status = head.lines().next().unwrap_or("");
    if status.contains(" 3") {
        return Err(Error("wallet RPC redirect refused".into()));
    }
    if !status.contains(" 200 ") && !status.ends_with(" 200") {
        return Err(Error(format!("wallet RPC HTTP {status}")));
    }
    let json: serde_json::Value =
        serde_json::from_str(body.trim()).map_err(|e| Error(format!("wallet RPC json: {e}")))?;
    if let Some(err) = json.get("error") {
        let msg = err
            .get("message")
            .and_then(|m| m.as_str())
            .unwrap_or("rpc error");
        return Err(Error(msg.into()));
    }
    json.get("result")
        .cloned()
        .ok_or_else(|| Error("wallet RPC missing result".into()))
}

fn json_str(v: &serde_json::Value, key: &str) -> Result<String> {
    v.get(key)
        .and_then(|x| x.as_str())
        .map(str::to_string)
        .ok_or_else(|| Error(format!("wallet RPC missing {key}")))
}

fn json_amount(v: &serde_json::Value) -> String {
    match v.get("amount") {
        Some(serde_json::Value::Number(n)) => n
            .as_u64()
            .map(|x| x.to_string())
            .or_else(|| n.as_i64().map(|x| x.to_string()))
            .unwrap_or_default(),
        Some(serde_json::Value::String(s)) => s.clone(),
        _ => String::new(),
    }
}