crank/crates/crank-adapter-rest/tests/integration/external_reference_fetcher.rs

use std::{
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
    time::Duration,
};

use axum::{Router, http::StatusCode, response::Redirect, routing::get};
use crank_adapter_rest::{
    ExternalReferenceFetchError, ExternalReferenceFetcher, OutboundHttpPolicy,
};
use tokio::net::TcpListener;

#[tokio::test]
async fn external_references_are_default_off_before_any_request() {
    let requests = Arc::new(AtomicUsize::new(0));
    let base_url = spawn_server(Arc::clone(&requests)).await;
    let fetcher = ExternalReferenceFetcher::try_new(
        OutboundHttpPolicy::allowing_hosts(["127.0.0.1"]),
        Vec::new(),
        1024,
        Duration::from_secs(1),
    )
    .unwrap();

    let error = fetcher
        .get(&format!("{base_url}/document"))
        .await
        .unwrap_err();

    assert!(matches!(error, ExternalReferenceFetchError::Disabled));
    assert_eq!(requests.load(Ordering::SeqCst), 0);
}

#[tokio::test]
async fn fetcher_uses_prefix_and_actual_address_policy_then_returns_bounded_bytes() {
    let base_url = spawn_server(Arc::new(AtomicUsize::new(0))).await;
    let fetcher = ExternalReferenceFetcher::try_new(
        OutboundHttpPolicy::allowing_hosts(["127.0.0.1"]),
        vec![base_url.clone()],
        8,
        Duration::from_secs(1),
    )
    .unwrap();

    assert_eq!(
        fetcher
            .get(&format!("{base_url}/document#/components/schemas/A"))
            .await
            .unwrap(),
        b"openapi".to_vec()
    );
    let error = fetcher.get(&format!("{base_url}/large")).await.unwrap_err();
    assert!(matches!(
        error,
        ExternalReferenceFetchError::ResponseTooLarge { limit_bytes: 8 }
    ));

    let private_without_explicit_outbound_allow = ExternalReferenceFetcher::try_new(
        OutboundHttpPolicy::default(),
        vec![base_url.clone()],
        1024,
        Duration::from_secs(1),
    )
    .unwrap();
    let error = private_without_explicit_outbound_allow
        .get(&format!("{base_url}/document"))
        .await
        .unwrap_err();
    assert!(matches!(
        error,
        ExternalReferenceFetchError::TargetNotAllowed
    ));

    let exact_path_fetcher = ExternalReferenceFetcher::try_new(
        OutboundHttpPolicy::allowing_hosts(["127.0.0.1"]),
        vec![format!("{base_url}/document")],
        1024,
        Duration::from_secs(1),
    )
    .unwrap();
    let error = exact_path_fetcher
        .get(&format!("{base_url}/document-unrelated"))
        .await
        .unwrap_err();
    assert!(matches!(
        error,
        ExternalReferenceFetchError::TargetNotAllowed
    ));
}

#[tokio::test]
async fn fetcher_rejects_redirects_and_userinfo_without_exposing_the_url() {
    let requests = Arc::new(AtomicUsize::new(0));
    let base_url = spawn_server(Arc::clone(&requests)).await;
    let fetcher = ExternalReferenceFetcher::try_new(
        OutboundHttpPolicy::allowing_hosts(["127.0.0.1"]),
        vec![base_url.clone()],
        1024,
        Duration::from_secs(1),
    )
    .unwrap();

    let error = fetcher
        .get(&format!("{base_url}/redirect"))
        .await
        .unwrap_err();
    assert!(matches!(
        error,
        ExternalReferenceFetchError::RedirectNotAllowed
    ));

    let userinfo_url = base_url.replacen("http://", "http://user:credential@", 1);
    let error = fetcher
        .get(&format!("{userinfo_url}/document"))
        .await
        .unwrap_err();
    assert!(matches!(
        error,
        ExternalReferenceFetchError::TargetNotAllowed
    ));
    let rendered = format!("{error:?} {error}");
    assert!(!rendered.contains("credential"));
    assert_eq!(requests.load(Ordering::SeqCst), 0);
}

async fn spawn_server(requests: Arc<AtomicUsize>) -> String {
    let app = Router::new()
        .route(
            "/document",
            get({
                let requests = Arc::clone(&requests);
                move || {
                    let requests = Arc::clone(&requests);
                    async move {
                        requests.fetch_add(1, Ordering::SeqCst);
                        (StatusCode::OK, "openapi")
                    }
                }
            }),
        )
        .route("/large", get(|| async { "response too large" }))
        .route("/redirect", get(|| async { Redirect::to("/document") }));
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let address = listener.local_addr().unwrap();
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });
    format!("http://{address}")
}