featherbit/plugins/native/

upstream.rs

//! The `upstream` node — forwards the request to a backend target over HTTP,
//! with round-robin, least-connections, or IP-hash load balancing across the
//! configured targets, and writes the backend's reply into `Context.response`.

use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use bytes::Bytes;

use crate::balancer::{Balancer, Strategy, Target};
use crate::context::stream::ResponseStream;
use crate::context::{Context, GatewayError, Protocol};
use crate::outbound::idle::{body_holding, idle_timeout_body};
use crate::outbound::{OutboundClient, OutboundError, OutboundRequest};
use crate::plugins::resources::PluginResources;
use crate::plugins::{Plugin, PluginExecutionError, PluginOutput, PluginResult};

/// Proxies the request to one of the configured backend targets and populates
/// `Context.response` with the upstream's status, headers, and body.
///
/// Connection failures, request-build failures, and body-read failures are
/// returned as [`PluginExecutionError`]s so the graph engine can route them
/// through this node's error port.
pub struct UpstreamPlugin {
    /// Backend pool + load-balancing strategy (shared with the L4 stream proxy).
    /// `Arc`-wrapped so the streaming branch can take an `owned_acquire`
    /// in-flight guard that outlives this call's stack frame — it travels
    /// with the response stream, released only when the body finishes.
    balancer: Arc<Balancer>,
    /// Shared pooled HTTP client (from `PluginResources`).
    client: Arc<OutboundClient>,
    /// Whole-call deadline per proxied request.
    timeout: Duration,
    /// Idle bound on a streaming response body: no frame for this long and
    /// the stream is reaped. Only consulted when `__may_stream` is set.
    stream_idle_timeout: Duration,
    /// Connect to the upstream over TLS (`https`/`wss`); default false.
    tls: bool,
    /// Verify the upstream's TLS certificate; default true. Only meaningful
    /// when `tls` is set.
    ssl_verify: bool,
    /// Per-upstream TLS identity (client cert / private CA); None = shared
    /// clients, exactly the pre-mTLS behavior.
    tls_identity: Option<Arc<crate::outbound::tls::UpstreamTls>>,
}

// Manual `Debug`, scoped to what's useful in test/error output: `balancer`
// and `client` don't implement it (pooled hyper client, atomics-backed pool
// state), so a derive isn't available here.
impl std::fmt::Debug for UpstreamPlugin {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("UpstreamPlugin")
            .field("timeout", &self.timeout)
            .field("tls", &self.tls)
            .field("ssl_verify", &self.ssl_verify)
            .field("tls_identity_set", &self.tls_identity.is_some())
            .finish()
    }
}

