product-os-proxy 0.0.19

use crate::mitm::{
    certificate_authority::CertificateAuthority, HttpContext, HttpHandler, RequestOrResponse,
    Rewind, WebSocketContext, WebSocketHandler,
};
use futures_util::{Sink, SinkExt, Stream, StreamExt};

use hyper_util::{
    client::legacy::{connect::Connect, Client},
    rt::TokioExecutor,
};

use product_os_http::{
    uri::{Authority, Scheme},
    HeaderValue, Method, Request, Response, StatusCode, Uri, Version,
};
use product_os_http_body::{BodyBytes, BodyExt, Full};
use product_os_request::ProductOSClient;

use hyper::body::Incoming;
use hyper::upgrade::Upgraded;
use hyper_util::rt::TokioIo;
use product_os_http_body::Bytes;
use std::{future::Future, net::SocketAddr, sync::Arc};

use hyper::service::service_fn;
use hyper_util::server::conn::auto::Builder;
use product_os_utilities::ProductOSError;
use tokio::{io::AsyncReadExt, net::TcpStream, task::JoinHandle};

use tokio_rustls::TlsAcceptor;
use tokio_tungstenite::{
    tungstenite::{self, Message},
    Connector, WebSocketStream,
};

#[cfg(feature = "tor")]
use std::fmt::{Debug, Formatter};

#[cfg(feature = "tor")]
use tokio::io::{AsyncRead, AsyncWrite};

#[cfg(feature = "tor")]
use rustls_pki_types::{CertificateDer, ServerName, UnixTime};

use tracing::{Instrument, Span};

#[cfg(feature = "tor")]
use crate::mitm::rustls;

#[cfg(feature = "tor")]
use crate::mitm::rustls::client::danger::{
    HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier,
};

#[cfg(feature = "tor")]
use crate::mitm::rustls::{DigitallySignedStruct, SignatureScheme};

fn spawn_with_trace<T: Send + Sync + 'static>(
    fut: impl Future<Output = T> + Send + 'static,
    span: Span,
) -> JoinHandle<T> {
    tokio::spawn(fut.instrument(span))
}

/// Builds an empty response with the given status code (for error/gateway responses).
#[must_use]
fn empty_response(status: StatusCode) -> Response<BodyBytes> {
    Response::builder()
        .status(status)
        .body(BodyBytes::empty())
        .unwrap_or_else(|_| Response::new(BodyBytes::empty()))
}

/// Returns true when the request should be handled as a WebSocket upgrade rather
/// than forwarded as a plain HTTP request or HTTPS tunnel.
fn is_websocket_request<B>(req: &Request<B>) -> bool {
    if hyper_tungstenite::is_upgrade_request(req) {
        return true;
    }

    if matches!(req.uri().scheme_str(), Some("ws") | Some("wss")) {
        return true;
    }

    if is_local_automation_socket(req) {
        return true;
    }

    is_h2_websocket_connect(req)
}

/// True when middleware rewrote the automation socket to the local WSS server.
fn is_local_automation_socket<B>(req: &Request<B>) -> bool {
    let Some(host) = req.uri().host() else {
        return false;
    };
    if host != "localhost" && host != "127.0.0.1" {
        return false;
    }
    req.uri().path().ends_with("/socket_endpoint")
}

fn is_localhost_authority(authority: &Authority) -> bool {
    match authority.host() {
        "localhost" | "127.0.0.1" => true,
        _ => false,
    }
}

#[cfg(feature = "http2")]
fn is_h2_websocket_connect<B>(req: &Request<B>) -> bool {
    req.method() == Method::CONNECT
        && req.version() == Version::HTTP_2
        && req
            .extensions()
            .get::<hyper::ext::Protocol>()
            .is_some_and(|p| p.as_str() == "websocket")
}

#[cfg(not(feature = "http2"))]
fn is_h2_websocket_connect<B>(_req: &Request<B>) -> bool {
    false
}

fn normalize_request_for_upstream(request: &mut Request<BodyBytes>) {
    // Browser traffic on the MITM TLS socket may arrive as HTTP/2 streams, but the
    // legacy hyper client forwards upstream over HTTP/1.1 (or its own TLS stack).
    if request.version() == Version::HTTP_2 {
        *request.version_mut() = Version::HTTP_11;
    }

    // HTTP/2-only headers must not be forwarded to upstream HTTP/1 servers.
    request.headers_mut().remove("priority");
    // Internal MITM tracking headers must not reach origin servers.
    request.headers_mut().remove("WindowID");
    request.headers_mut().remove("RequestID");
}

fn prepare_websocket_request(req: Request<BodyBytes>) -> Result<Request<()>, Response<BodyBytes>> {
    let (mut parts, _) = req.into_parts();

    let uri = {
        let mut uri_parts = parts.uri.into_parts();
        let scheme = if uri_parts.scheme.unwrap_or(Scheme::HTTP) == Scheme::HTTP {
            match "ws".try_into() {
                Ok(s) => s,
                Err(_) => return Err(empty_response(StatusCode::BAD_REQUEST)),
            }
        } else {
            match "wss".try_into() {
                Ok(s) => s,
                Err(_) => return Err(empty_response(StatusCode::BAD_REQUEST)),
            }
        };
        uri_parts.scheme = Some(scheme);
        match Uri::from_parts(uri_parts) {
            Ok(uri) => uri,
            Err(_) => return Err(empty_response(StatusCode::BAD_REQUEST)),
        }
    };

    parts.uri = uri;
    Ok(Request::from_parts(parts, ()))
}

