AnyBot/src/tunnel.rs

// 長連線: one lasting WebSocket from this AnyBot to the account server. Requests for this AnyBot
// (the phone, a screen of this brain, a brain using this computer's hands) come down it at once;
// this AnyBot's own requests to the account's other AnyBots go up it. Big bodies are compressed.
// A ping every 15 s keeps it (and any proxy on the way) awake; 45 s without a word and it is
// opened again. If the account server is older and has no WebSocket, the old way (asking
// "anything for me?" over and over) is used instead.

use crate::httpd::Req;
use crate::wire;
use serde_json::{json, Value};
use std::collections::HashMap;
use std::io::{ErrorKind, Read, Write};
use std::net::{TcpStream, ToSocketAddrs};
use std::sync::mpsc::{channel, Sender};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{Duration, Instant};
use tungstenite::{Message, WebSocket};

const PING_SECS: u64 = 15;
const DEAD_SECS: u64 = 45;
/// Bodies bigger than this are compressed.
const ZIP_OVER: usize = 1024;
const CUT: &str = "和帳號伺服器的連線斷了，正在重新連線";

/// Another AnyBot's answer to a call.
pub struct Answer {
    pub status: u16,
    pub body: Vec<u8>,
}

type Calls = Arc<Mutex<HashMap<String, Sender<Result<Answer, String>>>>>;

struct Link {
    tx: Sender<Vec<u8>>,
    calls: Calls,
    serial: u64,
}

fn link() -> &'static Mutex<Option<Link>> {
    static L: OnceLock<Mutex<Option<Link>>> = OnceLock::new();
    L.get_or_init(|| Mutex::new(None))
}


pub fn zip(b: &[u8]) -> Vec<u8> {
    let mut e = flate2::write::DeflateEncoder::new(Vec::with_capacity(b.len() / 3), flate2::Compression::fast());
    let _ = e.write_all(b);
    e.finish().unwrap_or_default()
}

/// The body to send and whether it is compressed: only when that saves a good part of it
/// (pictures and PDFs are compressed already).
fn maybe_zip(b: Vec<u8>) -> (Vec<u8>, bool) {
    if b.len() <= ZIP_OVER {
        return (b, false);
    }
    let z = zip(&b);
    if z.len() * 10 < b.len() * 9 {
        (z, true)
    } else {
        (b, false)
    }
}

pub fn unzip(b: &[u8]) -> Result<Vec<u8>, String> {
    let mut out = vec![];
    flate2::read::DeflateDecoder::new(b).take(512 * 1024 * 1024).read_to_end(&mut out).map_err(|_| "收到的資料壞掉了".to_string())?;
    Ok(out)
}

// ---------- the connection ----------

enum Stream {
    Plain(TcpStream),
    Tls(Box<rustls::StreamOwned<rustls::ClientConnection, TcpStream>>),
}

impl Stream {
    fn tcp(&self) -> &TcpStream {
        match self {
            Stream::Plain(t) => t,
            Stream::Tls(s) => &s.sock,
        }
    }
}

impl Read for Stream {
    fn read(&mut self, b: &mut [u8]) -> std::io::Result<usize> {
        match self {
            Stream::Plain(t) => t.read(b),
            Stream::Tls(s) => s.read(b),
        }
    }
}

impl Write for Stream {
    fn write(&mut self, b: &[u8]) -> std::io::Result<usize> {
        match self {
            Stream::Plain(t) => t.write(b),
            Stream::Tls(s) => s.write(b),
        }
    }
    fn flush(&mut self) -> std::io::Result<()> {
        match self {
            Stream::Plain(t) => t.flush(),
            Stream::Tls(s) => s.flush(),
        }
    }
}

