keyhog_verifier/oob/

client.rs

//! Low-level interactsh protocol client.
//!
//! A thin async wrapper around the projectdiscovery/interactsh-server register/
//! poll/deregister endpoints. Stateless aside from the RSA keypair, secret,
//! correlation id, and HTTP client - `OobSession` (in `session.rs`) layers
//! the per-finding subscription, polling loop, and notification fan-out on top.
//!
//! ## Crypto invariants
//!
//! - RSA-2048, OAEP padding, SHA-256 hash and MGF - interactsh-server speaks
//!   exactly this combination; `RSA_PKCS1_OAEP_PADDING` with SHA-256 in their
//!   Go code. Other parameters won't decrypt.
//! - AES-256-CFB with a 16-byte IV prepended to ciphertext. Each interaction
//!   carries an independent IV; the AES key is per-poll-batch.
//! - We never log credentials, public keys, or decrypted payloads. Errors
//!   carry stable strings - useful for support, opaque to leaks.

use std::sync::{LazyLock, Mutex};
use std::time::Duration;

use base64::{engine::general_purpose::STANDARD as B64, Engine as _};
use rand::{rngs::OsRng, Rng};
use reqwest::Client;
use rsa::pkcs8::{EncodePublicKey, LineEnding};
use rsa::{Oaep, RsaPrivateKey, RsaPublicKey};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use thiserror::Error;
use tracing::{debug, warn};

/// Stable bucket name for the global rate limiter. Every OOB call across
/// every detector shares this bucket so the aggregate request rate to the
/// upstream collector never exceeds the configured `--verify-rate`. Using
/// the literal string `"oob.interactsh"` (not the server URL) means the
/// budget covers all configured collectors collectively - the limit is
/// about our own machine not blasting traffic, not about per-host fairness.
const OOB_SERVICE: &str = "oob.interactsh";
const OOB_LIFECYCLE_BURST: usize = 3;
const DNS_TOKEN_ALPHABET: &[u8; 36] = b"abcdefghijklmnopqrstuvwxyz0123456789";
const CORRELATION_ID_LEN: usize = 24;
const UNIQUE_SUFFIX_LEN: usize = 24;

/// All errors that can arise from the OOB client. `Transient` errors mean the
/// caller should retry (network blip, rate-limit); everything else is final.
#[derive(Debug, Error)]
pub enum InteractshError {
    #[error("interactsh keypair generation failed: {0}")]
    KeyGen(String),
    #[error("interactsh public-key encoding failed: {0}")]
    KeyEncode(String),
    #[error("interactsh register failed (HTTP {status}): {body}")]
    Register { status: u16, body: String },
    #[error("interactsh deregister failed (HTTP {status}): {body}")]
    Deregister { status: u16, body: String },
    #[error("interactsh poll failed (HTTP {status}): {body}")]
    Poll { status: u16, body: String },
    #[error("interactsh response shape unexpected: {0}")]
    BadResponse(String),
    #[error("interactsh collector host blocked by SSRF guard: {0}")]
    BlockedCollector(String),
    #[error("interactsh AES key unwrap failed: {0}")]
    AesUnwrap(String),
    #[error("interactsh interaction decrypt failed: {0}")]
    Decrypt(String),
    #[error("interactsh transport error: {0}")]
    Transport(#[from] reqwest::Error),
    #[error("interactsh request timed out after {0:?}")]
    Timeout(Duration),
}

/// Protocol category of a received interaction.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum InteractionProtocol {
    Dns,
    Http,
    Smtp,
    Other,
}

impl InteractionProtocol {
    // `#[doc(hidden)] pub` rather than `pub(super)`: the OOB protocol-string
    // parser is exercised directly by the boundary test
    // `oob_interaction_protocol_parse_exact`. Hidden from the rendered API
    // it is an internal categorizer, not a semver-covered surface.
    #[doc(hidden)]
    pub fn parse(s: &str) -> Self {
        match s.to_ascii_lowercase().as_str() {
            "dns" => Self::Dns,
            "http" => Self::Http,
            "smtp" | "smtp-mail" => Self::Smtp,
            _ => Self::Other,
        }
    }
}