/// RFC 8441 browser-side requests stay as CONNECT; the local WSS server expects
/// a normal RFC 6455 GET upgrade from the proxy's outbound client connection.
fn prepare_upstream_websocket_client_request(req: Request<()>) -> Request<()> {
    let (mut parts, _) = req.into_parts();
    let is_h2_connect = parts.method == Method::CONNECT;

    if is_h2_connect {
        parts.method = Method::GET;
    }

    // Upstream WSS servers speak HTTP/1.1; H2 CONNECT inherits HTTP/2.0.
    parts.version = Version::HTTP_11;

    let headers = &mut parts.headers;
    headers.remove("priority");

    // H2 CONNECT carries :authority in the URI only; tungstenite requires Host.
    // After alter_request rewrites the URI, Host may still name the page origin.
    if let Some(authority) = parts.uri.authority() {
        headers.remove("host");
        if let Ok(v) = HeaderValue::from_str(authority.as_str()) {
            headers.insert("host", v);
        }
    }

    if is_h2_connect {
        headers.remove("sec-websocket-key");
        headers.remove("sec-websocket-version");
        headers.remove("upgrade");
        headers.remove("connection");

        let key = tungstenite::handshake::client::generate_key();
        if let Ok(v) = HeaderValue::from_str(&key) {
            headers.insert("sec-websocket-key", v);
        }
        headers.insert("sec-websocket-version", HeaderValue::from_static("13"));
        headers.insert("upgrade", HeaderValue::from_static("websocket"));
        headers.insert("connection", HeaderValue::from_static("Upgrade"));
    }

    Request::from_parts(parts, ())
}

/// Sets the `accept-encoding` header on the request based on proxy compression config.
fn set_accept_encoding(
    request: &mut Request<BodyBytes>,
    compression: crate::config::NetworkProxyCompression,
) {
    match compression {
        // Passthrough preserves upstream br/zstd/etc. The middleware only decompresses when
        // it needs to rewrite a body; forcing gzip here breaks dev servers like Vite.
        crate::config::NetworkProxyCompression::None => {}
        crate::config::NetworkProxyCompression::Gzip
        | crate::config::NetworkProxyCompression::Brotli => {
            request.headers_mut().remove("accept-encoding");
            let value = match compression {
                crate::config::NetworkProxyCompression::Gzip => "gzip",
                crate::config::NetworkProxyCompression::Brotli => "br",
                crate::config::NetworkProxyCompression::None => unreachable!(),
            };
            if let Ok(v) = HeaderValue::from_str(value) {
                request.headers_mut().insert("accept-encoding", v);
            }
        }
    }
}

pub struct InternalProxy<C, CA, H, W> {
    pub client_addr: SocketAddr,
    pub ca: Arc<CA>,
    pub client: Client<C, BodyBytes>,
    pub websocket_connector: Option<Connector>,

    pub http_handler: H,
    pub websocket_handler: W,

    pub server: Builder<TokioExecutor>,

    pub custom_requester: Option<product_os_request::ProductOSRequestClient>,
    pub compression: crate::config::NetworkProxyCompression,
    pub request_timeout_ms: Option<u64>,

    #[cfg(feature = "tor")]
    pub tor_client: Option<arti_client::TorClient<tor_rtcompat::PreferredRuntime>>,
    #[cfg(feature = "vpn")]
    pub vpn_client: Option<product_os_vpn::ProductOSVPN>,
}

impl<C, CA, H, W> Clone for InternalProxy<C, CA, H, W>
where
    C: Connect + Clone + Send + Sync + 'static,
    CA: CertificateAuthority,
    H: HttpHandler + Clone,
    W: WebSocketHandler + Clone,
{
    fn clone(&self) -> Self {
        InternalProxy {
            client_addr: self.client_addr,
            ca: Arc::clone(&self.ca),
            client: self.client.clone(),
            websocket_connector: self.websocket_connector.clone(),

            http_handler: self.http_handler.clone(),
            websocket_handler: self.websocket_handler.clone(),

            custom_requester: self.custom_requester.clone(),
            compression: self.compression,
            request_timeout_ms: self.request_timeout_ms,

            server: self.server.clone(),

            #[cfg(feature = "tor")]
            tor_client: self.tor_client.clone(),
            #[cfg(feature = "vpn")]
            vpn_client: self.vpn_client.clone(),
        }
    }
}

