include!("request_id/mod.rs");

fn wrap_observed_request(body: Body, exchange: &Arc<Exchange>, idle: Duration) -> Body {
    if body.is_end_stream() {
        body
    } else {
        Body::new(MeteredBody::new(
            body,
            exchange.clone(),
            Direction::Request,
            idle,
        ))
    }
}

fn log_prefix(input: &str, max_chars: usize) -> String {
    if input.len() <= max_chars || input.is_ascii() {
        input[..input.len().min(max_chars)].to_owned()
    } else {
        input.chars().take(max_chars).collect()
    }
}

fn should_replace_error_page(response: &Response, minimal_welcome: bool) -> bool {
    !minimal_welcome
        && response
            .extensions()
            .get::<ErrorFallbackResponse>()
            .is_none()
        && response
            .extensions()
            .get::<crate::pages::OriginProvidedResponse>()
            .is_none()
        && response
            .extensions()
            .get::<crate::pages::PageServed>()
            .is_none()
        && response
            .extensions()
            .get::<crate::auth_provider::AuthenticationResponse>()
            .is_none()
}

async fn observe(State(state): State<Arc<App>>, mut request: Request, next: Next) -> Response {
    let snapshot = state.config.load_full();
    let selected = response_context(&snapshot, &request);
    let frontend_context = selected.frontend.clone();
    let route_context = selected.route.clone();
    request.extensions_mut().insert(selected);
    let local = request.extensions().get::<LocalAddr>().map(|v| v.0);
    let transport_tls = request
        .extensions()
        .get::<Transport>()
        .copied()
        .unwrap_or_default()
        .tls;
    request.extensions_mut().insert(snapshot.clone());
    let original_headers =
        if snapshot.envelope.config.compression.enabled() && request.method() != Method::HEAD {
            crate::compression::request_headers(request.headers())
        } else {
            HeaderMap::new()
        };
    let method = request.method().clone();
    let is_head = method == Method::HEAD;
    let ws_attempt = request
        .headers()
        .get("upgrade")
        .is_some_and(|v| v.as_bytes().eq_ignore_ascii_case(b"websocket"))
        || (request.version() == Version::HTTP_2
            && request.method() == Method::CONNECT
            && websocket::h2_protocol(&request) == Some("websocket"))
        || (request.version() == Version::HTTP_3
            && request.method() == Method::CONNECT
            && websocket::h3_protocol(&request) == Some("websocket"));
    match request.version() {
        Version::HTTP_11 => {
            state.metrics.http1_requests.fetch_add(1, Ordering::Relaxed);
        }
        Version::HTTP_2 => {
            state.metrics.http2_requests.fetch_add(1, Ordering::Relaxed);
        }
        Version::HTTP_3 => {
            state.metrics.http3_requests.fetch_add(1, Ordering::Relaxed);
        }
        _ => {}
    }
    let id = state.metrics.requests.fetch_add(1, Ordering::Relaxed) + 1;
    let request_id = request_id_hash(id);
    if let Some(frontend) = &frontend_context {
        frontend.metrics.requests.fetch_add(1, Ordering::Relaxed);
    }
    let Ok(permit) = state.inflight.clone().try_acquire_owned() else {
        state
            .metrics
            .rejected_requests
            .fetch_add(1, Ordering::Relaxed);
        if let Some(frontend) = &frontend_context {
            frontend.metrics.rejected.fetch_add(1, Ordering::Relaxed);
        }
        let mut response = error(
            StatusCode::SERVICE_UNAVAILABLE,
            "RustProxy: in-flight limit reached\n",
        );
        if let Some(page) = frontend_context
            .as_ref()
            .and_then(|f| f.pages.errors.get(&503))
            .or_else(|| snapshot.envelope.config.pages.errors.get(&503))
        {
            crate::pages::replace_error(&mut response, page);
        }
        if is_head {
            *response.body_mut() = Body::empty();
        }
        apply_response_rules(
            response.headers_mut(),
            &snapshot,
            frontend_context.as_deref(),
            route_context.as_deref(),
            false,
        );
        finish_headers(
            &mut response,
            &snapshot,
            frontend_context.as_deref(),
            local,
            transport_tls,
            &request_id,
            &state.active_listeners,
        );
        apply_response_age_rules(
            response.headers_mut(),
            &snapshot,
            frontend_context.as_deref(),
            route_context.as_deref(),
        );
        return response;
    };
    state.metrics.active.fetch_add(1, Ordering::Relaxed);
    if let Some(frontend) = &frontend_context {
        frontend.metrics.active.fetch_add(1, Ordering::Relaxed);
    }
    let host = log_prefix(
        request
            .headers()
            .get(HOST)
            .and_then(|h| h.to_str().ok())
            .unwrap_or_else(|| request.uri().authority().map(|v| v.as_str()).unwrap_or("")),
        253,
    );
    let initial_client_ip = request
        .extensions()
        .get::<PeerAddr>()
        .map(|peer| peer.0.ip().to_string());
    let user_agent = state
        .logs
        .capture_user_agent()
        .then(|| {
            request
                .headers()
                .get(USER_AGENT)
                .and_then(|value| value.to_str().ok())
                .map(|value| log_prefix(value, 512))
        })
        .flatten();
    let (http_version, transport) = match request.version() {
        Version::HTTP_10 => ("1.0", "h1"),
        Version::HTTP_11 => ("1.1", "h1"),
        Version::HTTP_2 => ("2", "h2"),
        Version::HTTP_3 => ("3", "h3"),
        _ => ("unknown", "unknown"),
    };
    let websocket_bootstrap = ws_attempt.then_some(match request.version() {
        Version::HTTP_2 => "h2_extended_connect",
        Version::HTTP_3 => "h3_extended_connect",
        _ => "h1_upgrade",
    });
    let exchange = Arc::new(Exchange {
        vhost_metrics: frontend_context.as_ref().map(|f| f.metrics.clone()),
        vhost: frontend_context.as_ref().map(|f| f.id.clone()),
        route: Mutex::new(None),
        source: AtomicU8::new(ExchangeSource::Unresolved as u8),
        cache_effective_enabled: AtomicU8::new(0),
        transform_effective_enabled: AtomicU8::new(0),
        backend: Mutex::new(None),
        id,
        request_id: request_id.clone(),
        method: request.method().to_string(),
        host,
        path: log_prefix(request.uri().path(), 256),
        http_version,
        transport,
        websocket_bootstrap,
        websocket_phase: Mutex::new(websocket_bootstrap.map(|_| "request_received")),
        websocket_error: Mutex::new(None),
        h3_reset_code: Mutex::new(None),
        upstream_status: AtomicU16::new(0),
        upstream_error: Mutex::new(None),
        client_ip: Mutex::new(initial_client_ip),
        user_agent,
        country: Mutex::new(None),
        started: Instant::now(),
        status: AtomicU16::new(500),
        cache_status: AtomicU16::new(0),
        cache_reason: Mutex::new("not_proxy_response"),
        failed: AtomicBool::new(false),
        aborted: AtomicBool::new(false),
        request_bytes: AtomicU64::new(0),
        response_bytes: AtomicU64::new(0),
        metrics: state.metrics.clone(),
        logs: state.logs.clone(),
        traffic: state.traffic.clone(),
        _permit: permit,
    });
    if let Ok(value) = HeaderValue::from_str(&request_id) {
        request.headers_mut().insert("x-request-id", value);
    }
    request.extensions_mut().insert(exchange.clone());
    let request_send_timeout = frontend_context
        .as_ref()
        .and_then(|frontend| frontend.proxy_timeouts.send_timeout())
        .unwrap_or(state.body_idle_timeout);
    let request = request.map(|body| wrap_observed_request(body, &exchange, request_send_timeout));
    let mut response = next.run(request).await;
    let final_route_context = response
        .extensions()
        .get::<FinalRouteContext>()
        .map(|context| context.0.clone());
    let effective_route_context = final_route_context
        .as_deref()
        .or_else(|| route_context.as_deref());
    let minimal_welcome = response
        .extensions()
        .get::<crate::pages::MinimalWelcome>()
        .is_some();
    let code = response.status().as_u16();
    let custom_error = frontend_context
        .as_ref()
        .and_then(|f| f.pages.errors.get(&code))
        .or_else(|| snapshot.envelope.config.pages.errors.get(&code));
    if should_replace_error_page(&response, minimal_welcome) {
        if let Some(page) = custom_error {
            crate::pages::replace_error(&mut response, page);
        }
    }
    if minimal_welcome {
        for name in [
            "server",
            "server-timing",
            "alt-svc",
            "x-rustproxy-request-id",
            "x-rustproxy-revision",
            "x-rust-proxy-backend",
            "x-rustproxy-backend",
            "x-rust-proxy-backend-time",
            "x-rustproxy-backend-time",
            "x-cache",
            "x-cache-hits",
            "x-cache-status",
            "x-cache-resident-age",
            "x-cache-reason",
            "x-upstream-cache",
            "x-upstream-cache-hits",
            "x-upstream-cache-reason",
            "x-upstream-cache-resident-age",
            "cache-status",
            "content-encoding",
            "vary",
            "strict-transport-security",
            "content-security-policy",
            "permissions-policy",
            "referrer-policy",
            "x-frame-options",
            "x-content-type-options",
        ] {
            response.headers_mut().remove(name);
        }
        response
            .headers_mut()
            .insert("cache-control", HeaderValue::from_static("no-store"));
    } else {
        if response
            .extensions()
            .get::<ResponseRulesApplied>()
            .is_none()
            && response
                .extensions()
                .get::<crate::auth_provider::AuthenticationResponse>()
                .is_none()
        {
            apply_response_rules(
                response.headers_mut(),
                &snapshot,
                frontend_context.as_deref(),
                effective_route_context,
                false,
            );
        }
        apply_timing_headers(&mut response, frontend_context.as_deref(), &exchange);
        apply_backend_diagnostic_headers(&mut response, frontend_context.as_deref());
        finish_headers(
            &mut response,
            &snapshot,
            frontend_context.as_deref(),
            local,
            transport_tls,
            &request_id,
            &state.active_listeners,
        );
        if response
            .extensions()
            .get::<RedirectResponse>()
            .is_some_and(|redirect| !redirect.http3_enabled)
        {
            crate::http3::withdraw_advertisement(
                response.headers_mut(),
                local,
                transport_tls,
                &state.active_listeners,
            );
        }
        if response
            .extensions()
            .get::<LocalHostResponse>()
            .is_some_and(|item| !item.http3_enabled)
        {
            crate::http3::withdraw_advertisement(
                response.headers_mut(),
                local,
                transport_tls,
                &state.active_listeners,
            );
        }
        let global_compression = &snapshot.envelope.config.compression;
        if method != Method::HEAD
            && crate::compression::preferred_encoding(&original_headers, global_compression)
                .is_some()
        {
            let mut compression = global_compression.clone();
            if let Some(frontend) = frontend_context.as_ref() {
                compression = frontend
                    .compression
                    .resolve(global_compression, &compression);
            }
            if let Some(route) = effective_route_context {
                if let crate::config::CompiledSource::WebServer(service) = &route.source {
                    compression = service
                        .compression
                        .resolve(global_compression, &compression);
                }
            }
            response =
                crate::compression::encode(response, &original_headers, &method, &compression)
                    .await;
        }
        if let Some(encoding) = response
            .extensions()
            .get::<crate::compression::AppliedEncoding>()
            .copied()
        {
            state.metrics.record_compression(encoding);
            if let Some(route) = effective_route_context {
                if let crate::config::CompiledSource::WebServer(service) = &route.source {
                    service.web_metrics.record_encoding(encoding);
                }
            }
        }
        if response
            .extensions()
            .get::<ResponseRulesApplied>()
            .is_none()
            && response.extensions().get::<RedirectResponse>().is_none()
            && response.extensions().get::<LocalHostResponse>().is_none()
            && response
                .extensions()
                .get::<crate::pages::PageServed>()
                .is_none()
            && response
                .extensions()
                .get::<LocalGeneratedResponse>()
                .is_none()
            && response
                .extensions()
                .get::<StaticSourceResponse>()
                .is_none()
            && response
                .extensions()
                .get::<crate::auth_provider::AuthenticationResponse>()
                .is_none()
        {
            let outcome = match exchange.cache_status.load(Ordering::Relaxed) {
                1 => "HIT",
                2 => "MISS",
                _ => "BYPASS",
            };
            let cache = frontend_context
                .as_ref()
                .map(|f| &f.cache)
                .unwrap_or(&snapshot.envelope.config.cache);
            let reason = *exchange
                .cache_reason
                .lock()
                .unwrap_or_else(|e| e.into_inner());
            let transformed = response
                .extensions()
                .get::<crate::transform::Applied>()
                .is_some()
                || response
                    .extensions()
                    .get::<crate::transform::CachedTransformed>()
                    .is_some();
            cache.annotate_reason_transformed(
                response.headers_mut(),
                outcome,
                0,
                0,
                reason,
                transformed,
            );
        }
    }
    if !minimal_welcome {
        apply_response_age_rules(
            response.headers_mut(),
            &snapshot,
            frontend_context.as_deref(),
            effective_route_context,
        );
    }
    if response
        .extensions()
        .get::<crate::pages::PageServed>()
        .is_some()
        || response
            .extensions()
            .get::<LocalGeneratedResponse>()
            .is_some()
    {
        enforce_local_page_no_cache(&mut response);
    }
    if response.extensions().get::<LocalHostResponse>().is_some() {
        for name in [
            "cache-control",
            "age",
            "expires",
            "pragma",
            "etag",
            "last-modified",
            "cdn-cache-control",
            "surrogate-control",
            "x-cache",
            "x-cache-hits",
            "x-cache-status",
            "x-cache-resident-age",
            "x-cache-reason",
            "x-upstream-cache",
            "x-upstream-cache-hits",
            "x-upstream-cache-reason",
            "x-upstream-cache-resident-age",
            "cache-status",
        ] {
            response.headers_mut().remove(name);
        }
    }
    exchange
        .status
        .store(response.status().as_u16(), Ordering::Relaxed);
    let ws_success = response.status() == StatusCode::SWITCHING_PROTOCOLS
        || response.extensions().get::<websocket::H3Tunnel>().is_some();
    if ws_attempt && !ws_success {
        state
            .metrics
            .websocket_handshake_errors
            .fetch_add(1, Ordering::Relaxed);
    }
    if is_head {
        *response.body_mut() = Body::empty();
    }
    let response_read_timeout = frontend_context
        .as_ref()
        .and_then(|frontend| frontend.proxy_timeouts.read_timeout())
        .unwrap_or(state.body_idle_timeout);
    response.map(|body| {
        Body::new(MeteredBody::new(
            body,
            exchange,
            Direction::Response,
            response_read_timeout,
        ))
    })
}
fn effective_host(request: &Request, tls: bool) -> Result<HeaderValue, ()> {
    let count = request.headers().get_all(HOST).iter().count();
    if matches!(request.version(), Version::HTTP_2 | Version::HTTP_3) {
        let authority = request.uri().authority().ok_or(())?;
        if request.uri().scheme_str() != Some(if tls { "https" } else { "http" }) || count > 1 {
            return Err(());
        }
        if let Some(host) = request.headers().get(HOST) {
            if !host
                .to_str()
                .map_err(|_| ())?
                .eq_ignore_ascii_case(authority.as_str())
            {
                return Err(());
            }
        }
        HeaderValue::from_str(authority.as_str()).map_err(|_| ())
    } else {
        if request.uri().scheme().is_some() || request.uri().authority().is_some() || count != 1 {
            return Err(());
        }
        request.headers().get(HOST).cloned().ok_or(())
    }
}

#[cfg(test)]
mod auth_provider_page_interception_tests {
    use super::*;
    #[test]
    fn provider_failures_and_redirects_bypass_custom_http_error_pages() {
        for kind in [
            crate::auth_provider::GatewayFailure::Connection,
            crate::auth_provider::GatewayFailure::Timeout,
            crate::auth_provider::GatewayFailure::Verification,
        ] {
            let response = crate::auth_provider::gateway_failure("Future", kind);
            assert!(!should_replace_error_page(&response, false));
        }
        let response = crate::auth_provider::deny(StatusCode::UNAUTHORIZED);
        assert!(!should_replace_error_page(&response, false));
    }
}