/// One decrypted interaction returned by the collector.
#[derive(Debug, Clone)]
pub struct Interaction {
    /// Full unique id (correlation-id || per-finding suffix). This is
    /// what we match against per-finding URLs we minted.
    pub unique_id: String,
    pub protocol: InteractionProtocol,
    pub remote_address: String,
    pub timestamp: String,
    /// Raw protocol payload (HTTP request line + headers, DNS query, etc.).
    /// Sized - interactsh truncates server-side, but we cap to 16 KiB here as
    /// a defense-in-depth budget against memory abuse from a hostile server.
    pub raw_payload: String,
}

/// One interactsh registration. Cheap to clone (Arc-friendly fields only on
/// caller's side; here we hold owned values because the session pins this
/// for the lifetime of the engine).
pub struct InteractshClient {
    /// Collector client after the OOB SSRF/DNS policy is applied. On direct
    /// connections this client pins the collector host to the screened DNS
    /// answers; with an explicit proxy it is the caller-provided proxy client
    /// after the string-level private-host block has run.
    http: Client,
    server: String,
    correlation_id: String,
    secret_key: String,
    private_key: RsaPrivateKey,
    /// Length of the per-URL suffix. Production uses 24 DNS-safe characters so
    /// a known correlation id still leaves >122 bits of per-finding entropy.
    suffix_len: usize,
}

/// One shared 2048-bit RSA key for every `for_test` client. Generated ONCE
/// (lazily, on first test use) and cloned into each instance, so the many test
/// callers of [`InteractshClient::for_test`] pay a SINGLE 2048-bit keygen
/// instead of one per call, full weak-crypto hygiene (NIST-minimum modulus)
/// without the per-test keygen cost that a naive `RsaPrivateKey::new(_, 2048)`
/// in the constructor would impose. Never initialized in a production build
/// (`for_test` is test-only), so it costs nothing there.
static TEST_RSA_KEY: LazyLock<Result<RsaPrivateKey, String>> =
    LazyLock::new(|| RsaPrivateKey::new(&mut OsRng, 2048).map_err(|e| e.to_string()));
type GeneratedPrivateKey = Result<RsaPrivateKey, String>;

static PREWARMED_RSA_KEY: LazyLock<Mutex<Option<std::thread::JoinHandle<GeneratedPrivateKey>>>> =
    LazyLock::new(|| Mutex::new(None));

/// Start the OOB session key generation before findings reach verification.
///
/// Registration consumes this one-shot key. A second registration still
/// generates a fresh key, preserving the session-key isolation contract.
pub fn prewarm_key_generation() {
    let mut slot = PREWARMED_RSA_KEY
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner); // LAW10: poisoned one-shot key-slot recovery preserves the stored join handle; registration still returns key-generation failures.
    if slot.is_some() {
        return;
    }
    match std::thread::Builder::new()
        .name("keyhog-oob-keygen".to_owned())
        .spawn(generate_private_key)
    {
        Ok(handle) => *slot = Some(handle),
        Err(error) => {
            debug!(target: "keyhog::oob", %error, "could not prewarm OOB key generation");
        }
    }
}

fn generate_private_key() -> GeneratedPrivateKey {
    RsaPrivateKey::new(&mut OsRng, 2048).map_err(|error| error.to_string())
}

fn take_or_generate_private_key() -> Result<RsaPrivateKey, InteractshError> {
    let prewarmed = PREWARMED_RSA_KEY
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner) // LAW10: poisoned one-shot key-slot recovery takes the same handle; join and key-generation failures remain explicit.
        .take();
    let generated = match prewarmed {
        Some(handle) => handle
            .join()
            .map_err(|_| InteractshError::KeyGen("prewarm thread panicked".to_owned()))?,
        None => generate_private_key(),
    };
    generated.map_err(InteractshError::KeyGen)
}
pub(crate) fn prewarmed_key_pending_for_test() -> bool {
    PREWARMED_RSA_KEY
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner) // LAW10: test-only state inspection recovers the one-shot slot without affecting verifier behavior.
        .is_some()
}

