blit_proxy/

lib.rs

/// blit-proxy library — all proxy logic, usable in-process or as a binary.
///
/// Call [`proxy_socket_path`] to find the socket, then [`run`] to start the
/// proxy on the current thread (blocking, runs its own tokio runtime).
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicUsize, Ordering};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::sync::{Mutex, RwLock};

static VERBOSE: AtomicBool = AtomicBool::new(false);

macro_rules! log {
    ($($arg:tt)*) => {
        if VERBOSE.load(Ordering::Relaxed) {
            eprintln!($($arg)*);
        }
    };
}

type BoxRead = Box<dyn AsyncRead + Unpin + Send>;
type BoxWrite = Box<dyn AsyncWrite + Unpin + Send>;

// ---------------------------------------------------------------------------
// Proxy socket path (single stable path for the whole process)
// ---------------------------------------------------------------------------

pub fn proxy_socket_path() -> String {
    if let Ok(p) = std::env::var("BLIT_PROXY_SOCK") {
        return p;
    }
    #[cfg(unix)]
    {
        let dir = std::env::var("XDG_RUNTIME_DIR").unwrap_or_else(|_| "/tmp".into());
        format!("{dir}/blit-proxy.sock")
    }
    #[cfg(windows)]
    {
        r"\\.\pipe\blit-proxy".into()
    }
}

// ---------------------------------------------------------------------------
// Upstream connection
// ---------------------------------------------------------------------------

struct UpstreamConn {
    reader: BoxRead,
    writer: BoxWrite,
}

// ---------------------------------------------------------------------------
// Per-target pool
// ---------------------------------------------------------------------------

struct Pool {
    idle: Mutex<VecDeque<UpstreamConn>>,
    /// Number of currently active (proxied) downstream clients.
    active: AtomicUsize,
    /// Unix seconds of last client connect or disconnect.
    /// `i64::MAX` while any client is active.
    last_activity: AtomicI64,
    pool_size: usize,
    upstream_uri: String,
}

impl Pool {
    fn new(upstream_uri: String, pool_size: usize) -> Arc<Self> {
        Arc::new(Self {
            idle: Mutex::new(VecDeque::new()),
            active: AtomicUsize::new(0),
            last_activity: AtomicI64::new(now_secs()),
            pool_size,
            upstream_uri,
        })
    }

    async fn acquire(&self) -> Result<UpstreamConn, String> {
        {
            let mut idle = self.idle.lock().await;
            if let Some(conn) = idle.pop_front() {
                return Ok(conn);
            }
        }
        self.connect_one().await
    }

    async fn connect_one(&self) -> Result<UpstreamConn, String> {
        connect_upstream(&self.upstream_uri).await
    }

    fn client_connected(&self) {
        self.active.fetch_add(1, Ordering::Relaxed);
        self.last_activity.store(i64::MAX, Ordering::Relaxed);
    }

    fn client_disconnected(&self) {
        let prev = self.active.fetch_sub(1, Ordering::Relaxed);
        if prev == 1 {
            self.last_activity.store(now_secs(), Ordering::Relaxed);
        }
    }

    /// Background task: keep idle slots full.
    async fn refill_loop(self: Arc<Self>) {
        loop {
            let need = {
                let idle = self.idle.lock().await;
                self.pool_size.saturating_sub(idle.len())
            };
            for _ in 0..need {
                match self.connect_one().await {
                    Ok(conn) => {
                        self.idle.lock().await.push_back(conn);
                    }
                    Err(e) => {
                        log!(
                            "blit-proxy: [{uri}] upstream connect failed: {e}",
                            uri = self.upstream_uri
                        );
                        tokio::time::sleep(std::time::Duration::from_secs(2)).await;
                        break;
                    }
                }
            }
            tokio::time::sleep(std::time::Duration::from_millis(200)).await;
        }
    }
}

fn now_secs() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs() as i64
}

// ---------------------------------------------------------------------------
// Pool registry: one pool per distinct upstream URI
// ---------------------------------------------------------------------------

struct Registry {
    pools: RwLock<HashMap<String, Arc<Pool>>>,
    pool_size: usize,
}

