folk-plugin-http 0.3.2

use std::collections::HashMap;
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Instant;

use anyhow::Result;
use axum::Router;
use axum::body::Body;
use axum::extract::{ConnectInfo, State};
use axum::http::{self, Request, Response};
use axum::routing::any;
use ipnet::IpNet;
use tokio::net::TcpListener;
use tokio::sync::watch;
use tower::ServiceBuilder;
use tower_http::limit::RequestBodyLimitLayer;
use tower_http::timeout::TimeoutLayer;
use tracing::{error, info};

use crate::config::HttpConfig;
use crate::hooks::{HookEngine, HookResult, RequestContext, ResponseContext};
use crate::payload::encode_request;

#[derive(Clone)]
struct AppState {
    executor: Arc<dyn folk_api::Executor>,
    config: Arc<HttpConfig>,
    active_connections: Arc<AtomicU64>,
    hook_engine: Option<Arc<HookEngine>>,
}

pub struct HttpServer {
    config: HttpConfig,
    executor: Arc<dyn folk_api::Executor>,
    active_connections: Arc<AtomicU64>,
    hook_engine: Option<Arc<HookEngine>>,
}

impl HttpServer {
    pub fn new(
        config: HttpConfig,
        executor: Arc<dyn folk_api::Executor>,
        active_connections: Arc<AtomicU64>,
        hook_engine: Option<Arc<HookEngine>>,
    ) -> Self {
        Self {
            config,
            executor,
            active_connections,
            hook_engine,
        }
    }

    pub async fn run(self, shutdown: watch::Receiver<bool>) -> Result<()> {
        let state = AppState {
            executor: self.executor.clone(),
            config: Arc::new(self.config.clone()),
            active_connections: self.active_connections.clone(),
            hook_engine: self.hook_engine.clone(),
        };

        let mut app = Router::new()
            .route("/{*path}", any(handle))
            .route("/", any(handle))
            .with_state(state)
            .layer(
                ServiceBuilder::new()
                    .layer(RequestBodyLimitLayer::new(self.config.max_request_size))
                    .layer(TimeoutLayer::with_status_code(
                        http::StatusCode::GATEWAY_TIMEOUT,
                        self.config.write_timeout,
                    )),
            );

        if self.config.compression.enabled {
            app = app.layer(build_compression_layer(&self.config.compression));
        }

        #[cfg(feature = "tls")]
        if let Some(ref tls) = self.config.tls {
            return self.run_tls(app, tls, shutdown).await;
        }

        #[cfg(feature = "h2c")]
        if self.config.h2c {
            return self.run_h2c(app, shutdown).await;
        }

        self.run_plain(app, shutdown).await
    }

    async fn run_plain(&self, app: Router, shutdown: watch::Receiver<bool>) -> Result<()> {
        let listener = TcpListener::bind(self.config.listen).await?;

        axum::serve(
            listener,
            app.into_make_service_with_connect_info::<SocketAddr>(),
        )
        .with_graceful_shutdown(shutdown_signal(shutdown))
        .await?;

        Ok(())
    }

    #[cfg(feature = "tls")]
    async fn run_tls(
        &self,
        app: Router,
        tls: &crate::config::TlsConfig,
        shutdown: watch::Receiver<bool>,
    ) -> Result<()> {
        use axum_server::Handle;
        use axum_server::tls_rustls::RustlsConfig;

        let rustls_config = RustlsConfig::from_pem_file(&tls.cert, &tls.key).await?;

        info!(cert = %tls.cert.display(), "TLS enabled");

        let handle = Handle::new();
        let shutdown_handle = handle.clone();
        tokio::spawn(async move {
            shutdown_signal(shutdown).await;
            shutdown_handle.graceful_shutdown(None);
        });

        axum_server::bind_rustls(self.config.listen, rustls_config)
            .handle(handle)
            .serve(app.into_make_service_with_connect_info::<SocketAddr>())
            .await?;

        Ok(())
    }

    #[cfg(feature = "h2c")]
    async fn run_h2c(&self, app: Router, mut shutdown: watch::Receiver<bool>) -> Result<()> {
        use hyper_util::rt::{TokioExecutor, TokioIo};
        use hyper_util::server::conn::auto::Builder as AutoBuilder;

        info!("h2c (HTTP/2 cleartext) enabled");

        let listener = TcpListener::bind(self.config.listen).await?;
        let builder = Arc::new(AutoBuilder::new(TokioExecutor::new()));
        let mut tasks = tokio::task::JoinSet::new();

        loop {
            tokio::select! {
                result = listener.accept() => {
                    let (stream, remote_addr) = result?;
                    let app = app.clone();
                    let builder = builder.clone();
                    tasks.spawn(async move {
                        let svc = hyper::service::service_fn(move |mut req: Request<hyper::body::Incoming>| {
                            // Inject ConnectInfo manually since we bypass axum::serve.
                            req.extensions_mut().insert(ConnectInfo(remote_addr));
                            let app = app.clone();
                            async move {
                                let resp = tower::Service::call(&mut app.clone(), req).await;
                                resp.map_err(|e| match e {})
                            }
                        });
                        let _ = builder.serve_connection_with_upgrades(TokioIo::new(stream), svc).await;
                    });
                }
                _ = async {
                    loop {
                        if shutdown.changed().await.is_err() || *shutdown.borrow() {
                            break;
                        }
                    }
                } => {
                    break;
                }
            }
        }

        // Wait for active connections to finish
        while tasks.join_next().await.is_some() {}

        Ok(())
    }
}

