httproxide_client_util/

lib.rs

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::Mutex;
use std::task::{Context, Poll};
use std::time::Duration;

use futures_util::future::BoxFuture;
use hyper::body::HttpBody;
use hyper::client::connect::Connection;
use hyper::client::HttpConnector;
use hyper::{Client as HyperClient, Uri};
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tower::{Service, ServiceExt};

pub type Client<B> = HyperClient<Connector, B>;

#[derive(Clone, Debug)]
pub struct Connector {
    #[cfg(feature = "https")]
    http: hyper_rustls::HttpsConnector<HttpConnector>,
    #[cfg(not(feature = "https"))]
    http: HttpConnector,
    #[cfg(feature = "unix")]
    unix: hyperlocal::UnixConnector,
}

pub enum Stream {
    #[cfg(feature = "https")]
    Http(hyper_rustls::MaybeHttpsStream<tokio::net::TcpStream>),
    #[cfg(not(feature = "https"))]
    Http(tokio::net::TcpStream),
    #[cfg(feature = "unix")]
    Unix(hyperlocal::UnixStream),
}

impl Service<Uri> for Connector {
    type Response = Stream;
    type Error = tower::BoxError;
    type Future = BoxFuture<'static, Result<Stream, Self::Error>>;

    fn poll_ready(&mut self, _cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, dst: Uri) -> Self::Future {
        #[cfg(feature = "unix")]
        if dst.scheme_str() == Some("unix") {
            let clone = self.unix.clone();
            let unix = std::mem::replace(&mut self.unix, clone);
            return Box::pin(async move { Ok(Stream::Unix(unix.oneshot(dst).await?)) });
        }
        let clone = self.http.clone();
        let http = std::mem::replace(&mut self.http, clone);
        Box::pin(async move { Ok(Stream::Http(http.oneshot(dst).await?)) })
    }
}

impl AsyncRead for Stream {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context,
        buf: &mut ReadBuf,
    ) -> Poll<std::io::Result<()>> {
        match self.get_mut() {
            Stream::Http(s) => Pin::new(s).poll_read(cx, buf),
            #[cfg(feature = "unix")]
            Stream::Unix(s) => Pin::new(s).poll_read(cx, buf),
        }
    }
}

impl AsyncWrite for Stream {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        match self.get_mut() {
            Stream::Http(ref mut s) => Pin::new(s).poll_write(cx, buf),
            #[cfg(feature = "unix")]
            Stream::Unix(ref mut s) => Pin::new(s).poll_write(cx, buf),
        }
    }
    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<std::io::Result<()>> {
        match self.get_mut() {
            Stream::Http(ref mut s) => Pin::new(s).poll_flush(cx),
            #[cfg(feature = "unix")]
            Stream::Unix(ref mut s) => Pin::new(s).poll_flush(cx),
        }
    }
    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll<std::io::Result<()>> {
        match self.get_mut() {
            Stream::Http(ref mut s) => Pin::new(s).poll_flush(cx),
            #[cfg(feature = "unix")]
            Stream::Unix(ref mut s) => Pin::new(s).poll_flush(cx),
        }
    }
}

impl Connection for Stream {
    fn connected(&self) -> hyper::client::connect::Connected {
        match self {
            Stream::Http(s) => s.connected(),
            #[cfg(feature = "unix")]
            Stream::Unix(s) => s.connected(),
        }
    }
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, Default)]
pub struct ClientConfig {
    dangerous_skip_cert_check: bool,
}

lazy_static::lazy_static! {
    static ref CACHE: Mutex<HashMap<ClientConfig, Box<dyn std::any::Any + Send>>> = {
        Mutex::new(HashMap::new())
    };
}

pub fn clear_cache() {
    *(*CACHE).lock().unwrap() = HashMap::new();
}

#[cfg(feature = "https")]
struct NoCertVerifier {}

#[cfg(feature = "https")]
impl rustls::client::ServerCertVerifier for NoCertVerifier {
    fn verify_server_cert(
        &self,
        _end_entity: &rustls::Certificate,
        _intermediates: &[rustls::Certificate],
        _server_name: &rustls::client::ServerName,
        _scts: &mut dyn Iterator<Item = &[u8]>,
        _ocsp_response: &[u8],
        _now: std::time::SystemTime,
    ) -> Result<rustls::client::ServerCertVerified, rustls::Error> {
        Ok(rustls::client::ServerCertVerified::assertion())
    }
}

fn new_client<B>(cfg: ClientConfig) -> anyhow::Result<Client<B>>
where
    B: HttpBody + Send + 'static,
    B::Data: Send,
{
    let mut http = HttpConnector::new();

    #[cfg(feature = "https")]
    let https = {
        use hyper_rustls::ConfigBuilderExt;

        http.enforce_http(false);

        let tls_config = {
            rustls::ClientConfig::builder()
                .with_safe_defaults()
        };

        let tls_config = if cfg.dangerous_skip_cert_check {
            tls_config
                .with_custom_certificate_verifier(Arc::new(NoCertVerifier {}))
                .with_no_client_auth()
        } else {
            tls_config
                .with_native_roots()
                .with_no_client_auth()
        };

        let tls = hyper_rustls::HttpsConnectorBuilder::new()
            .with_tls_config(tls_config)
            .https_or_http()
            .enable_http1();

        #[cfg(feature = "http2")]
        {
            tls.enable_http2().wrap_connector(http)
        }

        #[cfg(not(feature = "http2"))]
        {
            tls.wrap_connector(http)
        }
    };

    let connector = Connector {
        #[cfg(feature = "https")]
        http: https,
        #[cfg(not(feature = "https"))]
        http,
        #[cfg(feature = "unix")]
        unix: hyperlocal::UnixConnector,
    };

    let client = HyperClient::builder()
        .pool_idle_timeout(Duration::from_secs(30))
        .build(connector);

    Ok(client)
}

pub fn get_client<B>(cfg: ClientConfig) -> anyhow::Result<Client<B>>
where
    B: HttpBody + Send + 'static,
    B::Data: Send,
{
    let mut cache = (*CACHE).lock().unwrap();
    if let Some(ref val) = cache.get(&cfg).and_then(|x| x.downcast_ref::<Client<B>>()) {
        Ok((*val).clone())
    } else {
        let new_val = new_client(cfg.clone())?;
        cache.insert(cfg, Box::new(new_val.clone()));
        Ok(new_val)
    }
}