ugi 0.2.1

Runtime-agnostic Rust request client with HTTP/1.1, HTTP/2, HTTP/3, H2C, WebSocket, SSE, and gRPC support
Documentation
use std::collections::HashMap;
use std::time::{Duration, Instant};

use crate::browser_emulation::BrowserProfile;
use crate::request::ProtocolPolicy;
use crate::tls::TlsConfig;
use crate::url::Url;

/// Configuration for the per-host idle connection pool.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PoolConfig {
    /// Maximum number of idle connections kept per host. Set to `0` to disable pooling.
    pub max_idle_per_host: usize,
    /// How long an idle connection is kept before being evicted. `None` means no expiry.
    pub idle_timeout: Option<Duration>,
}

impl Default for PoolConfig {
    fn default() -> Self {
        Self {
            max_idle_per_host: 8,
            idle_timeout: Some(Duration::from_secs(90)),
        }
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub(crate) struct PoolKey {
    scheme: String,
    host: String,
    port: u16,
    tls_backend: Option<&'static str>,
    accept_invalid_certs: bool,
    alpn_protocols: Vec<String>,
    emulation_connection: Option<u64>,
    tls_fingerprint: Option<u64>,
    http2_fingerprint: Option<u64>,
    http3_fingerprint: Option<u64>,
    h2_settings: Option<Vec<u8>>,
}

impl PoolKey {
    pub(crate) fn for_http1(
        url: &Url,
        tls_config: &TlsConfig,
        protocol_policy: ProtocolPolicy,
        browser_profile: Option<&BrowserProfile>,
    ) -> Self {
        Self::new(
            url,
            tls_config,
            tls_config.effective_alpn_protocols(protocol_policy),
            browser_profile,
            None,
        )
    }

    #[cfg(feature = "h2")]
    pub(crate) fn for_h2(
        url: &Url,
        tls_config: &TlsConfig,
        _h2_keepalive_config: crate::request::H2KeepAliveConfig,
        browser_profile: Option<&BrowserProfile>,
        settings_payload: &[u8],
    ) -> crate::Result<Self> {
        let protocols = tls_config
            .validate_h2_alpn()
            .map_err(|message| crate::Error::new(crate::ErrorKind::Transport, message))?;
        Ok(Self::new(
            url,
            tls_config,
            protocols,
            browser_profile,
            Some(settings_payload.to_vec()),
        ))
    }

    #[cfg(feature = "h2")]
    pub(crate) fn for_h2c(
        url: &Url,
        _h2_keepalive_config: crate::request::H2KeepAliveConfig,
        browser_profile: Option<&BrowserProfile>,
        settings_payload: &[u8],
    ) -> crate::Result<Self> {
        if url.scheme() != "http" {
            return Err(crate::Error::new(
                crate::ErrorKind::Transport,
                "h2c requires an http url",
            ));
        }
        Ok(Self {
            scheme: url.scheme().to_owned(),
            host: url.host().to_owned(),
            port: url.effective_port(),
            tls_backend: None,
            accept_invalid_certs: false,
            alpn_protocols: Vec::new(),
            emulation_connection: browser_profile
                .and_then(BrowserProfile::connection_identity_hash),
            tls_fingerprint: None,
            http2_fingerprint: browser_profile.and_then(BrowserProfile::http2_identity_hash),
            http3_fingerprint: None,
            h2_settings: Some(settings_payload.to_vec()),
        })
    }

    #[cfg(feature = "h3")]
    pub(crate) fn for_h3(
        url: &Url,
        tls_config: &TlsConfig,
        browser_profile: Option<&BrowserProfile>,
    ) -> crate::Result<Self> {
        let protocols = tls_config
            .validate_h3_alpn()
            .map_err(|message| crate::Error::new(crate::ErrorKind::Transport, message))?;
        Ok(Self::new(url, tls_config, protocols, browser_profile, None))
    }

    fn new(
        url: &Url,
        tls_config: &TlsConfig,
        alpn_protocols: Vec<String>,
        browser_profile: Option<&BrowserProfile>,
        h2_settings: Option<Vec<u8>>,
    ) -> Self {
        Self {
            scheme: url.scheme().to_owned(),
            host: url.host().to_owned(),
            port: url.effective_port(),
            tls_backend: tls_backend_name(tls_config),
            accept_invalid_certs: tls_config.accept_invalid_certs,
            alpn_protocols,
            emulation_connection: browser_profile
                .and_then(BrowserProfile::connection_identity_hash),
            tls_fingerprint: browser_profile.and_then(BrowserProfile::tls_identity_hash),
            http2_fingerprint: browser_profile.and_then(BrowserProfile::http2_identity_hash),
            http3_fingerprint: browser_profile.and_then(BrowserProfile::http3_identity_hash),
            h2_settings,
        }
    }
}

fn tls_backend_name(tls_config: &TlsConfig) -> Option<&'static str> {
    #[cfg(any(feature = "rustls", feature = "native-tls", feature = "btls-backend"))]
    {
        fn default_backend() -> crate::tls::TlsBackend {
            #[cfg(feature = "rustls")]
            {
                return crate::tls::TlsBackend::Rustls;
            }
            #[cfg(all(not(feature = "rustls"), feature = "native-tls"))]
            {
                return crate::tls::TlsBackend::Native;
            }
            #[cfg(all(
                not(feature = "rustls"),
                not(feature = "native-tls"),
                feature = "btls-backend"
            ))]
            {
                return crate::tls::TlsBackend::Boring;
            }
        }

        let backend = tls_config.backend.unwrap_or_else(default_backend);
        return Some(match backend {
            #[cfg(feature = "rustls")]
            crate::tls::TlsBackend::Rustls => "rustls",
            #[cfg(feature = "native-tls")]
            crate::tls::TlsBackend::Native => "native-tls",
            #[cfg(feature = "btls-backend")]
            crate::tls::TlsBackend::Boring => "boring",
        });
    }

    #[allow(unreachable_code)]
    {
        let _ = tls_config;
        None
    }
}