impl UpstreamPlugin {
    /// Builds the plugin from node config.
    ///
    /// Accepted keys:
    /// - `targets` (array of `{host: string, port: integer}`, **required**):
    ///   the backend pool. Entries missing `host` or `port` are skipped;
    ///   an empty resulting pool is a config error.
    /// - `load_balancing` (string, default `round_robin`): one of
    ///   `round_robin`, `least_connections`, or `ip_hash`. Hyphenated and
    ///   short spellings (`round-robin`, `least-conn`) are accepted, as is
    ///   the legacy key name `load_balancer` (see [`Strategy::parse`]).
    ///
    /// - `timeout_ms` (integer, default `60000`): whole-call deadline
    ///   (connect + request + response body) per proxied request; exceeding
    ///   it fails the node with `UPSTREAM_TIMEOUT` through the error port.
    ///   For a streaming response (see `__may_stream` on
    ///   [`crate::graph::engine`]) this bounds connect + request + response
    ///   headers only — the body is then bounded by `stream_idle_timeout_ms`
    ///   instead.
    /// - `stream_idle_timeout_ms` (integer, default `60000`): only consulted
    ///   when the node is permitted to stream; no frame arriving on the
    ///   response body for this long reaps the stream. Resets on every
    ///   frame, so a steady stream survives indefinitely.
    /// - `tls` (bool, default `false`): connect to the upstream over TLS
    ///   (`https` for the buffered path, `wss` for WebSocket).
    /// - `ssl_verify` (bool, default `true`): verify the upstream's TLS
    ///   certificate. Only meaningful when `tls` is set.
    /// - `client_cert_path` / `client_key_path` (string, optional): PEM
    ///   client certificate and private key presented to the upstream for
    ///   mutual TLS. Must be set together, and only with `tls: true`.
    /// - `ca_cert_path` (string, optional): PEM CA bundle used to verify the
    ///   upstream's certificate, *replacing* the native root store for this
    ///   upstream. Requires `tls: true`; rejected together with
    ///   `ssl_verify: false` (a CA bundle to verify with is contradictory
    ///   when verification is off).
    ///
    /// Errors if no valid target is configured, if the load-balancing value
    /// is not a string, if it names an unknown strategy, if any mTLS key is
    /// set without `tls: true`, if `client_cert_path`/`client_key_path` are
    /// not both set, if `ca_cert_path` is set with `ssl_verify: false`, or if
    /// the configured cert/key/CA files can't be read or parsed.
    ///
    /// ```yaml
    /// type: upstream
    /// config:
    ///   targets:
    ///     - host: backend-1
    ///       port: 3000
    ///     - host: backend-2
    ///       port: 3000
    ///   load_balancing: least_connections
    ///   timeout_ms: 60000
    ///   tls: true
    ///   client_cert_path: /etc/gateway/client.crt
    ///   client_key_path: /etc/gateway/client.key
    ///   ca_cert_path: /etc/gateway/ca.crt
    /// ```
    pub fn from_config(
        config: &HashMap<String, serde_json::Value>,
        resources: &Arc<PluginResources>,
    ) -> Result<Self, String> {
        // Tolerant parse: entries missing `host`/`port` are silently skipped
        // (an empty resulting pool is rejected by `Balancer::new`).
        let targets = config
            .get("targets")
            .and_then(|v| v.as_array())
            .map(|seq| {
                seq.iter()
                    .filter_map(|t| {
                        let mapping = t.as_object()?;
                        let host = mapping.get("host")?.as_str()?.to_string();
                        let port = mapping.get("port")?.as_u64()? as u16;
                        Some(Target { host, port })
                    })
                    .collect::<Vec<_>>()
            })
            .unwrap_or_default();

        // `load_balancing` is the canonical key; `load_balancer` is accepted
        // because earlier UI builds saved configs under that name.
        let strategy = match config
            .get("load_balancing")
            .or_else(|| config.get("load_balancer"))
        {
            None => Strategy::default(),
            Some(v) => {
                let s = v
                    .as_str()
                    .ok_or_else(|| "load_balancing must be a string".to_string())?;
                Strategy::parse(s)?
            }
        };

        let balancer = Arc::new(Balancer::new(targets, strategy)?);

        let timeout = Duration::from_millis(
            config
                .get("timeout_ms")
                .and_then(|v| v.as_u64())
                .unwrap_or(60_000),
        );
        let stream_idle_timeout = Duration::from_millis(
            config
                .get("stream_idle_timeout_ms")
                .and_then(|v| v.as_u64())
                .unwrap_or(60_000),
        );

        let tls = config.get("tls").and_then(|v| v.as_bool()).unwrap_or(false);
        let ssl_verify = config
            .get("ssl_verify")
            .and_then(|v| v.as_bool())
            .unwrap_or(true);

        let string_config_value = |key: &str| -> Result<Option<String>, String> {
            match config.get(key) {
                None => Ok(None),
                Some(v) => match v.as_str() {
                    Some(s) => Ok(Some(s.to_string())),
                    None => Err(format!("{} must be a string", key)),
                },
            }
        };

        let client_cert_path = string_config_value("client_cert_path")?;
        let client_key_path = string_config_value("client_key_path")?;
        let ca_cert_path = string_config_value("ca_cert_path")?;

        let any_mtls_key =
            client_cert_path.is_some() || client_key_path.is_some() || ca_cert_path.is_some();
        if any_mtls_key && !tls {
            return Err(
                "client_cert_path/client_key_path/ca_cert_path require tls: true".to_string(),
            );
        }
        if client_cert_path.is_some() != client_key_path.is_some() {
            return Err("client_cert_path and client_key_path must be set together".to_string());
        }
        // A CA bundle exists to *verify* the upstream; pairing it with
        // ssl_verify:false is contradictory, so reject rather than guess.
        if ca_cert_path.is_some() && !ssl_verify {
            return Err("ca_cert_path with ssl_verify: false is contradictory".to_string());
        }

        let tls_identity = if any_mtls_key {
            let client = client_cert_path.as_deref().zip(client_key_path.as_deref());
            let identity = crate::outbound::tls::UpstreamTls::load(
                client,
                ca_cert_path.as_deref(),
                ssl_verify,
            )?;
            crate::outbound::tls::UpstreamTls::register(&identity);
            Some(identity)
        } else {
            None
        };

        Ok(Self {
            balancer,
            client: resources.outbound.clone(),
            timeout,
            stream_idle_timeout,
            tls,
            ssl_verify,
            tls_identity,
        })
    }