pub(crate) fn consume_prewarmed_key_for_test() -> Result<Vec<u8>, InteractshError> {
    use rsa::traits::PublicKeyParts;

    let prewarmed = PREWARMED_RSA_KEY
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner) // LAW10: test-only key consumption recovers the slot while missing, panic, and generation failures remain explicit.
        .take()
        .ok_or_else(|| InteractshError::KeyGen("prewarmed key missing".to_owned()))?;
    let key = prewarmed
        .join()
        .map_err(|_| InteractshError::KeyGen("prewarm thread panicked".to_owned()))?
        .map_err(InteractshError::KeyGen)?;
    Ok(key.n().to_bytes_be())
}

impl InteractshClient {
    /// Test-only constructor without network registration. Returns
    /// `Err` if the shared test RSA keygen failed - which never happens on a
    /// healthy platform, but propagating the error keeps this constructor
    /// off the no-panic-in-production gate and matches the rest of the
    /// `InteractshError` surface. Test callers wrap with `.unwrap()` at
    /// the test boundary.
    pub(crate) fn for_test(server: &str) -> Result<Self, InteractshError> {
        // Clone the shared 2048-bit key (see `TEST_RSA_KEY`): NIST-minimum
        // modulus, one keygen amortized across every test caller.
        let private_key = TEST_RSA_KEY
            .as_ref()
            .map_err(|e| InteractshError::KeyGen(e.clone()))?
            .clone();
        Ok(Self {
            http: Client::new(),
            server: normalize_server(server),
            correlation_id: "abcdefghijklmnopqrstuvwx".to_string(),
            secret_key: "test-secret".to_string(),
            private_key,
            suffix_len: UNIQUE_SUFFIX_LEN,
        })
    }
}

/// JSON shapes from interactsh-server. Field names match the upstream Go
/// definitions (`pkg/server/types.go`). `serde(default)` keeps us forward-
/// compatible with future fields.
#[derive(Serialize)]
struct RegisterRequest<'a> {
    #[serde(rename = "public-key")]
    public_key: &'a str,
    #[serde(rename = "secret-key")]
    secret_key: &'a str,
    #[serde(rename = "correlation-id")]
    correlation_id: &'a str,
}

#[derive(Deserialize, Default)]
#[serde(default)]
struct PollResponse {
    /// Each entry is base64( AES-256-CFB( IV[16] || ciphertext ) ).
    data: Vec<String>,
    /// Auxiliary metadata; ignored.
    #[serde(rename = "extra")]
    _extra: Vec<String>,
    /// Base64( RSA-OAEP-SHA256( 32-byte AES key ) ). Server omits when there
    /// are no interactions; in that case `data` is also empty.
    aes_key: Option<String>,
}

/// Decrypted interaction shape. `serde(default)` because interactsh-server
/// sometimes ships partial events (failed protocol parse, etc.) and we'd
/// rather degrade gracefully than 500.
/// Hard cap on the body of a `/poll` response. Protects the process from a
/// hostile or misbehaving collector returning a multi-gigabyte JSON that
/// would force `serde_json::from_slice` to allocate the whole thing
/// in-memory before we can validate it. 4 MiB comfortably fits any
/// reasonable poll batch - see the rationale at the call site.
const MAX_POLL_BODY_BYTES: usize = 4 * 1024 * 1024;

/// Cap on error/diagnostic bodies. We only display the first 256 chars in
/// the error message anyway, but the cap prevents a server returning a
/// 500 with a 1 GiB body from spiking memory.
const ERROR_BODY_CAP: usize = 64 * 1024;