async fn shutdown_signal(mut shutdown: watch::Receiver<bool>) {
    loop {
        if shutdown.changed().await.is_err() || *shutdown.borrow() {
            break;
        }
    }
}

struct ConnectionGuard(Arc<AtomicU64>);

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        self.0.fetch_sub(1, Ordering::Relaxed);
    }
}

async fn handle(
    State(state): State<AppState>,
    connect_info: ConnectInfo<SocketAddr>,
    req: Request<Body>,
) -> Response<Body> {
    state.active_connections.fetch_add(1, Ordering::Relaxed);
    let _conn_guard = ConnectionGuard(state.active_connections.clone());
    let start = Instant::now();
    let method = req.method().clone();
    let uri = req.uri().clone();
    let peer_addr = connect_info.0;

    let client_ip = resolve_client_ip(
        peer_addr.ip(),
        req.headers()
            .get("x-forwarded-for")
            .and_then(|v| v.to_str().ok()),
        &state.config.trusted_proxies,
    );

    let (response, request_id) = handle_inner(&state, req, client_ip).await;

    if state.config.access_log {
        let duration = start.elapsed();
        let status = response.status().as_u16();
        let response_bytes = response
            .headers()
            .get(http::header::CONTENT_LENGTH)
            .and_then(|v| v.to_str().ok())
            .and_then(|v| v.parse::<u64>().ok())
            .unwrap_or(0);
        info!(
            request_id = %request_id,
            client_ip = %client_ip,
            method = %method,
            uri = %uri,
            status = status,
            duration_ms = duration.as_millis() as u64,
            response_bytes = response_bytes,
            "http request",
        );
    }

    // _conn_guard drop handles fetch_sub
    response
}

/// Handles the request and returns the response together with the `request_id`
/// under which it was dispatched (empty before dispatch, e.g. on read errors).
async fn handle_inner(
    state: &AppState,
    req: Request<Body>,
    client_ip: IpAddr,
) -> (Response<Body>, Arc<str>) {
    let no_id: Arc<str> = Arc::from("");
    let max_body = state.config.max_request_size;
    let read_timeout = state.config.read_timeout;

    // ── Extract request parts for hooks before consuming the body ────────────
    let (parts, body) = req.into_parts();

    let req_method = parts.method.to_string();
    let req_path = parts.uri.path().to_string();
    let req_query = parts.uri.query().unwrap_or("").to_string();
    let req_headers: HashMap<String, String> = parts
        .headers
        .iter()
        .filter_map(|(k, v)| Some((k.to_string(), v.to_str().ok()?.to_string())))
        .collect();

    // Reassemble so encode_request can consume the body.
    let req_reassembled = Request::from_parts(parts, body);

    // ── Read + encode body ───────────────────────────────────────────────────
    let payload =
        match tokio::time::timeout(read_timeout, encode_request(req_reassembled, max_body)).await {
            Ok(Ok(p)) => p,
            Ok(Err(e)) => {
                error!(error = ?e, "encode request");
                return (
                    Response::builder()
                        .status(500)
                        .body(Body::from("encode error"))
                        .unwrap(),
                    no_id,
                );
            }
            Err(_) => {
                return (
                    Response::builder()
                        .status(408)
                        .body(Body::from("request body read timeout"))
                        .unwrap(),
                    no_id,
                );
            }
        };

    // ── request.before hooks ──────────────────────────────────────────────────
    let mut req_ctx = RequestContext {
        method: req_method,
        path: req_path,
        query: req_query,
        client_ip: client_ip.to_string(),
        request_id: String::new(), // filled in after dispatch
        headers: req_headers,
        extra: HashMap::new(),
        error: None,
        short_circuited: false,
    };

    if let Some(ref engine) = state.hook_engine {
        match engine.run_request_before(&mut req_ctx) {
            HookResult::ShortCircuit(resp) => {
                // Async hooks already spawned inside run_request_before.
                return (resp, no_id);
            }
            HookResult::Continue => {}
        }
    }

    // ── Dispatch to PHP worker ───────────────────────────────────────────────
    let (response_value, request_id) = match state
        .executor
        .execute_value_traced("http.handle", payload)
        .await
    {
        Ok(v) => v,
        Err(e) => {
            error!(error = ?e, "dispatch to worker");

            // Fire request.error hooks.
            if let Some(ref engine) = state.hook_engine {
                let mut err_ctx = req_ctx.clone();
                err_ctx.request_id = no_id.to_string();
                err_ctx.error = Some(e.to_string());
                engine.run_request_error(&err_ctx);
            }

            return (
                Response::builder()
                    .status(502)
                    .body(Body::from("worker error"))
                    .unwrap(),
                no_id,
            );
        }
    };

    // ── Decode PHP response into parts ───────────────────────────────────────
    let status = response_value
        .get("status")
        .and_then(|v| v.as_u64())
        .unwrap_or(200) as u16;

    let resp_headers: HashMap<String, String> = response_value
        .get("headers")
        .and_then(|v| v.as_object())
        .map(|obj| {
            obj.iter()
                .filter_map(|(k, v)| Some((k.clone(), v.as_str()?.to_string())))
                .collect()
        })
        .unwrap_or_default();

    let body_str = response_value
        .get("body")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .to_string();
    let body_encoding = response_value
        .get("body_encoding")
        .and_then(|v| v.as_str())
        .map(str::to_string);

    // Update request_id in req_ctx for logging purposes.
    req_ctx.request_id = request_id.to_string();

    // ── response.headers hooks ────────────────────────────────────────────────
    let mut resp_ctx = ResponseContext {
        status,
        resp_headers,
        body: None, // not yet available at this stage
        short_circuited: false,
    };

    if let Some(ref engine) = state.hook_engine {
        match engine.run_response_headers(&mut resp_ctx) {
            HookResult::ShortCircuit(resp) => return (resp, request_id),
            HookResult::Continue => {}
        }
    }

    // ── Decode body bytes ────────────────────────────────────────────────────
    let body_bytes = if body_encoding.as_deref() == Some("base64") {
        use base64::Engine;
        match base64::engine::general_purpose::STANDARD.decode(&body_str) {
            Ok(b) => b,
            Err(e) => {
                error!(error = ?e, "decode base64 response body");
                return (
                    Response::builder()
                        .status(500)
                        .body(Body::from("decode error"))
                        .unwrap(),
                    request_id,
                );
            }
        }
    } else {
        body_str.clone().into_bytes()
    };

    // ── response.after hooks ──────────────────────────────────────────────────
    if let Some(ref engine) = state.hook_engine {
        // Only clone body_bytes when response.after hooks actually need it.
        if engine.has_event("response.after") {
            resp_ctx.body = Some(body_bytes.clone());
        }

        match engine.run_response_after(&mut resp_ctx) {
            HookResult::ShortCircuit(resp) => return (resp, request_id),
            HookResult::Continue => {}
        }
    }

    // ── Build final response ──────────────────────────────────────────────────
    let final_body_bytes: bytes::Bytes = resp_ctx
        .body
        .take()
        .map(bytes::Bytes::from)
        .unwrap_or_else(|| bytes::Bytes::from(body_bytes));

    let mut builder = Response::builder().status(resp_ctx.status);
    for (k, v) in &resp_ctx.resp_headers {
        builder = builder.header(k.as_str(), v.as_str());
    }

    let response = match builder.body(Body::from(final_body_bytes)) {
        Ok(r) => r,
        Err(e) => {
            error!(error = ?e, "build response");
            Response::builder()
                .status(500)
                .body(Body::from("build error"))
                .unwrap()
        }
    };

    (response, request_id)
}