impl Registry {
    fn new(pool_size: usize) -> Arc<Self> {
        Arc::new(Self {
            pools: RwLock::new(HashMap::new()),
            pool_size,
        })
    }

    /// Get or create the pool for `uri`, spawning its refill task if new.
    async fn get_or_create(self: &Arc<Self>, uri: &str) -> Arc<Pool> {
        // Fast path: pool already exists.
        {
            let pools = self.pools.read().await;
            if let Some(p) = pools.get(uri) {
                return p.clone();
            }
        }
        // Slow path: create and seed.
        let mut pools = self.pools.write().await;
        // Re-check under write lock.
        if let Some(p) = pools.get(uri) {
            return p.clone();
        }
        let pool = Pool::new(uri.to_string(), self.pool_size);
        pools.insert(uri.to_string(), pool.clone());
        drop(pools);

        // Seed eagerly in a background task.
        let pool_seed = pool.clone();
        tokio::spawn(async move {
            for _ in 0..pool_seed.pool_size {
                match pool_seed.connect_one().await {
                    Ok(conn) => pool_seed.idle.lock().await.push_back(conn),
                    Err(e) => {
                        log!(
                            "blit-proxy: [{uri}] initial connect: {e}",
                            uri = pool_seed.upstream_uri
                        );
                        break;
                    }
                }
            }
            pool_seed.refill_loop().await;
        });

        pool
    }

    /// Returns the most-recent `last_activity` across all pools.
    /// Returns `i64::MAX` if any pool has an active client.
    async fn latest_activity(&self) -> i64 {
        let pools = self.pools.read().await;
        if pools.is_empty() {
            return now_secs();
        }
        pools
            .values()
            .map(|p| p.last_activity.load(Ordering::Relaxed))
            .max()
            .unwrap_or_else(now_secs)
    }
}

// ---------------------------------------------------------------------------
// Upstream transport implementations
// ---------------------------------------------------------------------------

async fn connect_upstream(uri: &str) -> Result<UpstreamConn, String> {
    if let Some(rest) = uri.strip_prefix("share:") {
        return connect_share(rest).await;
    }

    #[cfg(unix)]
    if let Some(rest) = uri.strip_prefix("ssh:") {
        return connect_ssh(rest).await;
    }

    // Extract query parameters from URIs that support them.
    let (base_uri, passphrase, cert_hash) = extract_uri_params(uri);

    if let Some(path) = base_uri.strip_prefix("socket:") {
        return connect_socket(path).await;
    }
    if let Some(addr) = base_uri.strip_prefix("tcp:") {
        return connect_tcp(addr).await;
    }
    if base_uri.starts_with("ws://") || base_uri.starts_with("wss://") {
        return connect_ws(&base_uri, passphrase.as_deref()).await;
    }
    if let Some(rest) = base_uri.strip_prefix("wt://") {
        let cert_bytes = cert_hash.as_deref().and_then(parse_hex);
        return connect_wt(rest, passphrase.as_deref(), &cert_bytes).await;
    }
    Err(format!(
        "unknown upstream URI scheme in '{uri}' \
         (expected socket:, tcp:, ws://, wss://, wt://, share:, or ssh:)"
    ))
}

/// Split `share:` URI rest into (passphrase, hub_url).
///
/// Accepted forms:
///   `myphrase`                       — use default hub
///   `myphrase?hub=wss://custom.hub`  — use specific hub
fn parse_share_uri(rest: &str) -> (String, String) {
    let (passphrase_raw, hub) = if let Some(q_pos) = rest.find('?') {
        let phrase = &rest[..q_pos];
        let query = &rest[q_pos + 1..];
        let hub = query
            .split('&')
            .find_map(|kv| kv.strip_prefix("hub=").map(percent_decode))
            .unwrap_or_else(|| blit_webrtc_forwarder::DEFAULT_HUB_URL.to_string());
        (phrase.to_string(), hub)
    } else {
        (
            rest.to_string(),
            blit_webrtc_forwarder::DEFAULT_HUB_URL.to_string(),
        )
    };
    let passphrase = percent_decode(&passphrase_raw);
    let hub = blit_webrtc_forwarder::normalize_hub(&hub);
    (passphrase, hub)
}