/// Stream a response body into a Vec under a hard byte cap. Returns
/// `BadResponse` if the cap is exceeded - abort the read rather than
/// trust the server's framing.
async fn read_capped_bytes(
    resp: reqwest::Response,
    cap: usize,
) -> Result<Vec<u8>, InteractshError> {
    use futures_util::StreamExt;
    let mut stream = resp.bytes_stream();
    let mut buf: Vec<u8> = Vec::new();
    while let Some(chunk) = stream.next().await {
        let chunk = chunk.map_err(InteractshError::Transport)?;
        if buf.len().saturating_add(chunk.len()) > cap {
            return Err(InteractshError::BadResponse(format!(
                "response body exceeds {cap}-byte cap"
            )));
        }
        buf.extend_from_slice(&chunk);
    }
    Ok(buf)
}

/// Like `read_capped_bytes` but for diagnostic error messages - never
/// returns `Err`; on a stream failure or cap breach it returns whatever
/// was buffered so the error log can still surface something.
async fn read_capped_text(resp: reqwest::Response, cap: usize) -> String {
    use futures_util::StreamExt;
    let mut stream = resp.bytes_stream();
    let mut buf: Vec<u8> = Vec::new();
    while let Some(chunk) = stream.next().await {
        let Ok(chunk) = chunk else { break };
        if buf.len().saturating_add(chunk.len()) > cap {
            break;
        }
        buf.extend_from_slice(&chunk);
    }
    String::from_utf8_lossy(&buf).into_owned()
}

impl InteractshClient {
    /// Build, generate keys, and register with the collector. The returned
    /// client is ready to mint URLs and be polled.
    pub async fn register(http: Client, server: &str) -> Result<Self, InteractshError> {
        Self::register_with_network_policy(http, server, Duration::from_secs(30), false, false)
            .await
    }

    pub(crate) async fn register_with_network_policy(
        http: Client,
        server: &str,
        timeout: Duration,
        proxy_in_use: bool,
        insecure_tls: bool,
    ) -> Result<Self, InteractshError> {
        // RSA-2048 keygen runs before verification when the CLI can overlap it
        // with scanning. Direct library callers retain the blocking-pool path.
        let private_key = tokio::task::spawn_blocking(take_or_generate_private_key)
            .await
            .map_err(|error| InteractshError::KeyGen(format!("join error: {error}")))??;

        let public_key = RsaPublicKey::from(&private_key);
        let pem = public_key
            .to_public_key_pem(LineEnding::LF)
            .map_err(|e| InteractshError::KeyEncode(e.to_string()))?;
        let public_key_b64 = B64.encode(pem.as_bytes());

        // Correlation id is 24 lowercase alphanumerics - interactsh-server
        // matches incoming subdomains by this prefix, so the ID space must
        // be wide enough that collisions are statistically impossible across
        // every concurrent scanner sharing the collector. 36^24 ≈ 2.2e37.
        let correlation_id = random_dns_token(CORRELATION_ID_LEN);
        let secret_key = uuid::Uuid::new_v4().to_string();

        let server = normalize_server(server);
        let collector_http =
            collector_http_client(&http, &server, timeout, proxy_in_use, insecure_tls).await?;

        let body = RegisterRequest {
            public_key: &public_key_b64,
            secret_key: &secret_key,
            correlation_id: &correlation_id,
        };
        // One OOB session has a bounded three-request lifecycle: register,
        // poll, and deregister. A three-token burst avoids serial startup
        // sleeps while the shared bucket still enforces the configured
        // sustained collector rate for every later poll.
        crate::rate_limit::get_rate_limiter()
            .wait_with_burst(OOB_SERVICE, OOB_LIFECYCLE_BURST)
            .await;
        let resp = collector_http
            .post(format!("{server}/register"))
            .json(&body)
            .send()
            .await?;
        let status = resp.status();
        if !status.is_success() {
            let body = read_capped_text(resp, ERROR_BODY_CAP).await;
            return Err(InteractshError::Register {
                status: status.as_u16(),
                body: body.chars().take(256).collect(),
            });
        }
        // Drain (and discard) the register success body under a cap. Some
        // interactsh deployments echo registration metadata; we don't need
        // it but must not let the connection sit half-read indefinitely.
        let _ = read_capped_bytes(resp, ERROR_BODY_CAP).await; // LAW10: unused-binding marker; no runtime effect, not a fallback
        debug!(target: "keyhog::oob", correlation_id = %correlation_id, server = %server, "registered with interactsh collector");

        Ok(Self {
            http: collector_http,
            server,
            correlation_id,
            secret_key,
            private_key,
            suffix_len: UNIQUE_SUFFIX_LEN,
        })
    }

