zc2 0.0.29

//! `ANY /serve/:deployment_id/*path` — the broker's reverse proxy into a
//! locally-reconciled deployment's container, addressed by the container's
//! own network IP (`LocalState`'s `ip`/`port`, resolved by
//! `reconcile::drive` via `ContainerRuntime::container_ip` — see
//! `runtime`'s module doc comment for why: a broker driving this reconciler
//! is usually itself a container, so a host-published port is unreachable
//! from it). Every proxied 2xx response bills the caller via
//! `Ledger::record_serve_transaction` — fire-and-forget, off the response
//! path, exactly like the WAL flush worker keeps fsync off `/execute`'s.

use std::sync::Arc;
use std::time::Duration;

use crate::broker::http_adapter::{BrokerRequest, BrokerResponse};
use crate::broker::BrokerState;

use super::state::LocalState;

/// Hop-by-hop headers that must not be blindly forwarded (RFC 7230 §6.1),
/// plus `host`, which must be re-derived for the upstream request rather than
/// carried over from the inbound one.
const HOP_BY_HOP: &[&str] = &[
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
    "host",
];

/// Split `/serve/:deployment_id/*path` into `(deployment_id, "/*path")`. The
/// remainder always starts with `/` (empty tail → `/`), matching what a
/// proxied container serving at its own root expects. `None` when the path
/// doesn't start with `/serve/` or names no deployment id.
pub fn parse_serve_path(path: &str) -> Option<(&str, String)> {
    let rest = path.strip_prefix("/serve/")?;
    let (id, tail) = rest.split_once('/').unwrap_or((rest, ""));
    if id.is_empty() {
        return None;
    }
    Some((id, format!("/{tail}")))
}

/// `duration_ms/1000 * price_per_second`, rounded to 6 decimals per the
/// billing contract — credits are a currency figure, and an unrounded f64
/// product carries binary-float noise (e.g. `0.1000000000000001`) into the
/// ledger. Pure — unit tested without spinning up billing/threads.
pub fn compute_credits_amount(duration_ms: f64, price_per_second: f64) -> f64 {
    ((duration_ms / 1000.0) * price_per_second * 1_000_000.0).round() / 1_000_000.0
}

/// The upstream URL this proxy sends a request to, given the deployment's
/// container IP/port (`LocalState`'s `ip`/`port` — see this module's doc
/// comment) and the tail path from `parse_serve_path`. Pure — unit tested
/// without a real container.
pub fn upstream_url(host: &str, port: u16, tail_path: &str, query: &str) -> String {
    if query.is_empty() {
        format!("http://{host}:{port}{tail_path}")
    } else {
        format!("http://{host}:{port}{tail_path}?{query}")
    }
}

/// Map the transport-agnostic `tiny_http::Method` `BrokerRequest` carries
/// onto an `http::Method`, so the proxied request can be built as a plain
/// `http::Request` and run on the broker's shared `ureq::Agent`
/// (`state.http_client`) via `Agent::run` — the one `ureq::Agent` entry point
/// that isn't limited to a fixed verb, unlike `agent.get(..)`/`.post(..)`.
fn to_http_method(method: &tiny_http::Method) -> http::Method {
    use tiny_http::Method as M;
    match method {
        M::Get => http::Method::GET,
        M::Head => http::Method::HEAD,
        M::Post => http::Method::POST,
        M::Put => http::Method::PUT,
        M::Delete => http::Method::DELETE,
        M::Connect => http::Method::CONNECT,
        M::Options => http::Method::OPTIONS,
        M::Trace => http::Method::TRACE,
        M::Patch => http::Method::PATCH,
        M::NonStandard(s) => {
            http::Method::from_bytes(s.as_str().as_bytes()).unwrap_or(http::Method::GET)
        }
    }
}