async fn connect_share(rest: &str) -> Result<UpstreamConn, String> {
    let (passphrase, hub) = parse_share_uri(rest);
    let stream = blit_webrtc_forwarder::client::connect(&passphrase, &hub)
        .await
        .map_err(|e| format!("share:{rest}: {e}"))?;
    let (r, w) = tokio::io::split(stream);
    Ok(UpstreamConn {
        reader: Box::new(r),
        writer: Box::new(w),
    })
}

// ---------------------------------------------------------------------------
// SSH stdio bridge
// ---------------------------------------------------------------------------

/// Shell snippet that locates the blit-server Unix socket on the remote host.
#[cfg(unix)]
const SSH_SOCK_SEARCH: &str = r#"if [ -n "$BLIT_SOCK" ]; then S="$BLIT_SOCK"; elif [ -n "$TMPDIR" ] && [ -S "$TMPDIR/blit.sock" ]; then S="$TMPDIR/blit.sock"; elif [ -S "/tmp/blit-$(id -un).sock" ]; then S="/tmp/blit-$(id -un).sock"; elif [ -S "/run/blit/$(id -un).sock" ]; then S="/run/blit/$(id -un).sock"; elif [ -n "$XDG_RUNTIME_DIR" ] && [ -S "$XDG_RUNTIME_DIR/blit.sock" ]; then S="$XDG_RUNTIME_DIR/blit.sock"; else S=/tmp/blit.sock; fi"#;

/// Shell snippet that starts blit-server on the remote if the socket is absent/stale.
#[cfg(unix)]
const SSH_AUTOSTART: &str = "if [ -S \"$S\" ]; then if command -v nc >/dev/null 2>&1; then nc -z -U \"$S\" 2>/dev/null || rm -f \"$S\"; elif command -v socat >/dev/null 2>&1; then socat /dev/null \"UNIX-CONNECT:$S\" 2>/dev/null || rm -f \"$S\"; fi; fi; if ! [ -S \"$S\" ]; then if command -v blit >/dev/null 2>&1; then blit server & i=0; while ! [ -S \"$S\" ] && [ $i -lt 50 ]; do sleep 0.1; i=$((i+1)); done; elif command -v blit-server >/dev/null 2>&1; then blit-server & i=0; while ! [ -S \"$S\" ] && [ $i -lt 50 ]; do sleep 0.1; i=$((i+1)); done; fi; fi";

/// Escape a string for use inside double quotes in a POSIX shell.
/// Handles `\`, `$`, `` ` ``, and `"`.  Used instead of single-quote
/// escaping because the result is embedded inside an outer `sh -c '…'`
/// wrapper where nested single quotes would break.
#[cfg(unix)]
fn dq_escape(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for ch in s.chars() {
        match ch {
            '\\' | '$' | '`' | '"' => {
                out.push('\\');
                out.push(ch);
            }
            _ => out.push(ch),
        }
    }
    out
}

/// Parse `[user@]host[:socket_path]` from an `ssh:` URI remainder.
/// Returns `(ssh_target, remote_socket)`.
#[cfg(unix)]
fn parse_ssh_target(rest: &str) -> (&str, Option<&str>) {
    let colon_search_start = rest.find('@').map(|a| a + 1).unwrap_or(0);
    if let Some(rel) = rest[colon_search_start..].find(':') {
        let pos = colon_search_start + rel;
        let path = &rest[pos + 1..];
        if path.is_empty() {
            (rest, None)
        } else {
            (&rest[..pos], Some(path))
        }
    } else {
        (rest, None)
    }
}

