volo-http 0.5.5

use std::borrow::Cow;

use faststr::FastStr;
use http::{header, uri::Uri};
use motore::{layer::Layer, service::Service};
use volo::{client::Apply, context::Context, net::Address};

use crate::{
    client::Target,
    context::ClientContext,
    error::{ClientError, client::Result},
    request::Request,
};

/// Set a [`Target`] as the destination forcely.
///
/// The layer only sets destination address, including target address, target domain name (for DNS)
/// and SNI of target host (if using HTTPS).
///
/// Note that this layer MUST be set as an outer layer because outer layers runs before service
/// discover (DNS).
pub struct TargetLayer {
    target: Target,
    service_name: FastStr,
}

impl TargetLayer {
    /// Create a [`TargetLayer`] via a [`Target`].
    ///
    /// This layer will set the [`Target`] as request destination.
    pub const fn new(target: Target) -> Self {
        Self {
            target,
            service_name: FastStr::empty(),
        }
    }

    /// Create a [`TargetLayer`] via an address.
    pub fn new_address<A>(addr: A) -> Self
    where
        A: Into<Address>,
    {
        let addr = addr.into();
        let target = Target::from(addr);
        Self {
            target,
            service_name: FastStr::empty(),
        }
    }

    /// Create a [`TargetLayer`] via a host name.
    pub fn new_host<S>(host: S) -> Self
    where
        S: Into<Cow<'static, str>>,
    {
        let target = Target::from_host(host);
        Self {
            target,
            service_name: FastStr::empty(),
        }
    }

    /// Create a [`Target`] from a [`Uri`].
    ///
    /// ## Panics
    ///
    /// This function will panic if the given [`Uri`] cannot be parsed to a [`Target`].
    pub fn from_uri(uri: &Uri) -> Self {
        let target = Target::from_uri(uri).expect("invalid uri for building target");
        Self {
            target,
            service_name: FastStr::empty(),
        }
    }

    /// Set a service name for current [`TargetLayer`].
    ///
    /// When this layer sets [`Target`] as the destination, the service name is also set, and if the
    /// request uses HTTPS, the service name will be used as the SNI.
    pub fn with_service_name<S>(mut self, service_name: S) -> Self
    where
        S: Into<FastStr>,
    {
        self.service_name = service_name.into();
        self
    }
}

impl<S> Layer<S> for TargetLayer {
    type Service = TargetService<S>;

    fn layer(self, inner: S) -> Self::Service {
        TargetService {
            inner,
            target: self.target,
            service_name: self.service_name,
        }
    }
}

/// [`Service`] generated by [`TargetLayer`].
///
/// For more details, please refer to [`TargetLayer`].
pub struct TargetService<S> {
    inner: S,
    target: Target,
    service_name: FastStr,
}

impl<B, S> Service<ClientContext, Request<B>> for TargetService<S>
where
    B: Send,
    S: Service<ClientContext, Request<B>, Error = ClientError> + Send + Sync,
{
    type Response = S::Response;
    type Error = S::Error;

    async fn call(
        &self,
        cx: &mut ClientContext,
        mut req: Request<B>,
    ) -> Result<Self::Response, Self::Error> {
        self.target.clone().apply(cx)?;
        if !self.service_name.is_empty() {
            cx.rpc_info_mut()
                .callee_mut()
                .set_service_name(self.service_name.clone());
        }
        // Since `Host` is one of the outermost layers, it cannot know that `Target` has been
        // modified here, so we need to manually update `Host` here.
        if !req.headers().contains_key(header::HOST) {
            if let Some(host) = super::header::gen_host(&self.target) {
                req.headers_mut().insert(header::HOST, host);
            }
        }
        self.inner.call(cx, req).await
    }
}