impl<C, CA, H, W> InternalProxy<C, CA, H, W>
where
    C: Connect + Clone + Send + Sync + 'static,
    CA: CertificateAuthority,
    H: HttpHandler + Clone,
    W: WebSocketHandler + Clone,
{
    fn context(&self) -> HttpContext {
        HttpContext {
            client_addr: self.client_addr,
        }
    }

    /// Main proxy handler that processes incoming HTTP requests
    ///
    /// Steps:
    /// 1. Applies middleware transformations to the request
    /// 2. Determines the request type (WebSocket, CONNECT, or standard HTTP)
    /// 3. Routes the request appropriately (tunnel, upgrade, or forward)
    /// 4. Applies middleware transformations to the response
    #[allow(clippy::too_many_lines)]
    pub async fn proxy(
        mut self,
        req: Request<Incoming>,
    ) -> Result<Response<BodyBytes>, hyper::Error> {
        // Create context with client address for logging and filtering
        let ctx = HttpContext {
            client_addr: self.client_addr,
        };

        let method = req.method().clone();
        let uri = req.uri().clone();
        let span = tracing::trace_span!("handle_request", method = %method, uri = %uri);

        // Apply request middleware - may return early with a response
        let req = match self
            .http_handler
            .handle_request(&ctx, req.map(BodyBytes::from))
            .instrument(span)
            .await
        {
            RequestOrResponse::Request(req) => req,
            // Middleware returned early response, skip normal processing
            RequestOrResponse::Response(res) => return Ok(res),
        };

        // Check if this is a WebSocket upgrade request (HTTP/1 Upgrade or HTTP/2 CONNECT)
        if is_websocket_request(&req) {
            let route = if hyper_tungstenite::is_upgrade_request(&req) {
                "http1_upgrade"
            } else if is_h2_websocket_connect(&req) {
                "h2_connect"
            } else if is_local_automation_socket(&req) {
                "local_automation"
            } else {
                "fallback_upgrade"
            };
            tracing::info!(
                "Routing WebSocket request via {}: method={} version={:?} uri={}",
                route,
                req.method(),
                req.version(),
                req.uri()
            );
            let res = if hyper_tungstenite::is_upgrade_request(&req) {
                self.upgrade_websocket(req)
            } else if is_h2_websocket_connect(&req) {
                self.upgrade_websocket_h2(req)
            } else {
                self.upgrade_websocket(req)
            };
            Ok(res)
        }
        // Check if this is an HTTP CONNECT request (for HTTPS tunneling)
        else if req.method() == Method::CONNECT {
            // Extract the authority (hostname:port) from the request
            match req.uri().authority().cloned() {
                Some(authority) => {
                    // Spawn async task to handle the tunnel connection
                    self.process_connect(req, authority);
                    // Return empty response immediately as connection upgrade happens separately
                    Ok(Response::new(BodyBytes::empty()))
                }
                None => Ok(empty_response(StatusCode::BAD_REQUEST)),
            }
        }
        // Standard HTTP request - forward to destination
        else {
            // Collect the request body into bytes for processing (Bytes supports cheap clone)
            let (parts, body) = req.into_parts();
            let body_bytes: Bytes = match body.map_frame(|b| b).collect().await {
                Ok(collected) => collected.to_bytes(),
                Err(e) => {
                    tracing::error!("Failed to collect request body: {:?}", e);
                    return Ok(empty_response(StatusCode::BAD_GATEWAY));
                }
            };

            // Two requests: one for middleware (original), one for forwarding. Bytes::clone() is cheap (Arc).
            let request_original = Request::from_parts(
                parts.clone(),
                BodyBytes::from(Full::new(body_bytes.clone())),
            );
            let mut request = Request::from_parts(parts, BodyBytes::from(Full::new(body_bytes)));

            normalize_request_for_upstream(&mut request);
            set_accept_encoding(&mut request, self.compression);

            let res;

            // Standard HTTP client (non-Tor) path
            #[cfg(not(feature = "tor"))]
            {
                tracing::debug!("Attempting request: {:?}", request);

                match self.custom_requester {
                    // Use standard hyper client
                    None => {
                        let req_fut = self.client.request(request);
                        res = match self.request_timeout_ms {
                            Some(ms) => match tokio::time::timeout(
                                std::time::Duration::from_millis(ms),
                                req_fut,
                            )
                            .await
                            {
                                Ok(Ok(response)) => {
                                    tracing::debug!("Response being made: {:?}", response);
                                    let (parts, body) = response.into_parts();
                                    let body_bytes: Bytes =
                                        match body.map_frame(|b| b).collect().await {
                                            Ok(collected) => collected.to_bytes(),
                                            Err(e) => {
                                                tracing::error!(
                                                    "Failed to collect response body: {:?}",
                                                    e
                                                );
                                                return Ok(empty_response(StatusCode::BAD_GATEWAY));
                                            }
                                        };
                                    Response::from_parts(
                                        parts,
                                        BodyBytes::from(Full::new(body_bytes)),
                                    )
                                }
                                Ok(Err(e)) => {
                                    tracing::error!("Error attempting request: {:?}", e);
                                    empty_response(StatusCode::BAD_GATEWAY)
                                }
                                Err(_) => {
                                    tracing::warn!(
                                        "Request timed out after {} ms",
                                        self.request_timeout_ms.unwrap_or(0)
                                    );
                                    empty_response(StatusCode::GATEWAY_TIMEOUT)
                                }
                            },
                            None => match req_fut.await {
                                Ok(response) => {
                                    tracing::debug!("Response being made: {:?}", response);
                                    let (parts, body) = response.into_parts();
                                    let body_bytes: Bytes =
                                        match body.map_frame(|b| b).collect().await {
                                            Ok(collected) => collected.to_bytes(),
                                            Err(e) => {
                                                tracing::error!(
                                                    "Failed to collect response body: {:?}",
                                                    e
                                                );
                                                return Ok(empty_response(StatusCode::BAD_GATEWAY));
                                            }
                                        };
                                    Response::from_parts(
                                        parts,
                                        BodyBytes::from(Full::new(body_bytes)),
                                    )
                                }
                                Err(e) => {
                                    tracing::error!("Error attempting request: {:?}", e);
                                    empty_response(StatusCode::BAD_GATEWAY)
                                }
                            },
                        };
                    }
                    // Use custom requester (e.g., for VPN routing)
                    Some(custom_requester) => {
                        tracing::info!("Request being made: {:?}", request);
                        let build_response_fut = async {
                            match custom_requester.request_raw(request).await {
                                Ok(res) => {
                                    let status = res
                                        .try_status()
                                        .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
                                    let headers = res.try_get_headers().unwrap_or_default();
                                    match custom_requester.bytes(res).await {
                                        Ok(bytes) => {
                                            let mut response = Response::builder().status(status);
                                            for (key, value) in headers {
                                                response = response.header(key, value);
                                            }
                                            match response.body(BodyBytes::from(Full::new(bytes))) {
                                                Ok(r) => r,
                                                Err(e) => {
                                                    tracing::error!(
                                                        "Error building response: {:?}",
                                                        e
                                                    );
                                                    empty_response(
                                                        StatusCode::INTERNAL_SERVER_ERROR,
                                                    )
                                                }
                                            }
                                        }
                                        Err(e) => {
                                            tracing::error!(
                                                "Error reading response bytes: {:?}",
                                                e
                                            );
                                            empty_response(StatusCode::INTERNAL_SERVER_ERROR)
                                        }
                                    }
                                }
                                Err(e) => {
                                    tracing::error!("Error making custom request: {:?}", e);
                                    empty_response(StatusCode::INTERNAL_SERVER_ERROR)
                                }
                            }
                        };
                        res = if let Some(ms) = self.request_timeout_ms {
                            if let Ok(response) = tokio::time::timeout(
                                std::time::Duration::from_millis(ms),
                                build_response_fut,
                            )
                            .await
                            {
                                response
                            } else {
                                tracing::warn!("Custom request timed out after {} ms", ms);
                                empty_response(StatusCode::GATEWAY_TIMEOUT)
                            }
                        } else {
                            build_response_fut.await
                        };
                    }
                }
            }
            #[cfg(feature = "tor")]
            {
                res = if let Some(tor_client) = self.tor_client {
                    let uri = request.uri();

                    let host = match uri.host() {
                        None => "unknown",
                        Some(host) => host,
                    };

                    let port: u16 = uri.port_u16().unwrap_or_else(|| match uri.scheme_str() {
                        Some("https" | "HTTPS") => 443,
                        _ => 80,
                    });

                    match tor_client.connect((host, port)).await {
                        Ok(stream) => {
                            if port == 443 {
                                tracing::info!(
                                    "Attempting TLS connect: {:?}, port: {:?}",
                                    host,
                                    port
                                );
                                let tls_connector_builder = rustls::ClientConfig::builder();
                                let tls_connector_config = tls_connector_builder
                                    .dangerous()
                                    .with_custom_certificate_verifier(Arc::new(
                                        CustomCertVerifier::new(),
                                    ))
                                    .with_no_client_auth();
                                let tls_connector = tokio_rustls::TlsConnector::from(Arc::new(
                                    tls_connector_config,
                                ));

                                let server_name = match ServerName::try_from(host.to_owned()) {
                                    Ok(sn) => sn,
                                    Err(e) => {
                                        tracing::error!("Invalid server name '{}': {:?}", host, e);
                                        return Ok(empty_response(StatusCode::BAD_GATEWAY));
                                    }
                                };
                                match tls_connector.connect(server_name, stream).await {
                                    Ok(stream) => {
                                        tracing::info!(
                                            "TLS connect success: {:?}, address: {:?}",
                                            host,
                                            port
                                        );

                                        make_stream_request_h2(request, stream).await
                                    }
                                    Err(e) => {
                                        tracing::error!(
                                            "Error attempting request via Tor: {:?}",
                                            e
                                        );
                                        empty_response(StatusCode::BAD_GATEWAY)
                                    }
                                }
                            } else {
                                make_stream_request_h1(request, stream).await
                            }
                        }
                        Err(e) => {
                            tracing::error!(
                                "Error attempting to create connection via Tor: {:?}",
                                e
                            );
                            empty_response(StatusCode::BAD_GATEWAY)
                        }
                    }
                } else {
                    tracing::debug!("Attempting request via standard channel: {:?}", request);
                    match self.custom_requester {
                        None => {
                            let req_fut = self.client.request(request);
                            match self.request_timeout_ms {
                                Some(ms) => match tokio::time::timeout(
                                    std::time::Duration::from_millis(ms),
                                    req_fut,
                                )
                                .await
                                {
                                    Ok(Ok(res)) => {
                                        let (parts, body) = res.into_parts();
                                        match body.collect().await {
                                            Ok(collected) => {
                                                let bytes = collected.to_bytes();
                                                Response::from_parts(parts, BodyBytes::new(bytes))
                                            }
                                            Err(e) => {
                                                tracing::error!(
                                                    "Failed to collect response body: {:?}",
                                                    e
                                                );
                                                empty_response(StatusCode::BAD_GATEWAY)
                                            }
                                        }
                                    }
                                    Ok(Err(e)) => {
                                        tracing::error!(
                                            "Error attempting request via Vpn: {:?}",
                                            e
                                        );
                                        empty_response(StatusCode::BAD_GATEWAY)
                                    }
                                    Err(_) => {
                                        tracing::warn!(
                                            "Request timed out after {} ms",
                                            self.request_timeout_ms.unwrap_or(0)
                                        );
                                        empty_response(StatusCode::GATEWAY_TIMEOUT)
                                    }
                                },
                                None => match req_fut.await {
                                    Ok(res) => {
                                        let (parts, body) = res.into_parts();
                                        match body.collect().await {
                                            Ok(collected) => {
                                                let bytes = collected.to_bytes();
                                                Response::from_parts(parts, BodyBytes::new(bytes))
                                            }
                                            Err(e) => {
                                                tracing::error!(
                                                    "Failed to collect response body: {:?}",
                                                    e
                                                );
                                                empty_response(StatusCode::BAD_GATEWAY)
                                            }
                                        }
                                    }
                                    Err(e) => {
                                        tracing::error!(
                                            "Error attempting request via Vpn: {:?}",
                                            e
                                        );
                                        empty_response(StatusCode::BAD_GATEWAY)
                                    }
                                },
                            }
                        }
                        Some(custom_requester) => {
                            let build_response_fut = async {
                                match custom_requester.request_raw(request).await {
                                    Ok(res) => {
                                        let status = res
                                            .try_status()
                                            .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
                                        let headers = res.try_get_headers().unwrap_or_default();
                                        match custom_requester.bytes(res).await {
                                            Ok(bytes) => {
                                                let mut response =
                                                    Response::builder().status(status);
                                                for (key, value) in headers {
                                                    response = response.header(key, value);
                                                }
                                                match response
                                                    .body(BodyBytes::from(Full::new(bytes)))
                                                {
                                                    Ok(r) => r,
                                                    Err(e) => {
                                                        tracing::error!(
                                                            "Error building response: {:?}",
                                                            e
                                                        );
                                                        empty_response(
                                                            StatusCode::INTERNAL_SERVER_ERROR,
                                                        )
                                                    }
                                                }
                                            }
                                            Err(e) => {
                                                tracing::error!(
                                                    "Error reading response bytes: {:?}",
                                                    e
                                                );
                                                empty_response(StatusCode::INTERNAL_SERVER_ERROR)
                                            }
                                        }
                                    }
                                    Err(e) => {
                                        tracing::error!("Error making custom request: {:?}", e);
                                        empty_response(StatusCode::INTERNAL_SERVER_ERROR)
                                    }
                                }
                            };
                            if let Some(ms) = self.request_timeout_ms {
                                if let Ok(response) = tokio::time::timeout(
                                    std::time::Duration::from_millis(ms),
                                    build_response_fut,
                                )
                                .await
                                {
                                    response
                                } else {
                                    tracing::warn!("Custom request timed out after {} ms", ms);
                                    empty_response(StatusCode::GATEWAY_TIMEOUT)
                                }
                            } else {
                                build_response_fut.await
                            }
                        }
                    }
                };
            }

            let method = request_original.method().clone();
            let uri = request_original.uri().clone();
            let span = tracing::trace_span!("handle_response", method = %method, uri = %uri);
            Ok(self
                .http_handler
                .handle_response(&ctx, request_original, res)
                .instrument(span)
                .await)
        }
    }

    async fn upgrade_connect(
        &self,
        req: &mut Request<BodyBytes>,
    ) -> Result<(Rewind<TokioIo<Upgraded>>, [u8; 1024], usize, [u8; 4]), ProductOSError> {
        tracing::debug!("Connect request {:?}", req);
        match hyper::upgrade::on(req).await {
            Ok(upgraded) => {
                let mut upgraded = TokioIo::new(upgraded);
                let mut buffer = [0; 1024];

                let bytes_read = match upgraded.read(&mut buffer).await {
                    Ok(bytes_read) => bytes_read,
                    Err(e) => {
                        tracing::error!("Failed to read from upgraded connection: {}", e);
                        return Err(ProductOSError::GenericError(e.to_string()));
                    }
                };

                tracing::debug!("Connect data {:02X?}", &buffer[..bytes_read]);

                let mut first_bytes = [0u8; 4];
                let copy_len = bytes_read.min(4);
                first_bytes[..copy_len].copy_from_slice(&buffer[..copy_len]);

                let upgraded = Rewind::new(
                    upgraded,
                    Bytes::copy_from_slice(buffer[..bytes_read].as_ref()),
                );

                Ok((upgraded, buffer, bytes_read, first_bytes))
            }
            Err(e) => {
                tracing::error!("Upgrade error: {}", e);
                Err(ProductOSError::GenericError(e.to_string()))
            }
        }
    }

    fn process_connect(mut self, mut req: Request<BodyBytes>, authority: Authority) {
        let span = tracing::trace_span!("process_connect", authority = %authority);
        let fut = async move {
            match self.upgrade_connect(&mut req).await {
                Ok((mut upgraded, buffer, data_len, first_bytes)) => {
                    if data_len == 0 {
                        tracing::warn!(
                            "Connection closed immediately after CONNECT upgrade (0 bytes read)"
                        );
                        return;
                    }

                    if self
                        .http_handler
                        .should_intercept(&self.context(), &req)
                        .await
                        && !is_localhost_authority(&authority)
                    {
                        if data_len >= 4 && first_bytes[..4] == *b"GET " {
                            match self
                                .serve_stream(TokioIo::new(upgraded), Scheme::HTTP, authority)
                                .await
                            {
                                Ok(()) => {}
                                Err(e) => {
                                    tracing::error!("WebSocket connect error: {}", e);
                                }
                            }

                            return;
                        }
                        if data_len >= 2 && first_bytes[..2] == *b"\x16\x03" {
                            let server_config = match self
                                .ca
                                .gen_server_config(&authority)
                                .instrument(tracing::info_span!("gen_server_config", authority = %authority))
                                .await
                            {
                                Ok(cfg) => cfg,
                                Err(e) => {
                                    tracing::error!("Failed to generate server config: {}", e);
                                    return;
                                }
                            };

                            let stream = match TlsAcceptor::from(server_config)
                                .accept(upgraded)
                                .await
                            {
                                Ok(stream) => TokioIo::new(stream),
                                Err(e) => {
                                    tracing::error!("Failed to establish TLS connection: {}", e);
                                    return;
                                }
                            };

                            if let Err(e) =
                                self.serve_stream(stream, Scheme::HTTPS, authority).await
                            {
                                if !e.to_string().starts_with("error shutting down connection") {
                                    tracing::error!("HTTPS connect error: {}", e);
                                }
                            }

                            return;
                        }
                        if data_len >= 1
                            && (first_bytes[0] == 0x81
                                || first_bytes[0] == 0x82
                                || first_bytes[0] == 0x88)
                        {
                            // Websocket data
                            tracing::warn!(
                                "Websocket request: {:?} don't know yet what to do - request: {:02X?}",
                                req,
                                &buffer[..data_len]
                            );

                            return;
                        }
                        tracing::warn!(
                            "Unknown protocol, read '{:02X?}' from upgraded connection",
                            &buffer[..data_len]
                        );
                    }

                    let mut server = match TcpStream::connect(authority.as_ref()).await {
                        Ok(server) => server,
                        Err(e) => {
                            tracing::error!("Failed to connect to {}: {}", authority, e);
                            return;
                        }
                    };

                    if let Err(e) = tokio::io::copy_bidirectional(&mut upgraded, &mut server).await
                    {
                        tracing::error!("Failed to tunnel to {}: {}", authority, e);
                    }
                }
                Err(e) => {
                    tracing::error!("Problem upgrading request: {:?}", e);
                }
            }
        };
        spawn_with_trace(fut, span);
    }

    fn upgrade_websocket(self, req: Request<BodyBytes>) -> Response<BodyBytes> {
        let req = match prepare_websocket_request(req) {
            Ok(req) => req,
            Err(res) => return res,
        };

        let mut upgrade_req = req.clone();
        match hyper_tungstenite::upgrade(&mut upgrade_req, None) {
            Ok((res, websocket)) => {
                let span = tracing::info_span!("websocket");
                let fut = async move {
                    match websocket.await {
                        Ok(ws) => {
                            if let Err(e) = self.handle_websocket(ws, req).await {
                                tracing::error!("Failed to handle WebSocket: {}", e);
                            }
                        }
                        Err(e) => {
                            tracing::error!(
                                "Failed to upgrade to WebSocket {:?}: {}",
                                req.uri(),
                                e
                            );
                        }
                    }
                };

                spawn_with_trace(fut, span);
                let (parts, body) = res.into_parts();
                Response::from_parts(parts, BodyBytes::from(body))
            }
            Err(e) => {
                tracing::error!("WebSocket upgrade failed: {:?}", e);
                empty_response(StatusCode::BAD_REQUEST)
            }
        }
    }

    fn upgrade_websocket_h2(self, req: Request<BodyBytes>) -> Response<BodyBytes> {
        // RFC 8441: HTTP/2 WebSocket uses CONNECT + :protocol=websocket. The opening
        // handshake does not use Sec-WebSocket-Key / Sec-WebSocket-Accept; after 200
        // the stream carries WebSocket frames directly (state OPEN).
        if req
            .headers()
            .get("sec-websocket-version")
            .map(|v| v.as_bytes())
            != Some(b"13")
            && req
                .headers()
                .get("Sec-WebSocket-Version")
                .map(|v| v.as_bytes())
                != Some(b"13")
        {
            tracing::error!(
                "HTTP/2 WebSocket CONNECT missing or invalid Sec-WebSocket-Version (uri: {:?})",
                req.uri()
            );
            return empty_response(StatusCode::BAD_REQUEST);
        }

        let mut req = match prepare_websocket_request(req) {
            Ok(req) => req,
            Err(res) => return res,
        };

        let upstream_uri = req.uri().clone();
        tracing::info!(
            "HTTP/2 WebSocket CONNECT accepted (RFC 8441), bridging to upstream {}",
            upstream_uri
        );

        let on_upgrade = hyper::upgrade::on(&mut req);

        let response = match Response::builder()
            .status(StatusCode::OK)
            .body(BodyBytes::empty())
        {
            Ok(res) => res,
            Err(e) => {
                tracing::error!("Failed to build HTTP/2 WebSocket response: {}", e);
                return empty_response(StatusCode::INTERNAL_SERVER_ERROR);
            }
        };

        let span = tracing::info_span!("websocket_h2");
        let fut = async move {
            match on_upgrade.await {
                Ok(upgraded) => {
                    let ws = WebSocketStream::from_raw_socket(
                        TokioIo::new(upgraded),
                        tungstenite::protocol::Role::Server,
                        None,
                    )
                    .await;

                    if let Err(e) = self.handle_websocket(ws, req).await {
                        tracing::error!("Failed to handle HTTP/2 WebSocket: {}", e);
                    }
                }
                Err(e) => {
                    tracing::error!("Failed to upgrade HTTP/2 WebSocket: {}", e);
                }
            }
        };

        spawn_with_trace(fut, span);
        response
    }

    async fn handle_websocket(
        self,
        client_socket: WebSocketStream<TokioIo<Upgraded>>,
        req: Request<()>,
    ) -> Result<(), tungstenite::Error> {
        let uri = req.uri().clone();
        let upstream_req = prepare_upstream_websocket_client_request(req);
        tracing::info!(
            "Connecting upstream automation WebSocket to {} via {}",
            uri,
            upstream_req.method()
        );

        #[cfg(feature = "rustls_client")]
        let (server_socket, _) = tokio_tungstenite::connect_async_tls_with_config(
            upstream_req,
            None,
            false,
            self.websocket_connector,
        )
        .await?;

        #[cfg(not(feature = "rustls_client"))]
        let (server_socket, _) = tokio_tungstenite::connect_async(upstream_req).await?;

        tracing::info!("Upstream automation WebSocket connected to {}", uri);

        let (server_sink, server_stream) = server_socket.split();
        let (client_sink, client_stream) = client_socket.split();

        let InternalProxy {
            websocket_handler, ..
        } = self;

        spawn_message_forwarder(
            server_stream,
            client_sink,
            websocket_handler.clone(),
            WebSocketContext::ServerToClient {
                src: uri.clone(),
                dst: self.client_addr,
            },
        );

        spawn_message_forwarder(
            client_stream,
            server_sink,
            websocket_handler,
            WebSocketContext::ClientToServer {
                src: self.client_addr,
                dst: uri,
            },
        );

        Ok(())
    }

    async fn serve_stream<I>(
        self,
        stream: I,
        _scheme: Scheme,
        _authority: Authority,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>
    where
        I: hyper::rt::Read + hyper::rt::Write + Unpin + Send + 'static,
    {
        let service = service_fn(|req| self.clone().proxy(req));

        self.server
            .serve_connection_with_upgrades(stream, service)
            .await
    }
}

fn spawn_message_forwarder(
    mut stream: impl Stream<Item = Result<Message, tungstenite::Error>> + Unpin + Send + 'static,
    mut sink: impl Sink<Message, Error = tungstenite::Error> + Unpin + Send + 'static,
    mut handler: impl WebSocketHandler,
    ctx: WebSocketContext,
) {
    let span = tracing::trace_span!("message_forwarder", context = ?ctx);
    let fut = async move {
        while let Some(message) = stream.next().await {
            match message {
                Ok(message) => {
                    let Some(message) = handler.handle_message(&ctx, message).await else {
                        continue;
                    };

                    match sink.send(message).await {
                        Err(tungstenite::Error::ConnectionClosed) | Ok(()) => (),
                        Err(e) => tracing::error!("Websocket send error: {:?}", e),
                    }
                }
                Err(e) => {
                    tracing::error!("Websocket message error: {:?}", e);

                    match sink.send(Message::Close(None)).await {
                        Err(tungstenite::Error::ConnectionClosed) | Ok(()) => (),
                        Err(e) => tracing::error!("Websocket close error: {:?}", e),
                    }

                    break;
                }
            }
        }
    };

    spawn_with_trace(fut, span);
}

#[cfg(feature = "tor")]
pub struct CustomCertVerifier;

#[cfg(feature = "tor")]
impl CustomCertVerifier {
    pub fn new() -> Self {
        Self
    }
}

#[cfg(feature = "tor")]
impl Debug for CustomCertVerifier {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CustomCertVerifier").finish()
    }
}