/// Connect to a remote blit-server via an SSH stdio bridge.
///
/// Spawns `ssh <host> -- sh -c '…find-socket…; nc -U "$S"'` and wraps
/// the child's stdin/stdout as an `UpstreamConn`.
#[cfg(unix)]
async fn connect_ssh(rest: &str) -> Result<UpstreamConn, String> {
    if rest.is_empty() {
        return Err("ssh: destination requires a host".into());
    }
    let (host, remote_socket) = parse_ssh_target(rest);
    if host.starts_with('-') {
        return Err(format!(
            "invalid ssh host '{host}': must not start with '-'"
        ));
    }
    let resolve = match remote_socket {
        Some(path) => format!("S=\"{}\"", dq_escape(path)),
        None => SSH_SOCK_SEARCH.to_string(),
    };
    let bridge = format!(
        "sh -c '{resolve}; {SSH_AUTOSTART}; if command -v nc >/dev/null 2>&1; then exec nc -U \"$S\"; else exec socat - \"UNIX-CONNECT:$S\"; fi'"
    );
    let mut child = tokio::process::Command::new("ssh")
        .arg("-T")
        .arg("-o")
        .arg("ControlMaster=auto")
        .arg("-o")
        .arg("ControlPath=/tmp/blit-ssh-%r@%h:%p")
        .arg("-o")
        .arg("ControlPersist=300")
        .arg(host)
        .arg("--")
        .arg(&bridge)
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::inherit())
        .spawn()
        .map_err(|e| format!("ssh: {e}"))?;
    let stdin = child
        .stdin
        .take()
        .ok_or_else(|| "ssh: could not get stdin".to_string())?;
    let stdout = child
        .stdout
        .take()
        .ok_or_else(|| "ssh: could not get stdout".to_string())?;
    // Keep the child alive by moving it into a background task that waits on it.
    tokio::spawn(async move {
        let _ = child.wait().await;
    });
    Ok(UpstreamConn {
        reader: Box::new(stdout),
        writer: Box::new(stdin),
    })
}

/// Split a URI into (base, passphrase, certHash) by parsing query params.
/// Only applies to ws://, wss://, wt:// — socket: and tcp: are returned as-is.
fn extract_uri_params(uri: &str) -> (String, Option<String>, Option<String>) {
    if !uri.starts_with("ws://") && !uri.starts_with("wss://") && !uri.starts_with("wt://") {
        return (uri.to_string(), None, None);
    }
    // Find '?'.
    let (base, query) = match uri.find('?') {
        Some(pos) => (&uri[..pos], Some(&uri[pos + 1..])),
        None => (uri, None),
    };
    let mut passphrase = None;
    let mut cert_hash = None;
    if let Some(q) = query {
        for param in q.split('&') {
            if let Some(v) = param.strip_prefix("passphrase=") {
                passphrase = Some(percent_decode(v));
            } else if let Some(v) = param.strip_prefix("certHash=") {
                cert_hash = Some(v.to_string());
            }
        }
    }
    (base.to_string(), passphrase, cert_hash)
}

fn percent_decode(s: &str) -> String {
    // Minimal %XX decoder sufficient for passphrases.
    let mut out = String::with_capacity(s.len());
    let mut chars = s.chars().peekable();
    while let Some(c) = chars.next() {
        if c == '%' {
            let h1 = chars.next().unwrap_or('0');
            let h2 = chars.next().unwrap_or('0');
            if let Ok(b) = u8::from_str_radix(&format!("{h1}{h2}"), 16) {
                out.push(b as char);
                continue;
            }
        }
        out.push(c);
    }
    out
}

async fn connect_socket(path: &str) -> Result<UpstreamConn, String> {
    #[cfg(unix)]
    {
        let stream = tokio::net::UnixStream::connect(path)
            .await
            .map_err(|e| format!("socket:{path}: {e}"))?;
        let (r, w) = tokio::io::split(stream);
        Ok(UpstreamConn {
            reader: Box::new(r),
            writer: Box::new(w),
        })
    }
    #[cfg(not(unix))]
    {
        use tokio::net::windows::named_pipe::ClientOptions;
        let pipe = ClientOptions::new()
            .open(path)
            .map_err(|e| format!("socket:{path}: {e}"))?;
        let (r, w) = tokio::io::split(pipe);
        Ok(UpstreamConn {
            reader: Box::new(r),
            writer: Box::new(w),
        })
    }
}

async fn connect_tcp(addr: &str) -> Result<UpstreamConn, String> {
    let stream = tokio::net::TcpStream::connect(addr)
        .await
        .map_err(|e| format!("tcp:{addr}: {e}"))?;
    let _ = stream.set_nodelay(true);
    let (r, w) = tokio::io::split(stream);
    Ok(UpstreamConn {
        reader: Box::new(r),
        writer: Box::new(w),
    })
}

