weevil/tests/proxy_integration.rs

//! Integration tests for proxy support.
//!
//! Architecture per test:
//!
//!   weevil (or ProxyConfig::client())
//!       │  --proxy http://127.0.0.1:<proxy_port>
//!       ▼
//!   TestProxy          ← real forwarding proxy, hyper HTTP/1.1 server
//!       │  forwards absolute-form URI to origin
//!       ▼
//!   mockito origin     ← fake dl.google.com / github.com, returns canned bytes
//!
//! This proves traffic actually traverses the proxy, not just that the download
//! works.  The TestProxy struct is the only custom code; everything else is
//! standard mockito + assert_cmd.

use std::convert::Infallible;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Mutex;

// Tests that mutate HTTPS_PROXY / HTTP_PROXY must not run concurrently —
// those env vars are process-global.  cargo test runs tests in parallel
// within a single binary by default, so we serialize access with this lock.
// Tests that don't touch env vars (or that only use --proxy flag) skip it.
static ENV_MUTEX: Mutex<()> = Mutex::new(());

use hyper::body::Bytes;
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper_util::client::legacy::Client;
use hyper_util::client::legacy::connect::HttpConnector;
use hyper_util::rt::{TokioExecutor, TokioIo};
use http_body_util::{BodyExt, Full};
use tokio::net::TcpListener;

use weevil::sdk::proxy::ProxyConfig;

// ---------------------------------------------------------------------------
// TestProxy — a minimal HTTP forward proxy for use in tests.
//
// Binds to 127.0.0.1:0 (OS picks the port), spawns a tokio task to serve
// connections, and shuts down when dropped.  Also counts how many requests
// it forwarded so tests can assert the proxy was actually hit.
// ---------------------------------------------------------------------------

/// A live forwarding proxy bound to a random local port.
struct TestProxy {
    addr: SocketAddr,
    request_count: Arc<AtomicU64>,
    /// Dropping this handle shuts the server down.
    _shutdown: tokio::sync::mpsc::Sender<()>,
}

/// The actual proxy handler, extracted to a named async fn.
///
/// You cannot put a return-type annotation on an `async move { }` block, and
/// you cannot use `-> impl Future` on a closure inside `service_fn` because
/// the opaque future type is unnameable.  A named `async fn` is the one place
/// Rust lets you write both `async` and an explicit return type in the same
/// spot — the compiler knows the concrete future type and can verify the
/// `Into<Box<dyn Error>>` bound that `serve_connection` requires.
async fn proxy_handler(
    client: Client<HttpConnector, Full<Bytes>>,
    counter: Arc<AtomicU64>,
    req: hyper::Request<hyper::body::Incoming>,
) -> Result<hyper::Response<Full<Bytes>>, Infallible> {
    counter.fetch_add(1, Ordering::Relaxed);

    // The client sends the request in absolute-form:
    //   GET http://origin:PORT/path HTTP/1.1
    // hyper parses that into req.uri() for us.
    let uri = req.uri().clone();

    // Collect the incoming body so we can forward it.
    let body_bytes = req.into_body().collect().await
        .map(|b| b.to_bytes())
        .unwrap_or_default();

    let forwarded = hyper::Request::builder()
        .method("GET")
        .uri(uri)
        .body(Full::new(body_bytes))
        .unwrap();

    match client.request(forwarded).await {
        Ok(upstream_resp) => {
            // Collect the upstream body and re-wrap so we return a concrete
            // body type (not Incoming).
            let status  = upstream_resp.status();
            let headers = upstream_resp.headers().clone();
            let collected = upstream_resp.into_body().collect().await
                .map(|b| b.to_bytes())
                .unwrap_or_default();

            let mut resp = hyper::Response::builder()
                .status(status)
                .body(Full::new(collected))
                .unwrap();
            *resp.headers_mut() = headers;
            Ok(resp)
        }
        Err(_) => Ok(
            hyper::Response::builder()
                .status(502)
                .body(Full::new(Bytes::from("Bad Gateway")))
                .unwrap()
        ),
    }
}

impl TestProxy {
    async fn start() -> Self {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let count = Arc::new(AtomicU64::new(0));
        let count_clone = count.clone();

        let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel::<()>(1);

        let hyper_client: Client<HttpConnector, Full<Bytes>> =
            Client::builder(TokioExecutor::new()).build_http();

        tokio::spawn(async move {
            loop {
                let stream = tokio::select! {
                    Ok((stream, _peer)) = listener.accept() => stream,
                    _ = shutdown_rx.recv() => break,
                    else => break,
                };

                let client = hyper_client.clone();
                let counter = count_clone.clone();

                tokio::spawn(async move {
                    let io = TokioIo::new(stream);
                    // The closure just delegates to proxy_handler.  The named
                    // fn's return type is visible to the compiler so it can
                    // resolve the Error associated type that serve_connection
                    // needs.
                    let _ = http1::Builder::new()
                        .serve_connection(io, service_fn(move |req| {
                            proxy_handler(client.clone(), counter.clone(), req)
                        }))
                        .await;
                });
            }
        });

        Self {
            addr,
            request_count: count,
            _shutdown: shutdown_tx,
        }
    }