struct IdleEntry<T> {
    value: T,
    idle_since: Instant,
}

pub(crate) struct IdlePool<T> {
    idle: HashMap<PoolKey, Vec<IdleEntry<T>>>,
}

impl<T> Default for IdlePool<T> {
    fn default() -> Self {
        Self {
            idle: HashMap::new(),
        }
    }
}

impl<T> IdlePool<T> {
    pub(crate) fn checkout(&mut self, key: &PoolKey, config: PoolConfig) -> Option<T> {
        let now = Instant::now();
        let connections = self.idle.get_mut(key)?;
        connections.retain(|entry| match config.idle_timeout {
            Some(timeout) => now.duration_since(entry.idle_since) <= timeout,
            None => true,
        });
        let value = connections.pop().map(|entry| entry.value);
        if connections.is_empty() {
            self.idle.remove(key);
        }
        value
    }

    pub(crate) fn insert(&mut self, key: PoolKey, value: T, config: PoolConfig) {
        if config.max_idle_per_host == 0 {
            return;
        }
        let connections = self.idle.entry(key).or_default();
        if connections.len() >= config.max_idle_per_host {
            connections.remove(0);
        }
        connections.push(IdleEntry {
            value,
            idle_since: Instant::now(),
        });
    }

    pub(crate) fn clear(&mut self) {
        self.idle.clear();
    }
}

#[cfg(all(test, feature = "h2"))]
mod tests {
    use super::*;
    #[cfg(feature = "emulation")]
    use crate::browser_emulation::Emulation;

    #[test]
    fn pool_key_differs_by_h2_settings_payload() {
        let url = Url::parse("https://example.com/").unwrap();
        let tls = TlsConfig::default();
        let keepalive = crate::request::H2KeepAliveConfig::default();

        let key_a = PoolKey::for_h2(&url, &tls, keepalive, None, &[0x01, 0x02, 0x03]).unwrap();
        let key_b = PoolKey::for_h2(&url, &tls, keepalive, None, &[0x09, 0x09]).unwrap();
        let key_a2 = PoolKey::for_h2(&url, &tls, keepalive, None, &[0x01, 0x02, 0x03]).unwrap();

        assert_ne!(key_a, key_b);
        assert_eq!(key_a, key_a2);
    }

    #[cfg(feature = "emulation")]
    #[test]
    fn pool_key_differs_by_emulation_connection_identity() {
        let url = Url::parse("https://example.com/").unwrap();
        let tls = TlsConfig::default().ensure_emulation_backend().unwrap();
        let keepalive = crate::request::H2KeepAliveConfig::default();
        let chrome = Emulation::Chrome136.profile();
        let firefox = Emulation::Firefox128.profile();

        let chrome_key =
            PoolKey::for_h2(&url, &tls, keepalive, Some(&chrome), &[0x01, 0x02]).unwrap();
        let firefox_key =
            PoolKey::for_h2(&url, &tls, keepalive, Some(&firefox), &[0x01, 0x02]).unwrap();

        assert_ne!(chrome_key, firefox_key);
    }
}