async fn connect_ws(uri: &str, passphrase: Option<&str>) -> Result<UpstreamConn, String> {
    use futures_util::{SinkExt, StreamExt};
    use tokio_tungstenite::tungstenite::Message;

    let (mut ws, _) = tokio_tungstenite::connect_async(uri)
        .await
        .map_err(|e| format!("{uri}: {e}"))?;

    let pass = passphrase.unwrap_or("");
    ws.send(Message::Text(pass.into()))
        .await
        .map_err(|e| format!("{uri}: auth send: {e}"))?;
    match ws.next().await {
        Some(Ok(Message::Text(t))) if t.trim() == "ok" => {}
        Some(Ok(Message::Text(t))) => {
            return Err(format!("{uri}: auth rejected: {}", t.trim()));
        }
        other => {
            return Err(format!("{uri}: unexpected auth response: {other:?}"));
        }
    }

    let (ws_write, ws_read) = ws.split();
    Ok(UpstreamConn {
        reader: Box::new(WsFrameReader {
            inner: ws_read,
            buf: bytes::Bytes::new(),
        }),
        writer: Box::new(WsFrameWriter { inner: ws_write }),
    })
}

async fn connect_wt(
    rest: &str,
    passphrase: Option<&str>,
    cert_hash: &Option<Vec<u8>>,
) -> Result<UpstreamConn, String> {
    use web_transport_quinn as wt;

    // Build the URL for the WT session (must use https: scheme).
    let (host, port) = parse_wt_host_port(rest)?;
    let url: url::Url = format!("https://{host}:{port}/")
        .parse()
        .map_err(|e| format!("wt: url: {e}"))?;

    // Build the client with appropriate certificate verification.
    let client: wt::Client = if let Some(hash) = cert_hash {
        wt::ClientBuilder::new()
            .with_server_certificate_hashes(vec![hash.clone()])
            .map_err(|e| format!("wt: client build: {e}"))?
    } else {
        wt::ClientBuilder::new()
            .with_system_roots()
            .map_err(|e| format!("wt: client build: {e}"))?
    };

    let session = client
        .connect(url)
        .await
        .map_err(|e| format!("wt: connect {host}:{port}: {e}"))?;

    let (mut send, mut recv) = session
        .open_bi()
        .await
        .map_err(|e| format!("wt: open_bi: {e}"))?;

    // Auth: 2-byte-LE passphrase length + passphrase bytes, then read 1-byte response.
    let pass = passphrase.unwrap_or("").as_bytes();
    let mut auth_buf = Vec::with_capacity(2 + pass.len());
    auth_buf.extend_from_slice(&(pass.len() as u16).to_le_bytes());
    auth_buf.extend_from_slice(pass);
    send.write_all(&auth_buf)
        .await
        .map_err(|e| format!("wt: auth send: {e}"))?;

    let mut resp = [0u8; 1];
    recv.read_exact(&mut resp)
        .await
        .map_err(|e| format!("wt: auth recv: {e}"))?;
    if resp[0] != 1 {
        return Err(format!(
            "wt: auth rejected (response byte {:#04x})",
            resp[0]
        ));
    }

    Ok(UpstreamConn {
        reader: Box::new(recv),
        writer: Box::new(send),
    })
}

fn parse_wt_host_port(rest: &str) -> Result<(String, u16), String> {
    let without_path = rest.split('/').next().unwrap_or(rest);
    if let Some(colon) = without_path.rfind(':') {
        let host = without_path[..colon].to_string();
        let port: u16 = without_path[colon + 1..]
            .parse()
            .map_err(|_| format!("wt: invalid port in '{rest}'"))?;
        Ok((host, port))
    } else {
        Ok((without_path.to_string(), 443))
    }
}

fn parse_hex(s: &str) -> Option<Vec<u8>> {
    if !s.len().is_multiple_of(2) {
        return None;
    }
    (0..s.len())
        .step_by(2)
        .map(|i| u8::from_str_radix(&s[i..i + 2], 16).ok())
        .collect()
}