/// Same auth `/execute` applies (see `execute_prepare` in `server.rs`): a
/// non-empty Bearer token must resolve via `resolve_user_from_api_key`; in
/// local mode a missing token is allowed (free, single-node); otherwise it's
/// a 401. Kept as its own small copy rather than calling into `server.rs`'s
/// private `execute_prepare` — that function does far more than auth
/// (routing, WAL reserve, credit checks) and inlines the check rather than
/// exposing it, and duplicating four lines here is safer than widening that
/// hot path's surface for a caller it wasn't written for.
fn authenticate(
    state: &BrokerState,
    request: &BrokerRequest,
) -> Result<Option<String>, BrokerResponse> {
    let bearer = request
        .header("Authorization")
        .and_then(|a| a.strip_prefix("Bearer ").map(|t| t.trim().to_string()))
        .filter(|t| !t.is_empty());

    if let Some(token) = bearer {
        state
            .ledger
            .resolve_user_from_api_key(&token)
            .map(Some)
            .map_err(|e| error_response(&e.to_string(), 401))
    } else if state.is_local_mode() {
        Ok(None)
    } else {
        Err(error_response(
            "API key required. Set Authorization: Bearer <key>",
            401,
        ))
    }
}

fn error_response(message: &str, status: u16) -> BrokerResponse {
    let body = serde_json::json!({ "error": message }).to_string();
    BrokerResponse::json_bytes(body.into_bytes(), status)
}

/// Handle `ANY /serve/:deployment_id/*path`. Called from the broker's generic
/// fallback dispatcher (`http_adapter::handle`, itself run inside
/// `spawn_blocking` — see `server.rs`), so a blocking `ureq` round-trip here
/// is exactly as safe as every other synchronous handler in that dispatcher.
pub fn handle_serve(state: &Arc<BrokerState>, request: &BrokerRequest) -> BrokerResponse {
    let Some((deployment_id, tail_path)) = parse_serve_path(&request.path) else {
        return error_response("not a /serve/:deployment_id/*path request", 404);
    };

    let caller = match authenticate(state, request) {
        Ok(caller) => caller,
        Err(resp) => return resp,
    };

    let local = LocalState::load();
    let Some(record) = local.deployments.get(deployment_id) else {
        return error_response("deployment not found on this node", 502);
    };
    if record.phase != "healthy" {
        return error_response("deployment is not healthy on this node", 502);
    }
    let Some(ip) = record.ip.as_deref() else {
        return error_response(
            "deployment has no known container address on this node",
            502,
        );
    };
    let Some(port) = record.port else {
        return error_response("deployment has no known container port on this node", 502);
    };
    let price_per_second = record.price_per_second;

    let url = upstream_url(ip, port, &tail_path, &request.query);

    let mut builder = http::Request::builder()
        .method(to_http_method(&request.method))
        .uri(&url);
    for (k, v) in &request.headers {
        if HOP_BY_HOP.iter().any(|h| k.eq_ignore_ascii_case(h)) {
            continue;
        }
        builder = builder.header(k.as_str(), v.as_str());
    }
    let http_request = match builder.body(request.body.clone()) {
        Ok(r) => r,
        Err(e) => return error_response(&format!("bad proxied request: {e}"), 502),
    };
    let http_request = state
        .http_client
        .configure_request(http_request)
        .timeout_global(Some(Duration::from_secs(60)))
        .build();

    // Billing is per execution TIME (credits/second), not per call — see
    // `bill_served_request` — so the wall clock starts here, around the
    // actual proxied round-trip, not the auth/lookup work above it.
    let started = std::time::Instant::now();
    let response = match state.http_client.run(http_request) {
        Ok(resp) => resp,
        Err(e) => {
            eprintln!("  [SERVE] proxy to {deployment_id} failed: {e}");
            return error_response("upstream deployment unreachable", 502);
        }
    };

    let status = response.status().as_u16();
    let headers: Vec<(String, String)> = response
        .headers()
        .iter()
        .filter(|(k, _)| {
            !HOP_BY_HOP
                .iter()
                .any(|h| k.as_str().eq_ignore_ascii_case(h))
        })
        .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string()))
        .collect();
    let mut body_reader = response.into_body();
    let body = body_reader.read_to_vec().unwrap_or_default();
    let duration_ms = started.elapsed().as_secs_f64() * 1000.0;

    if (200..300).contains(&status) {
        bill_served_request(state, &caller, deployment_id, price_per_second, duration_ms);
    } else {
        eprintln!(
            "  [SERVE] {deployment_id} returned {status} (not billed), duration_ms={duration_ms:.1}"
        );
    }

    BrokerResponse {
        status,
        headers,
        body,
    }
}