    /// Builds the [`OutboundRequest`] to send to `target`: method, URL,
    /// forwarded headers (Host overridden to the target), body, timeout, and
    /// TLS settings. Shared by the buffered and streaming branches of
    /// `execute` so they can't drift — every field here (ssl_verify, tls
    /// identity, timeout) affects both equally.
    fn outbound_request(
        &self,
        ctx: &Context,
        uri: String,
        method: http::Method,
        target: &Target,
    ) -> OutboundRequest {
        OutboundRequest {
            method,
            url: uri,
            headers: forwarded_headers(ctx, target),
            body: ctx.request.body.clone(),
            timeout: self.timeout,
            ssl_verify: self.ssl_verify,
            tls: self.tls_identity.clone(),
        }
    }
}

/// Forwards the request's headers to the upstream, overriding `Host` with
/// the selected target. Shared by both `execute` branches via
/// [`UpstreamPlugin::outbound_request`].
fn forwarded_headers(ctx: &Context, target: &Target) -> Vec<(String, String)> {
    let mut headers: Vec<(String, String)> = Vec::new();
    for (key, values) in &ctx.request.headers {
        if key.eq_ignore_ascii_case("host") {
            continue;
        }
        for value in values {
            headers.push((key.clone(), value.clone()));
        }
    }
    headers.push((
        "host".to_string(),
        format!("{}:{}", target.host, target.port),
    ));
    headers
}

/// The outbound request-target: the request path plus the rebuilt query
/// string when the request carried one.
///
/// `query_params` holds values exactly as received — ingress splits the raw
/// query on `&`/`=` without percent-decoding — so rebuilding is lossless.
/// Parameter order is normalized (sorted) because the original order is
/// already lost in the `HashMap` at ingress; sorting at least makes the
/// outbound target deterministic.
fn request_target(ctx: &Context) -> String {
    let query = crate::vars::query_string(ctx);
    if query.is_empty() {
        ctx.request.path.clone()
    } else {
        format!("{}?{}", ctx.request.path, query)
    }
}

/// Maps an [`OutboundError`] to the `(code, message)` pair used on this
/// node's error port. Shared by the buffered and streaming branches so a fix
/// to one (e.g. a wording or timeout-classification change) can't land in
/// only one of them.
fn map_outbound_error(e: OutboundError, target: &Target) -> (&'static str, String) {
    match &e {
        OutboundError::Timeout(d) => (
            "UPSTREAM_TIMEOUT",
            format!(
                "Upstream {}:{} timed out after {:?}",
                target.host, target.port, d
            ),
        ),
        OutboundError::InvalidRequest(m) => (
            "UPSTREAM_REQUEST_BUILD_ERROR",
            format!("Failed to build upstream request: {}", m),
        ),
        OutboundError::Transport(m) => (
            "UPSTREAM_CONNECTION_ERROR",
            format!(
                "Failed to reach upstream {}:{}: {}",
                target.host, target.port, m
            ),
        ),
    }
}

#[async_trait]
impl Plugin for UpstreamPlugin {
    fn plugin_type(&self) -> &str {
        "upstream"
    }