fn tls_config() -> Arc<rustls::ClientConfig> {
    static C: OnceLock<Arc<rustls::ClientConfig>> = OnceLock::new();
    C.get_or_init(|| {
        let roots = rustls::RootCertStore { roots: webpki_roots::TLS_SERVER_ROOTS.to_vec() };
        Arc::new(
            rustls::ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider()))
                .with_safe_default_protocol_versions()
                .expect("TLS versions")
                .with_root_certificates(roots)
                .with_no_client_auth(),
        )
    })
    .clone()
}

fn would_block(e: &tungstenite::Error) -> bool {
    matches!(e, tungstenite::Error::Io(io) if matches!(io.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut))
}

/// Why it couldn't connect; `old` = the account server has no WebSocket (use the old way).
struct NoLink {
    old: bool,
    why: String,
}

fn open(url: &str, token: &str) -> Result<WebSocket<Stream>, NoLink> {
    use tungstenite::client::IntoClientRequest;
    let fail = |why: &str| NoLink { old: false, why: why.to_string() };
    let (tls, rest) = match (url.strip_prefix("https://"), url.strip_prefix("http://")) {
        (Some(r), _) => (true, r),
        (_, Some(r)) => (false, r),
        _ => return Err(fail("帳號伺服器的網址不對")),
    };
    let hostport = rest.split('/').next().unwrap_or("");
    let (host, port) = match hostport.rsplit_once(':') {
        Some((h, p)) if p.parse::<u16>().is_ok() => (h, p.parse().unwrap_or(443)),
        _ => (hostport, if tls { 443 } else { 80 }),
    };
    let addrs: Vec<_> = (host, port).to_socket_addrs().map_err(|_| fail("找不到伺服器：這台電腦可能沒連上網路，或伺服器網址打錯了"))?.collect();
    // Each address in turn ("localhost" is both ::1 and 127.0.0.1).
    let mut last = None;
    let mut tcp = None;
    for a in &addrs {
        match TcpStream::connect_timeout(a, Duration::from_secs(10)) {
            Ok(t) => {
                tcp = Some(t);
                break;
            }
            Err(e) => last = Some(e.kind()),
        }
    }
    let tcp = tcp.ok_or_else(|| {
        fail(match last {
            None => "找不到伺服器",
            Some(ErrorKind::ConnectionRefused) => "伺服器沒有開（連線被拒絕）",
            Some(ErrorKind::TimedOut | ErrorKind::WouldBlock) => "伺服器沒有回應：可能沒開機，或網路把連線擋住了",
            _ => "這台電腦沒有連上網路",
        })
    })?;
    let _ = tcp.set_nodelay(true);
    let _ = tcp.set_read_timeout(Some(Duration::from_secs(15)));
    let _ = tcp.set_write_timeout(Some(Duration::from_secs(30)));
    let stream = if tls {
        let name = rustls::pki_types::ServerName::try_from(host.to_string()).map_err(|_| fail("帳號伺服器的網址不對"))?;
        let conn = rustls::ClientConnection::new(tls_config(), name).map_err(|_| fail("和伺服器的安全連線失敗"))?;
        Stream::Tls(Box::new(rustls::StreamOwned::new(conn, tcp)))
    } else {
        Stream::Plain(tcp)
    };
    let mut req = format!("{}://{hostport}/relay/host/ws", if tls { "wss" } else { "ws" }).into_client_request().map_err(|_| fail("帳號伺服器的網址不對"))?;
    if let Ok(v) = format!("Bearer {token}").parse() {
        req.headers_mut().insert("Authorization", v);
    }
    let cfg = tungstenite::protocol::WebSocketConfig::default().max_message_size(Some(512 << 20)).max_frame_size(Some(512 << 20));
    match tungstenite::client::client_with_config(req, stream, Some(cfg)) {
        Ok((ws, _)) => Ok(ws),
        Err(tungstenite::HandshakeError::Failure(tungstenite::Error::Http(r))) => match r.status().as_u16() {
            401 | 403 => Err(fail("這台電腦的帳號登入失效了，請到「設定 › 手機與帳號」重新登入")),
            400 | 404 | 405 | 426 => Err(NoLink { old: true, why: String::new() }),
            _ => Err(fail("帳號伺服器暫時無法使用，稍後會自動再試")),
        },
        Err(tungstenite::HandshakeError::Failure(e)) if e.to_string().to_lowercase().contains("certificate") || e.to_string().to_lowercase().contains("tls") => Err(fail("和伺服器的安全連線失敗（伺服器的憑證有問題）")),
        Err(_) => Err(fail("連線到一半斷掉了")),
    }
}