    /// Mint a fresh callback URL bound to this session. The full unique-id
    /// subdomain is returned (unique-id) plus the host the service should
    /// hit. Caller is responsible for embedding it where the credential's
    /// API will follow.
    pub(crate) fn mint_url(&self) -> MintedUrl {
        let suffix = random_dns_token(self.suffix_len);
        let unique_id = format!("{}{}", self.correlation_id, suffix);
        let host = format!("{}.{}", unique_id, self.server_host());
        let url = format!("https://{host}");
        MintedUrl {
            unique_id,
            host,
            url,
        }
    }

    /// Poll once. Returns every interaction the collector has buffered for
    /// this correlation id since the last poll.
    pub async fn poll(&self) -> Result<Vec<Interaction>, InteractshError> {
        // The bounded lifecycle burst is shared with register and deregister;
        // subsequent polls remain spaced at the configured sustained rate.
        crate::rate_limit::get_rate_limiter()
            .wait_with_burst(OOB_SERVICE, OOB_LIFECYCLE_BURST)
            .await;
        let resp = self
            .http
            .get(format!("{}/poll", self.server))
            .query(&[("id", &self.correlation_id), ("secret", &self.secret_key)])
            .send()
            .await?;
        let status = resp.status();
        if !status.is_success() {
            let body = read_capped_text(resp, ERROR_BODY_CAP).await;
            return Err(InteractshError::Poll {
                status: status.as_u16(),
                body: body.chars().take(256).collect(),
            });
        }
        // Bound the response body before deserialization. A malicious or
        // misbehaving collector could otherwise blow process memory by
        // returning a multi-gigabyte JSON. 4 MiB comfortably fits even a
        // dense poll batch (≤100 interactions × ~16 KiB raw_payload each
        // base64-expanded ≈ 2 MiB) with headroom.
        let body = read_capped_bytes(resp, MAX_POLL_BODY_BYTES).await?;
        let parsed: PollResponse = serde_json::from_slice(&body)
            .map_err(|e| InteractshError::BadResponse(e.to_string()))?;
        if parsed.data.is_empty() {
            return Ok(Vec::new());
        }
        let aes_key_b64 = parsed.aes_key.ok_or_else(|| {
            InteractshError::BadResponse("data present but aes_key missing".into())
        })?;
        let aes_key = self.unwrap_aes_key(&aes_key_b64)?;
        if aes_key.len() != 32 {
            return Err(InteractshError::AesUnwrap(format!(
                "expected 32-byte AES-256 key, got {}",
                aes_key.len()
            )));
        }

        let mut out = Vec::with_capacity(parsed.data.len());
        for entry in parsed.data {
            match super::decrypt::decrypt_entry(&aes_key, &entry) {
                Ok(Some(interaction)) => out.push(interaction),
                Ok(None) => {} // decrypt_entry already warned for the dropped interaction
                Err(e) => {
                    warn!(target: "keyhog::oob", error = %e, "interactsh entry decrypt failed; skipping")
                }
            }
        }
        Ok(out)
    }