    async fn execute(&self, mut ctx: Context) -> PluginResult {
        let target_idx = self.balancer.select(&ctx.request.remote_addr);
        let target = self.balancer.target(target_idx);

        // WebSocket upgrade: don't do a buffered round-trip. Resolve the target
        // (load balancing works the same — selection is pure) and stash it for
        // the listener, which owns the raw connection and performs the upstream
        // handshake + bidirectional relay. Signal intent with 101. The in-flight
        // counter is intentionally skipped: a WS tunnel outlives this node, so
        // there is no round-trip lifecycle to bound it.
        if ctx.request.protocol == Protocol::WebSocket {
            // Computed before the first `ctx.message` mutable borrow below.
            let ws_target = request_target(&ctx);
            ctx.message.insert(
                "__ws_upstream_host".to_string(),
                serde_json::json!(target.host),
            );
            ctx.message.insert(
                "__ws_upstream_port".to_string(),
                serde_json::json!(target.port),
            );
            ctx.message.insert(
                "__ws_upstream_path".to_string(),
                serde_json::json!(ws_target),
            );
            ctx.message
                .insert("__ws_upstream_tls".to_string(), serde_json::json!(self.tls));
            ctx.message.insert(
                "__ws_upstream_verify".to_string(),
                serde_json::json!(self.ssl_verify),
            );
            if let Some(identity) = &self.tls_identity {
                ctx.message.insert(
                    "__ws_upstream_tls_key".to_string(),
                    serde_json::json!(identity.cache_key()),
                );
            }
            ctx.response.status_code = 101;
            return Ok(PluginOutput::success(ctx));
        }

        let scheme = if self.tls { "https" } else { "http" };
        let uri = format!(
            "{}://{}:{}{}",
            scheme,
            target.host,
            target.port,
            request_target(&ctx)
        );

        let method: http::Method = ctx.request.method.parse().unwrap_or(http::Method::GET);

        let may_stream = ctx
            .message
            .get("__may_stream")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);

        if may_stream {
            // Owned, not borrowed: this guard must outlive `execute`'s stack
            // frame — it travels with the response stream and is released
            // only when the body finishes, not when this node returns.
            let guard = self.balancer.owned_acquire(target_idx);
            let outbound = self.outbound_request(&ctx, uri, method, target);
            return match self.client.request_streaming(outbound).await {
                Ok(resp) => {
                    ctx.response.status_code = resp.status;
                    ctx.response.headers = resp.headers;
                    // Invariant: exactly one of `body`/`stream` carries
                    // content. This node is about to populate `stream`.
                    ctx.response.body = Bytes::new();
                    // The guard is bound into the body at construction, not
                    // attached after via `stream.hold(...)`: `into_parts`
                    // hands body and guards back as two independent values,
                    // so keeping them together end-to-end from the moment
                    // the stream exists removes any chance of the transport
                    // task dropping one half and not the other.
                    let idled = idle_timeout_body(resp.body, self.stream_idle_timeout);
                    let held = body_holding(idled, vec![Box::new(guard)]);
                    ctx.response.stream = Some(ResponseStream::new(held));
                    Ok(PluginOutput::success(ctx))
                }
                Err(e) => {
                    // No byte has reached the client yet (the deadline covers
                    // only connect + request + response headers), so this is
                    // an ordinary error-port exit exactly like the buffered
                    // path's failure below.
                    let (code, message) = map_outbound_error(e, target);
                    Err(PluginExecutionError {
                        context: ctx,
                        error: GatewayError {
                            node_id: String::new(),
                            code: code.to_string(),
                            message,
                            metadata: HashMap::new(),
                        },
                    })
                }
            };
        }

        let _in_flight_guard = self.balancer.acquire(target_idx);
        let outbound = self.outbound_request(&ctx, uri, method, target);

        let response = match self.client.request(outbound).await {
            Ok(resp) => resp,
            Err(e) => {
                let (code, message) = map_outbound_error(e, target);
                let error = GatewayError {
                    node_id: String::new(),
                    code: code.to_string(),
                    message,
                    metadata: HashMap::new(),
                };
                return Err(PluginExecutionError {
                    context: ctx,
                    error,
                });
            }
        };

        // Populate context.response from the upstream response
        ctx.response.status_code = response.status;
        ctx.response.headers = response.headers;
        ctx.response.body = response.body;
        // Invariant, enforced locally rather than argued globally: a
        // buffered response never carries a stream. Matters on a failover
        // policy shape where an earlier node's context (including a stale
        // `response.stream` from some prior hop) reaches this node.
        ctx.response.stream = None;