/// Keep the connection for as long as `alive()`. Returns false when the account server has no
/// WebSocket (the caller then uses the old way).
pub fn run(url: &str, token: &str, alive: &dyn Fn() -> bool, status: &dyn Fn(&str)) -> bool {
    static SERIAL: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
    let mut backoff = 1u64;
    while alive() {
        match open(url, token) {
            Err(NoLink { old: true, .. }) => return false,
            Err(e) => {
                status(&format!("連不上帳號伺服器：{}", e.why));
                // Wait, but notice at once when told to stop.
                let until = Instant::now() + Duration::from_secs(backoff);
                while Instant::now() < until && alive() {
                    std::thread::sleep(Duration::from_millis(200));
                }
                backoff = (backoff * 2).min(30);
            }
            Ok(ws) => {
                backoff = 1;
                status("已連線");
                let serial = SERIAL.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
                keep(ws, serial, alive);
                if alive() {
                    status("連線中斷，正在重新連線…");
                }
            }
        }
    }
    true
}

fn keep(mut ws: WebSocket<Stream>, serial: u64, alive: &dyn Fn() -> bool) {
    let _ = ws.get_ref().tcp().set_read_timeout(Some(Duration::from_millis(40)));
    let (tx, rx) = channel::<Vec<u8>>();
    let calls: Calls = Arc::default();
    *link().lock().unwrap() = Some(Link { tx: tx.clone(), calls: calls.clone(), serial });
    let s = crate::remote::hub().settings();
    let _ = tx.send(wire::pack(&json!({"t": "hello", "brain": s.brain_id, "hands": s.device_link}), &[]));
    let (mut heard, mut pinged) = (Instant::now(), Instant::now());
    'live: while alive() {
        match ws.read() {
            Ok(Message::Binary(b)) => {
                heard = Instant::now();
                if let Some((head, body)) = wire::unpack(&b) {
                    on_frame(head, body, &tx, &calls);
                }
            }
            Ok(Message::Close(_)) => break,
            Ok(_) => heard = Instant::now(),
            Err(e) if would_block(&e) => {}
            Err(_) => break,
        }
        while let Ok(f) = rx.try_recv() {
            if ws.write(Message::Binary(f.into())).is_err() {
                break 'live;
            }
        }
        if pinged.elapsed().as_secs() >= PING_SECS {
            pinged = Instant::now();
            if ws.write(Message::Ping(Vec::new().into())).is_err() {
                break;
            }
        }
        match ws.flush() {
            Err(e) if !would_block(&e) => break,
            _ => {}
        }
        if heard.elapsed().as_secs() >= DEAD_SECS {
            break;
        }
    }
    let _ = ws.close(None);
    let _ = ws.flush();
    let mut l = link().lock().unwrap();
    if l.as_ref().map(|l| l.serial) == Some(serial) {
        *l = None;
    }
    drop(l);
    // Whoever waits on an answer hears now, not after a timeout.
    for (_, w) in calls.lock().unwrap().drain() {
        let _ = w.send(Err(CUT.into()));
    }
}