    fn proxy_url(&self) -> String {
        format!("http://{}", self.addr)
    }

    fn requests_forwarded(&self) -> u64 {
        self.request_count.load(Ordering::Relaxed)
    }
}

// ---------------------------------------------------------------------------
// Tests: ProxyConfig::client() talking through TestProxy to mockito.
//
// reqwest's blocking::Client owns an internal tokio Runtime.  Dropping it
// inside an async context panics on tokio ≥ 1.49 ("Cannot drop a runtime in
// a context where blocking is not allowed").  Each test therefore does its
// reqwest work — client creation, HTTP calls, and drop — inside
// spawn_blocking, which runs on a thread-pool thread where blocking is fine.
// The mockito origin and TestProxy stay in the async body because they need
// the tokio runtime for their listeners.
// ---------------------------------------------------------------------------

#[tokio::test]
async fn proxy_forwards_request_to_origin() {
    let mut origin = mockito::Server::new_async().await;
    let _mock = origin
        .mock("GET", "/sdk.zip")
        .with_status(200)
        .with_header("content-type", "application/octet-stream")
        .with_body(b"fake-sdk-bytes".as_slice())
        .create_async()
        .await;

    let proxy = TestProxy::start().await;
    let proxy_url = proxy.proxy_url();
    let origin_url = origin.url();

    let (status, body) = tokio::task::spawn_blocking(move || {
        let config = ProxyConfig::resolve(Some(&proxy_url), false).unwrap();
        let client = config.client().unwrap();
        let resp = client.get(format!("{}/sdk.zip", origin_url)).send().unwrap();
        (resp.status().as_u16(), resp.text().unwrap())
    }).await.unwrap();

    assert_eq!(status, 200);
    assert_eq!(body, "fake-sdk-bytes");
    assert_eq!(proxy.requests_forwarded(), 1);
}

#[tokio::test]
async fn no_proxy_bypasses_proxy_entirely() {
    let _env_lock = ENV_MUTEX.lock().unwrap();

    let mut origin = mockito::Server::new_async().await;
    let _mock = origin
        .mock("GET", "/direct.txt")
        .with_status(200)
        .with_body("direct-hit")
        .create_async()
        .await;

    let proxy = TestProxy::start().await;
    let proxy_url = proxy.proxy_url();
    let origin_url = origin.url();

    let (status, body) = tokio::task::spawn_blocking(move || {
        std::env::set_var("HTTPS_PROXY", &proxy_url);
        let config = ProxyConfig::resolve(None, true).unwrap();
        std::env::remove_var("HTTPS_PROXY");
        let client = config.client().unwrap();
        let resp = client.get(format!("{}/direct.txt", origin_url)).send().unwrap();
        (resp.status().as_u16(), resp.text().unwrap())
    }).await.unwrap();

    assert_eq!(status, 200);
    assert_eq!(body, "direct-hit");
    assert_eq!(proxy.requests_forwarded(), 0);
}

#[tokio::test]
async fn env_var_proxy_is_picked_up() {
    let _env_lock = ENV_MUTEX.lock().unwrap();

    let mut origin = mockito::Server::new_async().await;
    let _mock = origin
        .mock("GET", "/env.txt")
        .with_status(200)
        .with_body("via-env-proxy")
        .create_async()
        .await;

    let proxy = TestProxy::start().await;
    let proxy_url = proxy.proxy_url();
    let origin_url = origin.url();

    let (status, body) = tokio::task::spawn_blocking(move || {
        std::env::remove_var("HTTP_PROXY");
        std::env::remove_var("http_proxy");
        std::env::remove_var("https_proxy");
        std::env::set_var("HTTPS_PROXY", &proxy_url);
        let config = ProxyConfig::resolve(None, false).unwrap();
        std::env::remove_var("HTTPS_PROXY");
        let client = config.client().unwrap();
        let resp = client.get(format!("{}/env.txt", origin_url)).send().unwrap();
        (resp.status().as_u16(), resp.text().unwrap())
    }).await.unwrap();

    assert_eq!(status, 200);
    assert_eq!(body, "via-env-proxy");
    assert_eq!(proxy.requests_forwarded(), 1);
}