#[cfg(feature = "tor")]
impl ServerCertVerifier for CustomCertVerifier {
    fn verify_server_cert(
        &self,
        _end_entity: &CertificateDer<'_>,
        _intermediates: &[CertificateDer<'_>],
        _server_name: &ServerName<'_>,
        _ocsp_response: &[u8],
        _now: UnixTime,
    ) -> Result<ServerCertVerified, rustls::Error> {
        Ok(ServerCertVerified::assertion())
    }

    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _cert: &CertificateDer<'_>,
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        Ok(HandshakeSignatureValid::assertion())
    }

    fn verify_tls13_signature(
        &self,
        _message: &[u8],
        _cert: &CertificateDer<'_>,
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        Ok(HandshakeSignatureValid::assertion())
    }

    fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
        // Return common signature schemes
        vec![
            SignatureScheme::RSA_PKCS1_SHA256,
            SignatureScheme::RSA_PKCS1_SHA384,
            SignatureScheme::RSA_PKCS1_SHA512,
            SignatureScheme::ECDSA_NISTP256_SHA256,
            SignatureScheme::ECDSA_NISTP384_SHA384,
            SignatureScheme::ECDSA_NISTP521_SHA512,
            SignatureScheme::RSA_PSS_SHA256,
            SignatureScheme::RSA_PSS_SHA384,
            SignatureScheme::RSA_PSS_SHA512,
            SignatureScheme::ED25519,
        ]
    }
}