/// Resolve the real client IP from X-Forwarded-For if the peer is a trusted proxy.
///
/// Walks the X-Forwarded-For chain from right to left, stopping at the first
/// IP that is NOT in a trusted subnet. This is the standard secure algorithm
/// (rightmost non-trusted).
pub fn resolve_client_ip(peer_ip: IpAddr, xff: Option<&str>, trusted: &[IpNet]) -> IpAddr {
    if trusted.is_empty() {
        return peer_ip;
    }

    if !is_trusted(peer_ip, trusted) {
        return peer_ip;
    }

    let Some(xff) = xff else {
        return peer_ip;
    };

    let addrs: Vec<&str> = xff.split(',').map(|s| s.trim()).collect();

    // Walk from right to left — the rightmost non-trusted IP is the client.
    for addr_str in addrs.iter().rev() {
        if let Ok(ip) = addr_str.parse::<IpAddr>() {
            if !is_trusted(ip, trusted) {
                return ip;
            }
        }
    }

    // All IPs in the chain are trusted — use peer IP.
    peer_ip
}

fn is_trusted(ip: IpAddr, trusted: &[IpNet]) -> bool {
    trusted.iter().any(|net| net.contains(&ip))
}

fn build_compression_layer(
    config: &crate::config::CompressionConfig,
) -> tower_http::compression::CompressionLayer<tower_http::compression::predicate::SizeAbove> {
    use crate::config::CompressionAlgorithm;
    use tower_http::compression::CompressionLayer;

    let mut layer = CompressionLayer::new()
        .no_gzip()
        .no_br()
        .no_zstd()
        .no_deflate();

    for algo in &config.algorithms {
        layer = match algo {
            CompressionAlgorithm::Gzip => layer.gzip(true),
            CompressionAlgorithm::Br => layer.br(true),
            CompressionAlgorithm::Zstd => layer.zstd(true),
            CompressionAlgorithm::Deflate => layer.deflate(true),
        };
    }

    #[allow(clippy::cast_possible_truncation)]
    let min_size = config.min_size as u16;
    layer.compress_when(tower_http::compression::predicate::SizeAbove::new(min_size))
}