fn on_frame(head: Value, body: Vec<u8>, tx: &Sender<Vec<u8>>, calls: &Calls) {
    match head["t"].as_str().unwrap_or("") {
        "job" => {
            let tx = tx.clone();
            std::thread::spawn(move || {
                let s = |k: &str| head[k].as_str().unwrap_or("").to_string();
                let id = s("id");
                let body = if head["z"].as_bool().unwrap_or(false) { unzip(&body).unwrap_or_default() } else { body };
                let req = Req {
                    method: s("method"),
                    path: s("path"),
                    query: head["query"].as_object().map(|o| o.iter().map(|(k, v)| (k.clone(), v.as_str().unwrap_or("").to_string())).collect()).unwrap_or_default(),
                    headers: [("x-anybot-user", s("user")), ("x-anybot-role", if head["role"].is_string() { s("role") } else { "owner".into() }), ("x-anybot-from", s("from"))].into_iter().map(|(k, v)| (k.to_string(), v)).collect(),
                    body,
                    peer: "relay".into(),
                };
                let resp = crate::remote::handle(req, true);
                let (out, z) = if head["z_ok"].as_bool().unwrap_or(false) { maybe_zip(resp.body) } else { (resp.body, false) };
                let _ = tx.send(wire::pack(&json!({"t": "reply", "id": id, "status": resp.status, "ctype": resp.ctype, "headers": resp.headers, "z": z}), &out));
            });
        }
        "ret" => {
            if let Some(w) = calls.lock().unwrap().remove(head["id"].as_str().unwrap_or("")) {
                let body = if head["z"].as_bool().unwrap_or(false) { unzip(&body) } else { Ok(body) };
                let _ = w.send(body.map(|body| Answer { status: head["status"].as_u64().unwrap_or(502) as u16, body }));
            }
        }
        _ => {}
    }
}

/// Tell the account server which brain this AnyBot shows (so the phone can go there).
pub fn hello(brain: &str, hands: bool) {
    if let Some(l) = link().lock().unwrap().as_ref() {
        let _ = l.tx.send(wire::pack(&json!({"t": "hello", "brain": brain, "hands": hands}), &[]));
    }
}

/// Ask another AnyBot of this account over the lasting connection. None when it isn't open
/// (the caller then uses an ordinary request).
pub fn call(to: &str, method: &str, path: &str, query: &[(&str, String)], body: &[u8], secs: u64) -> Option<Result<Answer, String>> {
    let (tx, calls) = {
        let l = link().lock().unwrap();
        let l = l.as_ref()?;
        (l.tx.clone(), l.calls.clone())
    };
    let id = crate::httpd::random_hex(8);
    let (rtx, rrx) = channel();
    calls.lock().unwrap().insert(id.clone(), rtx);
    let (out, z) = maybe_zip(body.to_vec());
    let q: serde_json::Map<String, Value> = query.iter().map(|(k, v)| (k.to_string(), json!(v))).collect();
    if tx.send(wire::pack(&json!({"t": "call", "id": id, "to": to, "method": method, "path": path, "query": q, "z": z}), &out)).is_err() {
        calls.lock().unwrap().remove(&id);
        return None;
    }
    Some(match rrx.recv_timeout(Duration::from_secs(secs)) {
        Ok(r) => r,
        Err(_) => {
            calls.lock().unwrap().remove(&id);
            Err("等太久沒有回應：網路很慢或斷斷續續".into())
        }
    })
}

#[cfg(test)]
mod tests {
    #[test]
    fn compression_round_trips_and_shrinks_text() {
        let text = "你好，這是一段會重複的聊天內容。".repeat(400);
        let z = super::zip(text.as_bytes());
        assert!(z.len() * 5 < text.len());
        assert_eq!(super::unzip(&z).unwrap(), text.as_bytes());
        // Already-compressed bytes go as they are.
        let mut x = 88172645463325252u64;
        let noise: Vec<u8> = (0..5000)
            .map(|_| {
                x ^= x << 13;
                x ^= x >> 7;
                x ^= x << 17;
                (x >> 24) as u8
            })
            .collect();
        assert!(!super::maybe_zip(noise).1);
        assert!(super::maybe_zip(text.into_bytes()).1);
    }
}