#[cfg(feature = "tor")]
async fn make_stream_request_h2(
    request: Request<BodyBytes>,
    stream: impl AsyncRead + AsyncWrite + Unpin + Send + 'static,
) -> Response<BodyBytes> {
    let handshake_result =
        hyper::client::conn::http2::handshake(TokioExecutor::new(), TokioIo::new(stream)).await;
    let (mut request_sender, connection) = match handshake_result {
        Ok(val) => val,
        Err(e) => {
            tracing::error!("Error generating tor h2 request: {e:?}");
            return empty_response(StatusCode::BAD_GATEWAY);
        }
    };

    tokio::spawn(async move {
        if let Err(e) = connection.await {
            tracing::error!("Error while polling tor h2 request: {:?}", e);
        }
    });

    match request_sender.send_request(request).await {
        Ok(response) => collect_tor_response(response).await,
        Err(e) => {
            tracing::error!("Error making tor h2 request: {:?}", e);
            empty_response(StatusCode::BAD_GATEWAY)
        }
    }
}

#[cfg(feature = "tor")]
async fn make_stream_request_h1(
    request: Request<BodyBytes>,
    stream: impl AsyncRead + AsyncWrite + Unpin + Send + 'static,
) -> Response<BodyBytes> {
    let handshake_result = hyper::client::conn::http1::handshake(TokioIo::new(stream)).await;
    let (mut request_sender, connection) = match handshake_result {
        Ok(val) => val,
        Err(e) => {
            tracing::error!("Error generating tor h1 request: {e:?}");
            return empty_response(StatusCode::BAD_GATEWAY);
        }
    };

    tokio::spawn(async move {
        if let Err(e) = connection.await {
            tracing::error!("Error while polling tor h1 request: {:?}", e);
        }
    });

    match request_sender.send_request(request).await {
        Ok(response) => collect_tor_response(response).await,
        Err(e) => {
            tracing::error!("Error making tor h1 request: {:?}", e);
            empty_response(StatusCode::BAD_GATEWAY)
        }
    }
}