/// Fire-and-forget billing for one verified (2xx) served request, by
/// execution time — spawned off the response path so a slow/unreachable
/// dashboard never adds latency to what the caller actually gets back.
fn bill_served_request(
    state: &Arc<BrokerState>,
    caller: &Option<String>,
    deployment_id: &str,
    price_per_second: f64,
    duration_ms: f64,
) {
    let user_id = caller.clone().unwrap_or_else(|| "anonymous".to_string());
    let state = state.clone();
    let deployment_id = deployment_id.to_string();
    let credits_amount = compute_credits_amount(duration_ms, price_per_second);
    std::thread::spawn(move || {
        let request_id = uuid::Uuid::new_v4().to_string();
        let job_name = deployment_id.clone();
        let source_node = state.config.node_name.clone();
        state.ledger.record_serve_transaction(
            &request_id,
            &user_id,
            &deployment_id,
            &job_name,
            credits_amount,
            duration_ms,
            source_node.as_deref(),
        );
    });
}

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

    #[test]
    fn compute_credits_amount_multiplies_seconds_by_price() {
        // 2500ms at 0.0001 credits/s = 2.5s * 0.0001 = 0.00025
        assert_eq!(compute_credits_amount(2500.0, 0.0001), 0.00025);
    }

    #[test]
    fn compute_credits_amount_rounds_to_six_decimals() {
        // 333ms at 0.001 credits/s = 0.333 * 0.001 = 0.000333, well inside
        // 6 decimals — pick a duration that WOULD carry float noise past
        // that without the rounding.
        let amount = compute_credits_amount(333.333, 0.001);
        assert_eq!(amount, (amount * 1_000_000.0).round() / 1_000_000.0);
    }

    #[test]
    fn compute_credits_amount_zero_duration_is_free() {
        assert_eq!(compute_credits_amount(0.0, 0.0001), 0.0);
    }

    #[test]
    fn parse_serve_path_splits_id_and_tail() {
        assert_eq!(
            parse_serve_path("/serve/dep_1/foo/bar"),
            Some(("dep_1", "/foo/bar".to_string()))
        );
    }

    #[test]
    fn parse_serve_path_defaults_tail_to_root() {
        assert_eq!(
            parse_serve_path("/serve/dep_1"),
            Some(("dep_1", "/".to_string()))
        );
        assert_eq!(
            parse_serve_path("/serve/dep_1/"),
            Some(("dep_1", "/".to_string()))
        );
    }

    #[test]
    fn parse_serve_path_rejects_non_serve_paths() {
        assert_eq!(parse_serve_path("/execute"), None);
        assert_eq!(parse_serve_path("/serve/"), None);
        assert_eq!(parse_serve_path("/serve"), None);
    }

    #[test]
    fn upstream_url_appends_query_when_present() {
        assert_eq!(
            upstream_url("172.17.0.5", 8000, "/foo", ""),
            "http://172.17.0.5:8000/foo"
        );
        assert_eq!(
            upstream_url("172.17.0.5", 8000, "/foo", "a=1&b=2"),
            "http://172.17.0.5:8000/foo?a=1&b=2"
        );
    }

    #[test]
    fn upstream_url_targets_the_container_ip_not_loopback() {
        // The whole point of the drill fix: a broker running as its own
        // container can't reach anything published to the docker HOST's
        // 127.0.0.1 — this must always address the container's own IP.
        let url = upstream_url("10.13.13.10", 8000, "/", "");
        assert!(url.starts_with("http://10.13.13.10:"));
        assert!(!url.contains("127.0.0.1"));
    }

    #[test]
    fn to_http_method_maps_common_verbs() {
        assert_eq!(to_http_method(&tiny_http::Method::Get), http::Method::GET);
        assert_eq!(to_http_method(&tiny_http::Method::Post), http::Method::POST);
        assert_eq!(
            to_http_method(&tiny_http::Method::Delete),
            http::Method::DELETE
        );
    }
}