onionwire/src/wallet.rs

use std::time::Duration;

use rand::RngCore;
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(Clone)]
struct Endpoint {
    host: String,
    port: u16,
    path: String,
    user: String,
    pass: String,
}

impl std::fmt::Debug for Endpoint {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Endpoint")
            .field("host", &self.host)
            .field("port", &self.port)
            .field("path", &self.path)
            .field("user", &self.user)
            .field("pass", &"<redacted>")
            .finish()
    }
}

#[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() => {
                let login = std::env::var("ONIONWIRE_WALLET_RPC_LOGIN")
                    .ok()
                    .map(|v| v.trim().to_string())
                    .filter(|v| !v.is_empty());
                match parse_http_url(s.trim(), login.as_deref()) {
                    Ok(endpoint) => Self {
                        endpoint: Some(endpoint),
                    },
                    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, None)?),
        })
    }

    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 raw = post_timeout(ep, &build_req(ep, &body, None)).await?;
        let status = http_status(&raw).unwrap_or(0);
        if status == 200 {
            return Err(Error(
                "wallet RPC requires digest auth (open RPC refused)".into(),
            ));
        }
        if status != 401 {
            return parse_json_rpc(&raw);
        }
        let challenge = www_authenticate(&raw).ok_or_else(|| {
            Error("wallet RPC digest required (--rpc-login / HTTP Digest)".into())
        })?;
        if !challenge.trim().to_ascii_lowercase().starts_with("digest") {
            return Err(Error(
                "wallet RPC digest required (--rpc-login / HTTP Digest)".into(),
            ));
        }
        let auth = digest_authorization(ep, &challenge)?;
        let raw = post_timeout(ep, &build_req(ep, &body, Some(&auth))).await?;
        parse_json_rpc(&raw)
    }
}

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

fn credentials_required() -> Error {
    Error("wallet RPC requires credentials (user:pass in URL or ONIONWIRE_WALLET_RPC_LOGIN)".into())
}

fn parse_login(login: &str) -> Result<(String, String)> {
    let (user, pass) = login.split_once(':').ok_or_else(credentials_required)?;
    if user.is_empty() || pass.is_empty() {
        return Err(credentials_required());
    }
    Ok((user.to_string(), pass.to_string()))
}

fn parse_http_url(url: &str, extra_login: Option<&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 (userinfo, rest) = match rest.rsplit_once('@') {
        Some((ui, hostpart)) => (Some(ui), hostpart),
        None => (None, rest),
    };
    let (user, pass) = match userinfo {
        Some(ui) => parse_login(ui)?,
        None => match extra_login {
            Some(login) => parse_login(login)?,
            None => return Err(credentials_required()),
        },
    };
    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 host.trim().to_ascii_lowercase().ends_with(".onion") {
        return Err(Error(
            "wallet RPC over .onion is not supported (loopback only; no Arti dial)".into(),
        ));
    }
    if !allowed_host(&host) {
        return Err(Error("wallet RPC host must be loopback".into()));
    }
    Ok(Endpoint {
        host,
        port,
        path,
        user,
        pass,
    })
}

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

const MAX_RPC_BYTES: usize = 1024 * 1024;

async fn post_timeout(ep: &Endpoint, req: &str) -> Result<Vec<u8>> {
    tokio::time::timeout(RPC_TIMEOUT, http_post(ep, req.as_bytes()))
        .await
        .map_err(|_| Error("wallet RPC timed out".into()))?
}

fn build_req(ep: &Endpoint, body: &str, authorization: Option<&str>) -> String {
    let host_hdr = host_header(&ep.host, ep.port);
    let mut req = format!(
        "POST {} HTTP/1.1\r\nHost: {host_hdr}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n",
        ep.path,
        body.len()
    );
    if let Some(auth) = authorization {
        req.push_str("Authorization: ");
        req.push_str(auth);
        req.push_str("\r\n");
    }
    req.push_str("\r\n");
    req.push_str(body);
    req
}

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();
    let mut tmp = [0u8; 8192];
    loop {
        let n = stream
            .read(&mut tmp)
            .await
            .map_err(|e| Error(format!("wallet read: {e}")))?;
        if n == 0 {
            break;
        }
        if buf.len().saturating_add(n) > MAX_RPC_BYTES {
            return Err(Error("wallet RPC response too large".into()));
        }
        buf.extend_from_slice(&tmp[..n]);
    }
    Ok(buf)
}