    /// Tear down the registration. Idempotent on the server side; a failure
    /// to deregister is non-fatal - the server prunes inactive sessions
    /// after its retention window.
    pub async fn deregister(&self) -> Result<(), InteractshError> {
        #[derive(Serialize)]
        struct DeregisterRequest<'a> {
            #[serde(rename = "correlation-id")]
            correlation_id: &'a str,
            #[serde(rename = "secret-key")]
            secret_key: &'a str,
        }
        // See `register` for the bounded lifecycle-burst rationale.
        crate::rate_limit::get_rate_limiter()
            .wait_with_burst(OOB_SERVICE, OOB_LIFECYCLE_BURST)
            .await;
        let resp = self
            .http
            .post(format!("{}/deregister", self.server))
            .json(&DeregisterRequest {
                correlation_id: &self.correlation_id,
                secret_key: &self.secret_key,
            })
            .send()
            .await?;
        let status = resp.status();
        if !status.is_success() {
            // Cap the diagnostic body exactly like register/poll. An uncapped
            // `resp.text()` here let a hostile or misbehaving collector force an
            // unbounded allocation by returning a multi-GiB body on a
            // deregister-failure status (the one error path that previously
            // skipped the shared `read_capped_text` budget). We only display
            // the first 256 chars in the error anyway.
            let body: String = read_capped_text(resp, ERROR_BODY_CAP)
                .await
                .chars()
                .take(256)
                .collect();
            warn!(target: "keyhog::oob", status = %status, body = %body, "interactsh deregister failed");
            return Err(InteractshError::Deregister {
                status: status.as_u16(),
                body,
            });
        }
        // Drain (and discard) the success body under a cap, exactly like
        // `register`. Some interactsh deployments echo deregister metadata;
        // leaving the connection half-read can poison it for the next pooled
        // request on the same host.
        let _ = read_capped_bytes(resp, ERROR_BODY_CAP).await; // LAW10: unused-binding marker; no runtime effect, not a fallback
        Ok(())
    }

    pub(crate) fn correlation_id(&self) -> &str {
        &self.correlation_id
    }

    /// `oast.fun` from `https://oast.fun/`.
    fn server_host(&self) -> &str {
        // strip scheme; we normalized at register time so no path component.
        self.server
            .split_once("://")
            .map(|(_, rest)| rest)
            .unwrap_or(&self.server) // LAW10: absent name/label => display default; reporting-only, recall-safe
            .trim_end_matches('/')
    }

    fn unwrap_aes_key(&self, b64: &str) -> Result<Vec<u8>, InteractshError> {
        let wrapped = B64
            .decode(b64.as_bytes())
            .map_err(|e| InteractshError::AesUnwrap(format!("base64: {e}")))?;
        let padding = Oaep::new::<Sha256>();
        self.private_key
            .decrypt(padding, &wrapped)
            .map_err(|e| InteractshError::AesUnwrap(format!("rsa-oaep: {e}")))
    }
}

fn random_dns_token(len: usize) -> String {
    let mut rng = OsRng;
    (0..len)
        .map(|_| {
            let idx = rng.gen_range(0..DNS_TOKEN_ALPHABET.len());
            DNS_TOKEN_ALPHABET[idx] as char
        })
        .collect()
}

/// One per-finding callback URL, returned from `InteractshClient::mint_url`.
#[derive(Debug, Clone)]
pub struct MintedUrl {
    /// Full id; the value the service will reflect in DNS/HTTP host.
    pub unique_id: String,
    /// `<unique_id>.<server-host>` - bare host without scheme.
    pub host: String,
    /// `https://<host>` - convenience for HTTP-shaped probes.
    pub url: String,
}