        Ok(PluginOutput::success(ctx))
    }
}

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

    fn plugin_with(strategy: Option<&str>, key: &str, n_targets: usize) -> UpstreamPlugin {
        let targets: Vec<serde_json::Value> = (0..n_targets)
            .map(|i| serde_json::json!({ "host": format!("backend-{}", i), "port": 3000 }))
            .collect();
        let mut config = HashMap::new();
        config.insert("targets".to_string(), serde_json::Value::Array(targets));
        if let Some(s) = strategy {
            config.insert(key.to_string(), serde_json::Value::String(s.to_string()));
        }
        UpstreamPlugin::from_config(&config, &PluginResources::empty()).unwrap()
    }

    #[test]
    fn test_load_balancing_parsing_and_aliases() {
        // canonical key, spec spelling
        assert_eq!(
            plugin_with(Some("least_connections"), "load_balancing", 2)
                .balancer
                .strategy(),
            Strategy::LeastConnections
        );
        // legacy UI key and hyphenated/short spellings
        assert_eq!(
            plugin_with(Some("round-robin"), "load_balancer", 2)
                .balancer
                .strategy(),
            Strategy::RoundRobin
        );
        assert_eq!(
            plugin_with(Some("least-conn"), "load_balancer", 2)
                .balancer
                .strategy(),
            Strategy::LeastConnections
        );
        assert_eq!(
            plugin_with(Some("ip_hash"), "load_balancing", 2)
                .balancer
                .strategy(),
            Strategy::IpHash
        );
        // absent -> default
        assert_eq!(
            plugin_with(None, "load_balancing", 2).balancer.strategy(),
            Strategy::RoundRobin
        );
    }

    #[test]
    fn test_load_balancing_rejects_unknown() {
        let mut config = HashMap::new();
        config.insert(
            "targets".to_string(),
            serde_json::json!([{ "host": "backend", "port": 3000 }]),
        );
        config.insert(
            "load_balancing".to_string(),
            serde_json::Value::String("random".to_string()),
        );
        assert!(UpstreamPlugin::from_config(&config, &PluginResources::empty()).is_err());
    }

    #[tokio::test]
    async fn test_websocket_branch_stashes_target_and_101() {
        use crate::context::GatewayRequest;

        let plugin = plugin_with(None, "load_balancing", 1);
        let mut req_headers = HashMap::new();
        req_headers.insert("upgrade".to_string(), vec!["websocket".to_string()]);
        let ctx = Context::new(GatewayRequest {
            method: "GET".into(),
            path: "/ws/chat".into(),
            host: "h".into(),
            scheme: "http".into(),
            headers: req_headers,
            query_params: HashMap::new(),
            body: bytes::Bytes::new(),
            remote_addr: "1.2.3.4:5".into(),
            protocol: Protocol::WebSocket,
        });

        // No target is reachable, but the WS branch must NOT do a round-trip —
        // it resolves the target and returns a 101 without any network call.
        let out = plugin.execute(ctx).await.unwrap();
        assert_eq!(out.context.response.status_code, 101);
        assert_eq!(
            out.context.message.get("__ws_upstream_host").unwrap(),
            "backend-0"
        );
        assert_eq!(out.context.message.get("__ws_upstream_port").unwrap(), 3000);
        assert_eq!(
            out.context.message.get("__ws_upstream_path").unwrap(),
            "/ws/chat"
        );
        // The in-flight counter was not touched for the WS path.
        assert_eq!(plugin.balancer.in_flight_count(0), 0);
        // TLS flags default to plaintext + verify-on.
        assert_eq!(out.context.message.get("__ws_upstream_tls").unwrap(), false);
        assert_eq!(
            out.context.message.get("__ws_upstream_verify").unwrap(),
            true
        );
    }

    #[test]
    fn test_tls_config_parses_and_defaults() {
        // Defaults: plaintext, verify on.
        let default = plugin_with(None, "load_balancing", 1);
        assert!(!default.tls);
        assert!(default.ssl_verify);

        // Explicit tls + ssl_verify:false.
        let mut config = HashMap::new();
        config.insert(
            "targets".to_string(),
            serde_json::json!([{ "host": "backend", "port": 443 }]),
        );
        config.insert("tls".to_string(), serde_json::json!(true));
        config.insert("ssl_verify".to_string(), serde_json::json!(false));
        let plugin = UpstreamPlugin::from_config(&config, &PluginResources::empty()).unwrap();
        assert!(plugin.tls);
        assert!(!plugin.ssl_verify);
    }

    fn write_identity(tag: &str) -> (String, String, String) {
        // Same helper as src/outbound/tls.rs tests: CA + leaf, PEM files in
        // temp_dir, returns (cert_path, key_path, ca_path).
        let mut ca_params = rcgen::CertificateParams::new(Vec::<String>::new()).unwrap();
        ca_params.is_ca = rcgen::IsCa::Ca(rcgen::BasicConstraints::Unconstrained);
        let ca_key = rcgen::KeyPair::generate().unwrap();
        let ca_cert = ca_params.self_signed(&ca_key).unwrap();
        let ca_issuer = rcgen::Issuer::from_ca_cert_der(ca_cert.der(), &ca_key).unwrap();
        let leaf_params = rcgen::CertificateParams::new(vec!["client".to_string()]).unwrap();
        let leaf_key = rcgen::KeyPair::generate().unwrap();
        let leaf_cert = leaf_params.signed_by(&leaf_key, &ca_issuer).unwrap();
        let dir = std::env::temp_dir();
        let pid = std::process::id();
        let cert = dir.join(format!("featherbit_up_{}_{}.crt", tag, pid));
        let key = dir.join(format!("featherbit_up_{}_{}.key", tag, pid));
        let ca = dir.join(format!("featherbit_up_{}_{}.ca.crt", tag, pid));
        std::fs::write(&cert, leaf_cert.pem()).unwrap();
        std::fs::write(&key, leaf_key.serialize_pem()).unwrap();
        std::fs::write(&ca, ca_cert.pem()).unwrap();
        (
            cert.to_str().unwrap().to_string(),
            key.to_str().unwrap().to_string(),
            ca.to_str().unwrap().to_string(),
        )
    }

    /// Base config with one target; callers add TLS keys.
    fn mtls_config() -> HashMap<String, serde_json::Value> {
        let mut config = HashMap::new();
        config.insert(
            "targets".to_string(),
            serde_json::json!([{"host": "backend", "port": 443}]),
        );
        config
    }

    #[test]
    fn test_mtls_config_requires_tls_true() {
        let (cert, key, _) = write_identity("needstls");
        let mut config = mtls_config();
        config.insert("client_cert_path".to_string(), serde_json::json!(cert));
        config.insert("client_key_path".to_string(), serde_json::json!(key));
        // tls defaults to false -> config error.
        let err = UpstreamPlugin::from_config(&config, &PluginResources::empty()).unwrap_err();
        assert!(err.contains("tls"), "err was: {}", err);
    }

    #[test]
    fn test_mtls_config_cert_and_key_must_pair() {
        let (cert, _, _) = write_identity("pair");
        let mut config = mtls_config();
        config.insert("tls".to_string(), serde_json::json!(true));
        config.insert("client_cert_path".to_string(), serde_json::json!(cert));
        let err = UpstreamPlugin::from_config(&config, &PluginResources::empty()).unwrap_err();
        assert!(err.contains("together"), "err was: {}", err);
    }

    #[test]
    fn test_mtls_config_ca_with_no_verify_rejected() {
        let (_, _, ca) = write_identity("contradiction");
        let mut config = mtls_config();
        config.insert("tls".to_string(), serde_json::json!(true));
        config.insert("ssl_verify".to_string(), serde_json::json!(false));
        config.insert("ca_cert_path".to_string(), serde_json::json!(ca));
        assert!(UpstreamPlugin::from_config(&config, &PluginResources::empty()).is_err());
    }

    #[test]
    fn test_mtls_config_loads_identity_and_registers() {
        let (cert, key, ca) = write_identity("loads");
        let mut config = mtls_config();
        config.insert("tls".to_string(), serde_json::json!(true));
        config.insert("client_cert_path".to_string(), serde_json::json!(cert));
        config.insert("client_key_path".to_string(), serde_json::json!(key));
        config.insert("ca_cert_path".to_string(), serde_json::json!(ca));
        let plugin = UpstreamPlugin::from_config(&config, &PluginResources::empty()).unwrap();
        let id = plugin.tls_identity.as_ref().expect("identity loaded");
        // Registered for the WebSocket relay to look up.
        assert!(crate::outbound::tls::UpstreamTls::lookup(id.cache_key()).is_some());
    }

    #[test]
    fn test_mtls_config_absent_means_no_identity() {
        let mut config = mtls_config();
        config.insert("tls".to_string(), serde_json::json!(true));
        let plugin = UpstreamPlugin::from_config(&config, &PluginResources::empty()).unwrap();
        assert!(plugin.tls_identity.is_none());
    }

    #[test]
    fn test_mtls_config_non_string_keys_rejected() {
        for key in ["client_cert_path", "client_key_path", "ca_cert_path"] {
            for bad_value in [
                serde_json::json!(123),
                serde_json::json!(true),
                serde_json::json!(["x"]),
            ] {
                let mut config = mtls_config();
                config.insert("tls".to_string(), serde_json::json!(true));
                config.insert(key.to_string(), bad_value.clone());
                let err =
                    UpstreamPlugin::from_config(&config, &PluginResources::empty()).unwrap_err();
                assert!(
                    err.contains(&format!("{} must be a string", key)),
                    "key {} value {:?} produced err: {}",
                    key,
                    bad_value,
                    err
                );
            }
        }
    }

    /// Minimal one-shot HTTP server that records the request line (method and
    /// request-target) of the first request it receives and answers `200`
    /// with a non-empty body — a `content-length: 0` reply would make any
    /// `body.is_empty()` assertion on the caller's side vacuous (it would
    /// pass whether or not the plugin actually populated `body`/`stream`
    /// correctly). Returns its port and a receiver for the captured line.
    async fn spawn_request_line_capture() -> (u16, tokio::sync::oneshot::Receiver<String>) {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();
        let (tx, rx) = tokio::sync::oneshot::channel();
        tokio::spawn(async move {
            if let Ok((mut stream, _)) = listener.accept().await {
                use tokio::io::{AsyncReadExt, AsyncWriteExt};
                let mut buf = [0u8; 4096];
                let n = stream.read(&mut buf).await.unwrap_or(0);
                let text = String::from_utf8_lossy(&buf[..n]).to_string();
                let line = text.lines().next().unwrap_or("").to_string();
                let _ = tx.send(line);
                let _ = stream
                    .write_all(
                        b"HTTP/1.1 200 OK
content-length: 2

ok",
                    )
                    .await;
                let _ = stream.shutdown().await;
            }
        });
        (port, rx)
    }

    fn ctx_with_query(path: &str, query: Vec<(&str, Vec<&str>)>) -> Context {
        use crate::context::GatewayRequest;
        let query_params: HashMap<String, Vec<String>> = query
            .into_iter()
            .map(|(k, vs)| {
                (
                    k.to_string(),
                    vs.into_iter().map(|v| v.to_string()).collect(),
                )
            })
            .collect();
        Context::new(GatewayRequest {
            method: "GET".into(),
            path: path.into(),
            host: "h".into(),
            scheme: "http".into(),
            headers: HashMap::new(),
            query_params,
            body: bytes::Bytes::new(),
            remote_addr: "1.2.3.4:5".into(),
            protocol: Protocol::Http1,
        })
    }

    fn plugin_at(port: u16) -> UpstreamPlugin {
        let mut config = HashMap::new();
        config.insert(
            "targets".to_string(),
            serde_json::json!([{ "host": "127.0.0.1", "port": port }]),
        );
        UpstreamPlugin::from_config(&config, &PluginResources::empty()).unwrap()
    }

    /// Regression: the outbound request-target must carry the query string.
    /// Building the URL from `ctx.request.path` alone silently dropped it on
    /// every proxied call — an OIDC authorize hop reached the IdP with no
    /// `client_id`, which the IdP reports as "parameter not present".
    #[tokio::test]
    async fn test_query_string_is_forwarded_to_upstream() {
        let (port, rx) = spawn_request_line_capture().await;
        let ctx = ctx_with_query(
            "/realms/example/protocol/openid-connect/auth",
            vec![("client_id", vec!["apisix"])],
        );

        plugin_at(port).execute(ctx).await.unwrap();

        let request_line = rx.await.unwrap();
        assert!(
            request_line.contains("client_id=apisix"),
            "query string dropped from outbound request-target: {request_line}"
        );
    }

    /// Regression: the WebSocket relay target must keep the query string too.
    /// `__ws_upstream_path` is consumed verbatim by the relay in
    /// `server::listener`, so dropping the query there breaks token-in-query
    /// upgrades (`wss://host/ws?token=...`) exactly as it broke plain HTTP.
    #[tokio::test]
    async fn test_websocket_upstream_path_keeps_query_string() {
        use crate::context::GatewayRequest;

        let mut req_headers = HashMap::new();
        req_headers.insert("upgrade".to_string(), vec!["websocket".to_string()]);
        let mut query_params = HashMap::new();
        query_params.insert("token".to_string(), vec!["abc123".to_string()]);
        let ctx = Context::new(GatewayRequest {
            method: "GET".into(),
            path: "/ws/chat".into(),
            host: "h".into(),
            scheme: "http".into(),
            headers: req_headers,
            query_params,
            body: bytes::Bytes::new(),
            remote_addr: "1.2.3.4:5".into(),
            protocol: Protocol::WebSocket,
        });

        let out = plugin_with(None, "load_balancing", 1)
            .execute(ctx)
            .await
            .unwrap();

        assert_eq!(
            out.context.message.get("__ws_upstream_path").unwrap(),
            "/ws/chat?token=abc123"
        );
    }

    /// With `__may_stream` set, the node must hand back a stream rather than a
    /// buffered body — and must leave `body` empty, per the invariant.
    #[tokio::test]
    async fn test_upstream_streams_when_permitted() {
        let (port, _rx) = spawn_request_line_capture().await;
        let mut ctx = ctx_with_query("/stream", vec![]);
        ctx.message
            .insert("__may_stream".to_string(), serde_json::json!(true));

        let out = plugin_at(port).execute(ctx).await.unwrap();

        assert!(out.context.response.stream.is_some(), "expected a stream");
        assert!(
            out.context.response.body.is_empty(),
            "invariant: body must be empty when stream is set"
        );
    }

    /// Without the key, behaviour is exactly as today: buffered body, no stream.
    #[tokio::test]
    async fn test_upstream_buffers_when_not_permitted() {
        let (port, _rx) = spawn_request_line_capture().await;
        let ctx = ctx_with_query("/stream", vec![]);

        let out = plugin_at(port).execute(ctx).await.unwrap();

        assert!(out.context.response.stream.is_none());
        assert_eq!(
            out.context.response.body.as_ref(),
            b"ok",
            "buffered path must actually carry the upstream's body"
        );
    }

    /// Sends response headers plus a first chunk immediately, then blocks
    /// until `resume` is signaled before sending the final chunk and the
    /// chunked terminator and closing — giving a test a window to observe
    /// the balancer's in-flight count while the stream is still open, before
    /// deciding when the body is allowed to finish.
    async fn spawn_pausable_stream_server() -> (u16, tokio::sync::oneshot::Sender<()>) {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();
        let (resume_tx, resume_rx) = tokio::sync::oneshot::channel();
        tokio::spawn(async move {
            if let Ok((mut stream, _)) = listener.accept().await {
                use tokio::io::{AsyncReadExt, AsyncWriteExt};
                let mut buf = [0u8; 4096];
                let _ = stream.read(&mut buf).await;
                let _ = stream
                    .write_all(
                        b"HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n5\r\nhello\r\n",
                    )
                    .await;
                let _ = resume_rx.await;
                let _ = stream.write_all(b"6\r\nworld!\r\n0\r\n\r\n").await;
                let _ = stream.shutdown().await;
            }
        });
        (port, resume_tx)
    }

    /// The balancer's in-flight guard must be bound to the streaming body's
    /// own lifetime (via `body_holding`), not to this node's `execute` call.
    /// Swapping that binding for a bare `drop(guard)` — the exact bug this
    /// guards against — would release the count the instant `execute`
    /// returns, long before the client has actually received the whole
    /// body; the whole rest of the test suite stays green either way, which
    /// is what makes this property worth testing directly rather than
    /// trusting it stayed wired correctly.
    #[tokio::test]
    async fn test_in_flight_guard_released_only_when_stream_completes() {
        use http_body_util::BodyExt;

        let (port, resume_tx) = spawn_pausable_stream_server().await;
        let plugin = plugin_at(port);
        let mut ctx = ctx_with_query("/stream", vec![]);
        ctx.message
            .insert("__may_stream".to_string(), serde_json::json!(true));

        let out = plugin.execute(ctx).await.unwrap();
        let stream = out
            .context
            .response
            .stream
            .expect("expected a stream when __may_stream is set");

        assert_eq!(
            plugin.balancer.in_flight_count(0),
            1,
            "in-flight count must stay held while the stream is still open"
        );

        let (body, _guards) = stream.into_parts();
        let _ = resume_tx.send(());
        let collected = body.collect().await.unwrap().to_bytes();
        assert_eq!(collected.as_ref(), b"helloworld!");

        assert_eq!(
            plugin.balancer.in_flight_count(0),
            0,
            "in-flight count must release once the stream body completes"
        );
    }
}