#[cfg(feature = "tor")]
async fn collect_tor_response(response: Response<hyper::body::Incoming>) -> Response<BodyBytes> {
    let (parts, mut response_body) = response.into_parts();

    let mut full_bytes = vec![];

    while let Some(frame) = response_body.frame().await {
        let bytes = match frame {
            Ok(frame) => frame.into_data().unwrap_or_else(|e| {
                tracing::error!("Problem converting frame to data from tor: {:?}", e);
                Bytes::new()
            }),
            Err(e) => {
                tracing::error!("Problem getting data frame from tor: {:?}", e);
                Bytes::new()
            }
        };
        full_bytes.append(&mut bytes.to_vec());
    }

    Response::from_parts(parts, BodyBytes::new(Bytes::from(full_bytes)))
}

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

    #[test]
    fn none_compression_preserves_client_accept_encoding() {
        let mut request = Request::builder()
            .uri("http://127.0.0.1:5173/")
            .header("accept-encoding", "br, zstd")
            .body(BodyBytes::empty())
            .unwrap();

        set_accept_encoding(&mut request, crate::config::NetworkProxyCompression::None);

        assert_eq!(
            request
                .headers()
                .get("accept-encoding")
                .and_then(|v| v.to_str().ok()),
            Some("br, zstd")
        );
    }

    #[test]
    fn gzip_compression_replaces_client_accept_encoding() {
        let mut request = Request::builder()
            .uri("http://example.com/")
            .header("accept-encoding", "br, zstd")
            .body(BodyBytes::empty())
            .unwrap();

        set_accept_encoding(&mut request, crate::config::NetworkProxyCompression::Gzip);

        assert_eq!(
            request
                .headers()
                .get("accept-encoding")
                .and_then(|v| v.to_str().ok()),
            Some("gzip")
        );
    }
}