/// Build the only HTTP client OOB collector traffic may use.
///
/// Direct connections mirror `resolved_client_for_url`: block private-looking
/// collector URLs, resolve once, reject any private resolved address, then pin
/// the accepted addresses into reqwest via `resolve_to_addrs`. Register, poll,
/// and deregister all use this stored client, so a collector host cannot pass a
/// first DNS screen and later rebind the unattended poller to an internal
/// service that receives the session secret.
///
/// With an explicit proxy, DNS resolution belongs to the proxy. We still run
/// the string-level private-host block locally, then keep the caller-provided
/// proxy client instead of rebuilding a direct client that would drop proxy
/// policy.
async fn collector_http_client(
    base_client: &Client,
    server: &str,
    timeout: Duration,
    proxy_in_use: bool,
    insecure_tls: bool,
) -> Result<Client, InteractshError> {
    // String-level block first: refuse a private/loopback/link-local *literal*
    // (or an unparseable/non-http(s)) collector URL before spending a DNS
    // lookup on it.
    if crate::ssrf::is_private_url(server) {
        return Err(InteractshError::BlockedCollector(format!(
            "{server} resolves to a private/loopback/link-local address"
        )));
    }

    // Resolve, then screen the collector's IPs on BOTH the direct and proxied
    // paths via the ONE decision owner below. Previously the proxied path
    // returned the caller's client BEFORE any DNS screen, so a collector host
    // that resolved to an internal address slipped past the string block and
    // the proxy forwarded the session secret to it (proxy-SSRF / DNS
    // rebinding). `collector_client_plan` screens first, then decides.
    let (host, host_port) = collector_host_and_port(server)?;
    let resolved = crate::ssrf::resolve_dns_cached(&host_port).await;

    match collector_client_plan(server, proxy_in_use, resolved)? {
        // DNS belongs to the proxy, so we cannot pin addresses into a proxied
        // client; the local screen above already rejected an internal resolve.
        // Keep the caller's proxy client after the screen.
        CollectorClientPlan::UseProxy => Ok(base_client.clone()),
        // ONE owner for the pinned rebuild, identical posture to the per-request
        // verify client; see `crate::build_pinned_verifier_client`.
        CollectorClientPlan::Pin(pinned_addrs) => {
            crate::build_pinned_verifier_client(&host, &pinned_addrs, timeout, insecure_tls)
                .map_err(|error| {
                    InteractshError::BlockedCollector(format!(
                        "{server} DNS pin client build failed ({error}); refusing an unpinned collector client"
                    ))
                })
        }
    }
}

/// Which client the OOB collector policy permits, after the resolved-IP screen.
enum CollectorClientPlan {
    /// Proxy in use: screen passed; reuse the caller's proxy client (DNS is the
    /// proxy's job; we cannot pin addresses into a proxied client).
    UseProxy,
    /// Direct connection: screen passed; pin these screened addresses.
    Pin(Vec<std::net::SocketAddr>),
}

/// ONE owner for the collector resolved-IP screen + client decision, applied
/// identically on the direct and proxied paths so neither can forward the
/// session secret to a host that resolved to an internal address.
fn collector_client_plan(
    server: &str,
    proxy_in_use: bool,
    resolved: std::io::Result<Vec<std::net::SocketAddr>>,
) -> Result<CollectorClientPlan, InteractshError> {
    let addrs = resolved.map_err(|error| collector_dns_failure(server, error))?;
    check_collector_resolved_addrs(server, &addrs)?;
    if proxy_in_use {
        Ok(CollectorClientPlan::UseProxy)
    } else {
        Ok(CollectorClientPlan::Pin(addrs))
    }
}

pub(crate) fn ssrf_check_collector_dns_result_for_test(
    server: &str,
    resolved: std::io::Result<Vec<std::net::SocketAddr>>,
) -> Result<(), InteractshError> {
    let _host_port = collector_host_and_port(server)?;
    collector_client_plan(server, false, resolved).map(|_plan| ())
}