// ---------------------------------------------------------------------------
// WS ↔ raw blit-frame adapters
// ---------------------------------------------------------------------------

use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message;

type WsSink = futures_util::stream::SplitSink<
    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>,
    Message,
>;
type WsStream = futures_util::stream::SplitStream<
    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>,
>;

struct WsFrameReader {
    inner: WsStream,
    buf: bytes::Bytes,
}

impl AsyncRead for WsFrameReader {
    fn poll_read(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        out: &mut tokio::io::ReadBuf<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        loop {
            if !self.buf.is_empty() {
                let n = out.remaining().min(self.buf.len());
                out.put_slice(&self.buf[..n]);
                self.buf = self.buf.slice(n..);
                return std::task::Poll::Ready(Ok(()));
            }
            match self.inner.poll_next_unpin(cx) {
                std::task::Poll::Pending => return std::task::Poll::Pending,
                std::task::Poll::Ready(None) => return std::task::Poll::Ready(Ok(())),
                std::task::Poll::Ready(Some(Err(e))) => {
                    return std::task::Poll::Ready(Err(std::io::Error::other(e)));
                }
                std::task::Poll::Ready(Some(Ok(msg))) => {
                    let data = match msg {
                        Message::Binary(d) => d,
                        Message::Close(_) => return std::task::Poll::Ready(Ok(())),
                        _ => continue,
                    };
                    let len = data.len() as u32;
                    let mut framed = Vec::with_capacity(4 + data.len());
                    framed.extend_from_slice(&len.to_le_bytes());
                    framed.extend_from_slice(&data);
                    self.buf = bytes::Bytes::from(framed);
                }
            }
        }
    }
}

struct WsFrameWriter {
    inner: WsSink,
}

impl AsyncWrite for WsFrameWriter {
    fn poll_write(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<std::io::Result<usize>> {
        if buf.len() < 4 {
            return std::task::Poll::Ready(Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "ws frame writer: buffer too small for length prefix",
            )));
        }
        let len = u32::from_le_bytes(buf[..4].try_into().unwrap()) as usize;
        if buf.len() < 4 + len {
            return std::task::Poll::Ready(Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!(
                    "ws frame writer: buffer ({}) too small for frame payload ({})",
                    buf.len(),
                    4 + len
                ),
            )));
        }
        let payload = buf[4..4 + len].to_vec();
        match self.inner.poll_ready_unpin(cx) {
            std::task::Poll::Pending => return std::task::Poll::Pending,
            std::task::Poll::Ready(Err(e)) => {
                return std::task::Poll::Ready(Err(std::io::Error::other(e)));
            }
            std::task::Poll::Ready(Ok(())) => {}
        }
        match self.inner.start_send_unpin(Message::Binary(payload.into())) {
            Err(e) => std::task::Poll::Ready(Err(std::io::Error::other(e))),
            Ok(()) => std::task::Poll::Ready(Ok(4 + len)),
        }
    }
    fn poll_flush(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        self.inner
            .poll_flush_unpin(cx)
            .map_err(std::io::Error::other)
    }
    fn poll_shutdown(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        self.inner
            .poll_close_unpin(cx)
            .map_err(std::io::Error::other)
    }
}

// ---------------------------------------------------------------------------
// Downstream listener
// ---------------------------------------------------------------------------

/// Read one line from the downstream socket (up to 4 KiB).
#[cfg(unix)]
async fn read_line(stream: &mut tokio::net::UnixStream) -> Option<String> {
    let mut buf = Vec::with_capacity(128);
    let mut byte = [0u8; 1];
    loop {
        match stream.read(&mut byte).await {
            Ok(0) | Err(_) => return None,
            Ok(_) => {
                if byte[0] == b'\n' {
                    break;
                }
                buf.push(byte[0]);
                if buf.len() > 4096 {
                    return None;
                }
            }
        }
    }
    Some(
        String::from_utf8_lossy(&buf)
            .trim_end_matches('\r')
            .to_string(),
    )
}