#[tokio::test]
async fn explicit_proxy_flag_overrides_env() {
    let _env_lock = ENV_MUTEX.lock().unwrap();

    let mut origin = mockito::Server::new_async().await;
    let _mock = origin
        .mock("GET", "/override.txt")
        .with_status(200)
        .with_body("flag-proxy-wins")
        .create_async()
        .await;

    let decoy = TestProxy::start().await;
    let real  = TestProxy::start().await;
    let decoy_url = decoy.proxy_url();
    let real_url  = real.proxy_url();
    let origin_url = origin.url();

    let (status, body) = tokio::task::spawn_blocking(move || {
        std::env::set_var("HTTPS_PROXY", &decoy_url);
        let config = ProxyConfig::resolve(Some(&real_url), false).unwrap();
        std::env::remove_var("HTTPS_PROXY");
        let client = config.client().unwrap();
        let resp = client.get(format!("{}/override.txt", origin_url)).send().unwrap();
        (resp.status().as_u16(), resp.text().unwrap())
    }).await.unwrap();

    assert_eq!(status, 200);
    assert_eq!(body, "flag-proxy-wins");
    assert_eq!(real.requests_forwarded(),  1);
    assert_eq!(decoy.requests_forwarded(), 0);
}

#[tokio::test]
async fn proxy_returns_502_when_origin_is_unreachable() {
    let proxy = TestProxy::start().await;
    let proxy_url = proxy.proxy_url();

    let status = tokio::task::spawn_blocking(move || {
        let config = ProxyConfig::resolve(Some(&proxy_url), false).unwrap();
        let client = config.client().unwrap();
        let resp = client.get("http://127.0.0.1:1/unreachable").send().unwrap();
        resp.status().as_u16()
    }).await.unwrap();

    assert_eq!(status, 502);
    assert_eq!(proxy.requests_forwarded(), 1);
}

#[tokio::test]
async fn multiple_sequential_requests_all_forwarded() {
    let mut origin = mockito::Server::new_async().await;
    let _m1 = origin.mock("GET", "/a").with_status(200).with_body("aaa").create_async().await;
    let _m2 = origin.mock("GET", "/b").with_status(200).with_body("bbb").create_async().await;
    let _m3 = origin.mock("GET", "/c").with_status(200).with_body("ccc").create_async().await;

    let proxy = TestProxy::start().await;
    let proxy_url = proxy.proxy_url();
    let origin_url = origin.url();

    let (a, b, c) = tokio::task::spawn_blocking(move || {
        let config = ProxyConfig::resolve(Some(&proxy_url), false).unwrap();
        let client = config.client().unwrap();
        let a = client.get(format!("{}/a", origin_url)).send().unwrap().text().unwrap();
        let b = client.get(format!("{}/b", origin_url)).send().unwrap().text().unwrap();
        let c = client.get(format!("{}/c", origin_url)).send().unwrap().text().unwrap();
        (a, b, c)
    }).await.unwrap();

    assert_eq!(a, "aaa");
    assert_eq!(b, "bbb");
    assert_eq!(c, "ccc");
    assert_eq!(proxy.requests_forwarded(), 3);
}

// ---------------------------------------------------------------------------
// CLI-level tests
// ---------------------------------------------------------------------------

#[allow(deprecated)] // cargo_bin_cmd! requires assert_cmd ≥ 2.2; we're on 2.1.2
#[test]
fn cli_help_shows_proxy_flags() {
    let mut cmd = assert_cmd::Command::cargo_bin("weevil").unwrap();
    let assert = cmd.arg("--help").assert();
    assert
        .success()
        .stdout(predicates::prelude::predicate::str::contains("--proxy"))
        .stdout(predicates::prelude::predicate::str::contains("--no-proxy"));
}

#[allow(deprecated)]
#[test]
fn cli_rejects_garbage_proxy_url() {
    let mut cmd = assert_cmd::Command::cargo_bin("weevil").unwrap();
    let assert = cmd
        .arg("--proxy")
        .arg("not-a-url-at-all")
        .arg("sdk")
        .arg("install")
        .assert();
    assert
        .failure()
        .stderr(predicates::prelude::predicate::str::contains("Invalid --proxy URL"));
}

#[allow(deprecated)]
#[test]
fn cli_proxy_and_no_proxy_are_mutually_exclusive_in_effect() {
    let mut cmd = assert_cmd::Command::cargo_bin("weevil").unwrap();
    let out = cmd
        .arg("--proxy")
        .arg("http://127.0.0.1:9999")
        .arg("--no-proxy")
        .arg("sdk")
        .arg("install")
        .output()
        .expect("weevil binary failed to execute");

    let stderr = String::from_utf8_lossy(&out.stderr);
    assert!(!stderr.contains("panic"), "binary panicked: {}", stderr);
}