/// Test seam for the proxy-aware screen decision: returns `true` when the plan
/// reuses the proxy client, `false` when it pins a direct client, and `Err`
/// when the screen rejects the resolved addresses. Proves the proxied path
/// screens resolved IPs (a rebinding host resolving to an internal address is
/// rejected even with `proxy_in_use = true`).
pub(crate) fn collector_reuses_proxy_client_for_test(
    server: &str,
    proxy_in_use: bool,
    resolved: std::io::Result<Vec<std::net::SocketAddr>>,
) -> Result<bool, InteractshError> {
    match collector_client_plan(server, proxy_in_use, resolved)? {
        CollectorClientPlan::UseProxy => Ok(true),
        CollectorClientPlan::Pin(_) => Ok(false),
    }
}

fn collector_host_and_port(server: &str) -> Result<(String, String), InteractshError> {
    let url = url::Url::parse(server).map_err(|_error| {
        InteractshError::BlockedCollector(format!("{server} is not a parseable collector URL"))
    })?;
    let host = url.host_str().ok_or_else(|| {
        InteractshError::BlockedCollector(format!("{server} has no collector host"))
    })?;
    let port = url
        .port_or_known_default()
        .unwrap_or(crate::DEFAULT_HTTPS_PORT); // LAW10: no explicit port => scheme default; recall-irrelevant
    Ok((host.to_string(), format!("{host}:{port}")))
}

fn collector_dns_failure(server: &str, error: std::io::Error) -> InteractshError {
    InteractshError::BlockedCollector(format!(
        "{server} DNS resolution failed before SSRF screening: {error}; collector was not contacted"
    ))
}

fn check_collector_resolved_addrs(
    server: &str,
    addrs: &[std::net::SocketAddr],
) -> Result<(), InteractshError> {
    if addrs.is_empty() {
        return Err(InteractshError::BlockedCollector(format!(
            "{server} DNS returned no addresses before SSRF screening; collector was not contacted"
        )));
    }
    if addrs
        .iter()
        .any(|addr| crate::ssrf::is_private_ip_addr(&addr.ip()))
    {
        return Err(InteractshError::BlockedCollector(format!(
            "{server} resolves to a private/loopback/link-local address"
        )));
    }
    Ok(())
}

/// Accept `oast.fun`, `oast.fun/`, `https://oast.fun`, `https://oast.fun/`.
/// Always return `https://<host>[:<port>]` with scheme/host/port ONLY. HTTP-only
/// is force-upgraded because the AES key flowing back must travel TLS-wrapped.
///
/// Keeping only scheme/host/port is load-bearing for host safety: a collector
/// string carrying a path (`oast.fun/evil`) or userinfo (`oast.fun@internal`)
/// would otherwise survive into `server_host()` and mint a malformed
/// `<id>.oast.fun/evil` callback host, or, worse, let the userinfo `@`
/// redirect the real connect target. We re-serialize from a parsed URL so the
/// stored `server` is exactly `https://<host>[:<port>]`.
///
/// An unparseable / hostless input is returned scheme-forced but otherwise
/// untouched; it is not silently "cleaned" into something connectable, the
/// downstream `is_private_url` / `collector_host_and_port` screens then reject
/// it (fail closed).
fn normalize_server(s: &str) -> String {
    let s = s.trim();
    // Force a scheme so `url::Url::parse` can split host/port; force https so we
    // never speak plaintext to a collector (the wrapped AES key would leak).
    let with_scheme = if let Some(rest) = s.strip_prefix("http://") {
        format!("https://{rest}")
    } else if s.starts_with("https://") {
        s.to_string()
    } else {
        format!("https://{s}")
    };
    match url::Url::parse(&with_scheme) {
        Ok(url) => match url.host_str() {
            // `host_str()` already excludes userinfo/path; `port()` is `None`
            // for the scheme default (443), which we then omit.
            Some(host) => match url.port() {
                Some(port) => format!("https://{host}:{port}"),
                None => format!("https://{host}"),
            },
            None => with_scheme.trim_end_matches('/').to_string(),
        },
        Err(_invalid_collector) => with_scheme.trim_end_matches('/').to_string(),
    }
}