#[cfg(unix)]
async fn run_listener(registry: Arc<Registry>, sock_path: &str) {
    let _ = std::fs::remove_file(sock_path);
    let listener = tokio::net::UnixListener::bind(sock_path).unwrap_or_else(|e| {
        eprintln!("blit-proxy: cannot bind to {sock_path}: {e}");
        std::process::exit(1);
    });
    log!("blit-proxy: listening on {sock_path}");
    loop {
        match listener.accept().await {
            Ok((stream, _)) => {
                let registry = registry.clone();
                tokio::spawn(async move {
                    handle_downstream(registry, stream).await;
                });
            }
            Err(e) => log!("blit-proxy: accept error: {e}"),
        }
    }
}

#[cfg(unix)]
async fn handle_downstream(registry: Arc<Registry>, mut downstream: tokio::net::UnixStream) {
    // Handshake: read "target <uri>\n" or "shutdown\n"
    let line = match read_line(&mut downstream).await {
        Some(l) => l,
        None => return,
    };
    if line == "shutdown" {
        let _ = downstream.write_all(b"ok\n").await;
        log!("blit-proxy: shutdown requested, exiting");
        std::process::exit(0);
    }
    let uri = match line.strip_prefix("target ") {
        Some(u) if !u.is_empty() => u.to_string(),
        _ => {
            let _ = downstream.write_all(b"error invalid handshake\n").await;
            return;
        }
    };

    let pool = registry.get_or_create(&uri).await;

    let upstream = match pool.acquire().await {
        Ok(u) => u,
        Err(e) => {
            log!("blit-proxy: [{uri}] upstream unavailable: {e}");
            let msg = format!("error {e}\n");
            let _ = downstream.write_all(msg.as_bytes()).await;
            return;
        }
    };

    if downstream.write_all(b"ok\n").await.is_err() {
        return;
    }

    pool.client_connected();

    let (mut ds_read, mut ds_write) = tokio::io::split(downstream);
    let (mut us_read, mut us_write) = (upstream.reader, upstream.writer);

    let mut d2u = tokio::spawn(async move { tokio::io::copy(&mut ds_read, &mut us_write).await });
    let mut u2d = tokio::spawn(async move { tokio::io::copy(&mut us_read, &mut ds_write).await });

    tokio::select! { _ = &mut d2u => { u2d.abort(); }, _ = &mut u2d => { d2u.abort(); } }

    pool.client_disconnected();
}

// ---------------------------------------------------------------------------
// Public entry point
// ---------------------------------------------------------------------------