fn http_status(raw: &[u8]) -> Option<u16> {
    let text = std::str::from_utf8(raw).ok()?;
    let line = text.lines().next()?;
    let mut parts = line.split_whitespace();
    let _http = parts.next()?;
    parts.next()?.parse().ok()
}

fn www_authenticate(raw: &[u8]) -> Option<String> {
    let text = std::str::from_utf8(raw).ok()?;
    let (head, _) = text
        .split_once("\r\n\r\n")
        .or_else(|| text.split_once("\n\n"))?;
    for line in head.lines().skip(1) {
        let (k, v) = match line.split_once(':') {
            Some(kv) => kv,
            None => continue,
        };
        if k.eq_ignore_ascii_case("www-authenticate") {
            return Some(v.trim().to_string());
        }
    }
    None
}

fn digest_param(challenge: &str, key: &str) -> Option<String> {
    let t = challenge.trim();
    let rest = if t.len() >= 6 && t[..6].eq_ignore_ascii_case("digest") {
        t[6..].trim()
    } else {
        return None;
    };
    for part in rest.split(',') {
        let part = part.trim();
        let (k, v) = match part.split_once('=') {
            Some(kv) => kv,
            None => continue,
        };
        if k.eq_ignore_ascii_case(key) {
            return Some(v.trim().trim_matches('"').to_string());
        }
    }
    None
}

fn md5_hex(s: &str) -> String {
    use md5::{Digest, Md5};
    hex_lower(&Md5::digest(s.as_bytes()))
}

fn hex_lower(bytes: &[u8]) -> String {
    const H: &[u8; 16] = b"0123456789abcdef";
    let mut out = String::with_capacity(bytes.len() * 2);
    for &b in bytes {
        out.push(H[(b >> 4) as usize] as char);
        out.push(H[(b & 0xf) as usize] as char);
    }
    out
}

fn digest_authorization(ep: &Endpoint, challenge: &str) -> Result<String> {
    if let Some(alg) = digest_param(challenge, "algorithm")
        && !alg.eq_ignore_ascii_case("MD5")
    {
        return Err(Error("wallet RPC digest algorithm not MD5".into()));
    }
    let realm = digest_param(challenge, "realm").unwrap_or_default();
    let nonce = digest_param(challenge, "nonce")
        .ok_or_else(|| Error("wallet RPC digest required (--rpc-login / HTTP Digest)".into()))?;
    let qop = digest_param(challenge, "qop");
    let ha1 = md5_hex(&format!("{}:{realm}:{}", ep.user, ep.pass));
    let ha2 = md5_hex(&format!("POST:{}", ep.path));
    let (qop_part, resp) = if qop
        .as_deref()
        .is_some_and(|q| q.split(',').any(|x| x.trim() == "auth"))
    {
        let mut cnonce_bytes = [0u8; 8];
        rand::thread_rng().fill_bytes(&mut cnonce_bytes);
        let cnonce = hex_lower(&cnonce_bytes);
        let nc = "00000001";
        let response = md5_hex(&format!("{ha1}:{nonce}:{nc}:{cnonce}:auth:{ha2}"));
        (
            format!(", qop=auth, nc={nc}, cnonce=\"{cnonce}\""),
            response,
        )
    } else {
        (String::new(), md5_hex(&format!("{ha1}:{nonce}:{ha2}")))
    };
    Ok(format!(
        "Digest username=\"{}\", realm=\"{realm}\", nonce=\"{nonce}\", uri=\"{}\", algorithm=MD5, response=\"{resp}\"{qop_part}",
        ep.user, ep.path
    ))
}

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