rama-http-backend 0.3.0-alpha.4

use rama_core::{
    Layer, Service,
    error::{ErrorContext, OpaqueError},
    extensions::{ExtensionsMut, ExtensionsRef},
    telemetry::tracing,
};
use rama_http::Request;
use rama_net::address::ProxyAddress;

#[derive(Debug, Clone, Default)]
/// A [`Layer`] which allows you to add a [`ProxyAddress`]
/// to the [`Context`] in order to have your client connector
/// make a connection via this proxy (e.g. by using [`HttpProxyConnectorLayer`]).
///
/// See [`HttpProxyAddressService`] for more information.
///
/// [`Context`]: rama_core::Context
/// [`HttpProxyConnectorLayer`]: crate::client::proxy::layer::HttpProxyConnectorLayer
pub struct HttpProxyAddressLayer {
    address: Option<ProxyAddress>,
    preserve: bool,
}

impl HttpProxyAddressLayer {
    /// Create a new [`HttpProxyAddressLayer`] that will create
    /// a service to set the given [`ProxyAddress`].
    #[must_use]
    pub fn new(address: ProxyAddress) -> Self {
        Self::maybe(Some(address))
    }

    /// Create a new [`HttpProxyAddressLayer`] which will create
    /// a service that will set the given [`ProxyAddress`] if it is not
    /// `None`.
    #[must_use]
    pub fn maybe(address: Option<ProxyAddress>) -> Self {
        Self {
            address,
            ..Default::default()
        }
    }

    /// Try to create a new [`HttpProxyAddressLayer`] which will establish
    /// a proxy connection over the environment variable `HTTP_PROXY`.
    pub fn try_from_env_default() -> Result<Self, OpaqueError> {
        Self::try_from_env("HTTP_PROXY")
    }

    /// Try to create a new [`HttpProxyAddressLayer`] which will establish
    /// a proxy connection over the given environment variable.
    pub fn try_from_env(key: impl AsRef<str>) -> Result<Self, OpaqueError> {
        let env_result = std::env::var(key.as_ref()).ok();
        let env_result_mapped = env_result.as_ref().and_then(|v| {
            let v = v.trim();
            if v.is_empty() { None } else { Some(v) }
        });

        let proxy_address = match env_result_mapped {
            Some(value) => Some(value.try_into().context("parse std env proxy info")?),
            None => None,
        };

        Ok(Self::maybe(proxy_address))
    }

    rama_utils::macros::generate_set_and_with! {
        /// Preserve the existing [`ProxyAddress`] in the context if it already exists.
        pub fn preserve(mut self, preserve: bool) -> Self {
            self.preserve = preserve;
            self
        }
    }
}

impl<S> Layer<S> for HttpProxyAddressLayer {
    type Service = HttpProxyAddressService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        HttpProxyAddressService::maybe(inner, self.address.clone()).with_preserve(self.preserve)
    }

    fn into_layer(self, inner: S) -> Self::Service {
        HttpProxyAddressService::maybe(inner, self.address).with_preserve(self.preserve)
    }
}

/// A [`Service`] which allows you to add a [`ProxyAddress`]
/// to the [`Context`] in order to have your client connector
/// make a connection via this proxy (e.g. by using [`HttpProxyConnectorLayer`]).
///
/// [`Context`]: rama_core::Context
/// [`HttpProxyConnectorLayer`]: crate::client::proxy::layer::HttpProxyConnectorLayer
#[derive(Debug, Clone)]
pub struct HttpProxyAddressService<S> {
    inner: S,
    proxy_info: Option<ProxyAddress>,
    preserve: bool,
}

impl<S> HttpProxyAddressService<S> {
    /// Create a new [`HttpProxyAddressService`] that will create
    /// a service to set the given [`ProxyAddress`].
    pub const fn new(inner: S, address: ProxyAddress) -> Self {
        Self::maybe(inner, Some(address))
    }

    /// Create a new [`HttpProxyAddressService`] which will create
    /// a service that will set the given [`ProxyAddress`] if it is not
    /// `None`.
    pub const fn maybe(inner: S, address: Option<ProxyAddress>) -> Self {
        Self {
            inner,
            proxy_info: address,
            preserve: false,
        }
    }

    /// Try to create a new [`HttpProxyAddressService`] which will establish
    /// a proxy connection over the environment variable `HTTP_PROXY`.
    pub fn try_from_env_default(inner: S) -> Result<Self, OpaqueError> {
        Self::try_from_env(inner, "HTTP_PROXY")
    }

    /// Try to create a new [`HttpProxyAddressService`] which will establish
    /// a proxy connection over the given environment variable.
    pub fn try_from_env(inner: S, key: impl AsRef<str>) -> Result<Self, OpaqueError> {
        let env_result = std::env::var(key.as_ref()).ok();
        let env_result_mapped = env_result.as_ref().and_then(|v| {
            let v = v.trim();
            if v.is_empty() { None } else { Some(v) }
        });

        let proxy_address = match env_result_mapped {
            Some(value) => Some(value.try_into().context("parse std env proxy info")?),
            None => None,
        };

        Ok(Self::maybe(inner, proxy_address))
    }

    rama_utils::macros::generate_set_and_with! {
        /// Preserve the existing [`ProxyAddress`] in the context if it already exists.
        pub fn preserve(mut self, preserve: bool) -> Self {
            self.preserve = preserve;
            self
        }
    }
}

impl<S, Body> Service<Request<Body>> for HttpProxyAddressService<S>
where
    S: Service<Request<Body>>,
    Body: Send + 'static,
{
    type Output = S::Output;
    type Error = S::Error;

    fn serve(
        &self,
        mut req: Request<Body>,
    ) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send + '_ {
        if let Some(ref proxy_info) = self.proxy_info
            && (!self.preserve || !req.extensions().contains::<ProxyAddress>())
        {
            tracing::trace!(
                server.address = %proxy_info.address.host,
                server.port = proxy_info.address.port,
                "setting proxy address",
            );
            req.extensions_mut().insert(proxy_info.clone());
        }
        self.inner.serve(req)
    }
}