/// Run the proxy on the current thread (blocks until the process exits).
///
/// Reads `BLIT_PROXY_SOCK`, `BLIT_PROXY_POOL`, and `BLIT_PROXY_IDLE` from the
/// environment. When called from within the `blit` binary instead of the
/// standalone `blit-proxy` binary, `verbose` should be `false`.
pub fn run(verbose: bool) {
    if verbose {
        VERBOSE.store(true, Ordering::Relaxed);
    }

    let sock_path = proxy_socket_path();

    let pool_size: usize = std::env::var("BLIT_PROXY_POOL")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(4);

    let idle_secs: Option<u64> = std::env::var("BLIT_PROXY_IDLE")
        .ok()
        .and_then(|v| v.parse().ok());

    tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .expect("blit-proxy: tokio runtime")
        .block_on(async move {
            rustls::crypto::ring::default_provider()
                .install_default()
                .ok(); // may already be installed by the CLI's runtime

            let registry = Registry::new(pool_size);

            // Idle-timeout watcher.
            if let Some(idle) = idle_secs {
                let reg = registry.clone();
                let sock = sock_path.clone();
                tokio::spawn(async move {
                    loop {
                        tokio::time::sleep(std::time::Duration::from_secs(5)).await;
                        let last = reg.latest_activity().await;
                        if last == i64::MAX {
                            continue;
                        }
                        let elapsed = now_secs().saturating_sub(last) as u64;
                        if elapsed >= idle {
                            log!("blit-proxy: idle for {elapsed}s (limit {idle}s), exiting");
                            let _ = std::fs::remove_file(&sock);
                            std::process::exit(0);
                        }
                    }
                });
            }

            #[cfg(unix)]
            {
                let sock_cleanup = sock_path.clone();
                tokio::spawn(async move {
                    let _ = tokio::signal::ctrl_c().await;
                    let _ = std::fs::remove_file(&sock_cleanup);
                    std::process::exit(0);
                });
                run_listener(registry, &sock_path).await;
            }
            #[cfg(not(unix))]
            {
                eprintln!("blit-proxy: Windows not yet supported");
                std::process::exit(1);
            }
        });
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn proxy_socket_path_default() {
        // Just verify it produces a non-empty string without panicking.
        let p = proxy_socket_path();
        assert!(!p.is_empty());
        assert!(p.contains("blit-proxy"));
    }

    #[test]
    fn parse_hex_valid() {
        assert_eq!(parse_hex("deadbeef"), Some(vec![0xde, 0xad, 0xbe, 0xef]));
        assert_eq!(parse_hex(""), Some(vec![]));
    }

    #[test]
    fn parse_hex_odd_length() {
        assert_eq!(parse_hex("abc"), None);
    }

    #[test]
    fn parse_hex_invalid_char() {
        assert_eq!(parse_hex("zz"), None);
    }

    #[test]
    fn extract_uri_params_no_query() {
        let (base, pass, cert) = extract_uri_params("wss://host:3264/");
        assert_eq!(base, "wss://host:3264/");
        assert_eq!(pass, None);
        assert_eq!(cert, None);
    }

    #[test]
    fn extract_uri_params_passphrase() {
        let (base, pass, cert) = extract_uri_params("wss://host:3264/?passphrase=secret");
        assert_eq!(base, "wss://host:3264/");
        assert_eq!(pass, Some("secret".into()));
        assert_eq!(cert, None);
    }

    #[test]
    fn extract_uri_params_both() {
        let (base, pass, cert) =
            extract_uri_params("wt://host:4433/?passphrase=abc&certHash=deadbeef");
        assert_eq!(base, "wt://host:4433/");
        assert_eq!(pass, Some("abc".into()));
        assert_eq!(cert, Some("deadbeef".into()));
    }

    #[test]
    fn extract_uri_params_socket_unchanged() {
        let (base, pass, cert) = extract_uri_params("socket:/tmp/blit.sock");
        assert_eq!(base, "socket:/tmp/blit.sock");
        assert_eq!(pass, None);
        assert_eq!(cert, None);
    }

    #[test]
    fn parse_wt_host_port_with_port() {
        assert_eq!(parse_wt_host_port("host:4433"), Ok(("host".into(), 4433)));
    }

    #[test]
    fn parse_wt_host_port_default() {
        assert_eq!(parse_wt_host_port("host"), Ok(("host".into(), 443)));
    }

    #[test]
    fn percent_decode_basic() {
        assert_eq!(percent_decode("hello%20world"), "hello world");
        assert_eq!(percent_decode("plain"), "plain");
    }

    #[test]
    fn parse_share_uri_no_hub() {
        let (pass, hub) = parse_share_uri("myphrase");
        assert_eq!(pass, "myphrase");
        assert_eq!(
            hub,
            blit_webrtc_forwarder::normalize_hub(blit_webrtc_forwarder::DEFAULT_HUB_URL)
        );
    }

    #[test]
    fn parse_share_uri_with_hub() {
        let (pass, hub) = parse_share_uri("myphrase?hub=wss://custom.example.com");
        assert_eq!(pass, "myphrase");
        assert_eq!(hub, "wss://custom.example.com");
    }

    #[test]
    fn parse_share_uri_hub_normalized() {
        let (pass, hub) = parse_share_uri("myphrase?hub=custom.example.com");
        assert_eq!(pass, "myphrase");
        assert_eq!(hub, "wss://custom.example.com");
    }

    #[test]
    fn parse_share_uri_percent_encoded_passphrase() {
        let (pass, hub) = parse_share_uri("my%3Fphrase");
        assert_eq!(pass, "my?phrase");
        assert_eq!(
            hub,
            blit_webrtc_forwarder::normalize_hub(blit_webrtc_forwarder::DEFAULT_HUB_URL)
        );
    }
}