litellm_rs/utils/net/
http.rs

1//! Shared HTTP client for optimal connection pooling
2//!
3//! This module provides a high-performance shared HTTP client with connection reuse.
4//! Using a shared client avoids the overhead of creating new connection pools and
5//! DNS resolution caches for each request.
6//!
7//! # Performance Benefits
8//!
9//! - **Connection Reuse**: Keeps TCP connections alive across requests
10//! - **DNS Caching**: Avoids repeated DNS lookups
11//! - **HTTP/2 Multiplexing**: Multiple requests over a single connection
12//! - **Reduced Latency**: 20-50% improvement in request latency
13//!
14//! # Usage
15//!
16//! ```rust,no_run
17//! # use litellm_rs::utils::net::http::get_shared_client;
18//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
19//! let client = get_shared_client();
20//! let response = client.get("https://api.openai.com").send().await?;
21//! # Ok(())
22//! # }
23//! ```
24
25use dashmap::DashMap;
26use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
27use reqwest::{
28    Body, Client, ClientBuilder, IntoUrl, Method, RequestBuilder, Version, multipart, redirect,
29};
30use serde::Serialize;
31use std::error::Error as _;
32use std::io;
33use std::net::{SocketAddr, ToSocketAddrs};
34use std::sync::{Arc, OnceLock};
35use std::time::Duration;
36use tracing::{debug, warn};
37
38use crate::core::net::{
39    ProviderEndpointAccess, ProviderEndpointPolicy, SsrfError, is_provider_endpoint_ip_allowed,
40    validate_outbound_url_without_resolution,
41};
42
43/// DNS resolver that rejects private/reserved IP addresses at resolution time.
44///
45/// This mitigates DNS-rebinding attacks: even if a hostname resolves to a public IP
46/// at config-validation time, every actual request re-validates the resolved address,
47/// so a later rebind to an internal IP will be caught and rejected.
48trait HostResolver: Send + Sync {
49    fn resolve(&self, host: &str) -> io::Result<Vec<SocketAddr>>;
50}
51
52struct SystemHostResolver;
53
54impl HostResolver for SystemHostResolver {
55    fn resolve(&self, host: &str) -> io::Result<Vec<SocketAddr>> {
56        (host, 0u16)
57            .to_socket_addrs()
58            .map(|addresses| addresses.collect())
59    }
60}
61
62struct PolicyDnsResolver {
63    access: ProviderEndpointAccess,
64    resolver: Arc<dyn HostResolver>,
65}
66
67#[derive(Debug, thiserror::Error)]
68enum ProviderEndpointPolicyRequestError {
69    #[error("Host resolution returned no addresses (SSRF protection)")]
70    EmptyDnsAnswer,
71    #[error("Host resolves to a disallowed address (SSRF protection)")]
72    DisallowedDnsAddress,
73    #[error("Redirect target failed SSRF validation")]
74    RedirectTarget,
75}
76
77fn reqwest_error_is_endpoint_policy(error: &reqwest::Error) -> bool {
78    let mut source = error.source();
79    while let Some(current) = source {
80        if current
81            .downcast_ref::<ProviderEndpointPolicyRequestError>()
82            .is_some()
83            || current
84                .downcast_ref::<io::Error>()
85                .and_then(io::Error::get_ref)
86                .and_then(|inner| inner.downcast_ref::<ProviderEndpointPolicyRequestError>())
87                .is_some()
88        {
89            return true;
90        }
91        source = current.source();
92    }
93    false
94}
95
96impl reqwest::dns::Resolve for PolicyDnsResolver {
97    fn resolve(&self, name: reqwest::dns::Name) -> reqwest::dns::Resolving {
98        let host = name.as_str().to_owned();
99        let resolver = Arc::clone(&self.resolver);
100        let access = self.access;
101        Box::pin(async move {
102            let addrs = tokio::task::spawn_blocking(move || resolver.resolve(&host))
103                .await
104                .map_err(io::Error::other)??;
105            let safe = validate_provider_addresses(access, addrs)?;
106
107            Ok(Box::new(safe.into_iter()) as reqwest::dns::Addrs)
108        })
109    }
110}
111
112fn validate_provider_addresses(
113    access: ProviderEndpointAccess,
114    addrs: Vec<SocketAddr>,
115) -> io::Result<Vec<SocketAddr>> {
116    if addrs.is_empty() {
117        return Err(io::Error::other(
118            ProviderEndpointPolicyRequestError::EmptyDnsAnswer,
119        ));
120    }
121    if addrs
122        .iter()
123        .any(|addr| !is_provider_endpoint_ip_allowed(access, &addr.ip()))
124    {
125        return Err(io::Error::other(
126            ProviderEndpointPolicyRequestError::DisallowedDnsAddress,
127        ));
128    }
129    Ok(addrs)
130}
131
132fn ssrf_safe_redirect_policy() -> redirect::Policy {
133    redirect::Policy::custom(|attempt| {
134        if validate_outbound_url_without_resolution(attempt.url()).is_err() {
135            debug!("blocked redirect target by outbound endpoint policy");
136            return attempt.error(ProviderEndpointPolicyRequestError::RedirectTarget);
137        }
138
139        redirect::Policy::limited(10).redirect(attempt)
140    })
141}
142
143fn provider_redirect_policy(policy: ProviderEndpointPolicy) -> redirect::Policy {
144    if policy.access() == ProviderEndpointAccess::PrivateNetwork {
145        return redirect::Policy::none();
146    }
147
148    redirect::Policy::custom(move |attempt| {
149        if policy
150            .validate_url_without_resolution(attempt.url())
151            .is_err()
152        {
153            debug!("blocked provider redirect target by endpoint policy");
154            return attempt.error(ProviderEndpointPolicyRequestError::RedirectTarget);
155        }
156        redirect::Policy::limited(10).redirect(attempt)
157    })
158}
159
160#[derive(Debug, thiserror::Error)]
161pub enum ProviderHttpClientError {
162    #[error(transparent)]
163    Endpoint(#[from] SsrfError),
164    #[error(transparent)]
165    Request(#[from] reqwest::Error),
166}
167
168impl ProviderHttpClientError {
169    pub(crate) fn is_endpoint_policy(&self) -> bool {
170        matches!(self, Self::Endpoint(_))
171    }
172}
173
174/// Request builder that cannot expose or replace its policy-bound HTTP client.
175pub struct ProviderRequestBuilder {
176    inner: RequestBuilder,
177}
178
179impl ProviderRequestBuilder {
180    fn map(self, transform: impl FnOnce(RequestBuilder) -> RequestBuilder) -> Self {
181        Self {
182            inner: transform(self.inner),
183        }
184    }
185
186    pub fn header<K, V>(self, key: K, value: V) -> Self
187    where
188        HeaderName: TryFrom<K>,
189        <HeaderName as TryFrom<K>>::Error: Into<http::Error>,
190        HeaderValue: TryFrom<V>,
191        <HeaderValue as TryFrom<V>>::Error: Into<http::Error>,
192    {
193        self.map(|inner| inner.header(key, value))
194    }
195
196    pub fn headers(self, headers: HeaderMap) -> Self {
197        self.map(|inner| inner.headers(headers))
198    }
199
200    pub fn basic_auth<U: std::fmt::Display, P: std::fmt::Display>(
201        self,
202        username: U,
203        password: Option<P>,
204    ) -> Self {
205        self.map(|inner| inner.basic_auth(username, password))
206    }
207
208    pub fn bearer_auth<T: std::fmt::Display>(self, token: T) -> Self {
209        self.map(|inner| inner.bearer_auth(token))
210    }
211
212    pub fn query<T: Serialize + ?Sized>(self, query: &T) -> Self {
213        self.map(|inner| inner.query(query))
214    }
215
216    pub fn body<T: Into<Body>>(self, body: T) -> Self {
217        self.map(|inner| inner.body(body))
218    }
219
220    pub fn timeout(self, timeout: Duration) -> Self {
221        self.map(|inner| inner.timeout(timeout))
222    }
223
224    pub fn version(self, version: Version) -> Self {
225        self.map(|inner| inner.version(version))
226    }
227
228    pub fn form<T: Serialize + ?Sized>(self, form: &T) -> Self {
229        self.map(|inner| inner.form(form))
230    }
231
232    pub fn json<T: Serialize + ?Sized>(self, json: &T) -> Self {
233        self.map(|inner| inner.json(json))
234    }
235
236    pub fn multipart(self, form: multipart::Form) -> Self {
237        self.map(|inner| inner.multipart(form))
238    }
239
240    pub async fn send(self) -> Result<reqwest::Response, reqwest::Error> {
241        self.inner.send().await
242    }
243}
244
245#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
246enum ProviderClientMode {
247    Request,
248    Streaming,
249    StreamingNoRedirect,
250    NoRedirect,
251}
252
253#[derive(Debug, Clone, PartialEq, Eq, Hash)]
254struct ProviderClientCacheKey {
255    policy: ProviderEndpointPolicy,
256    timeout: Duration,
257    mode: ProviderClientMode,
258}
259
260static PROVIDER_CLIENT_CACHE: OnceLock<DashMap<ProviderClientCacheKey, Arc<Client>>> =
261    OnceLock::new();
262
263/// HTTP client that enforces one provider endpoint policy at every request boundary.
264#[derive(Clone)]
265pub struct ProviderHttpClient {
266    client: Arc<Client>,
267    policy: ProviderEndpointPolicy,
268}
269
270impl std::fmt::Debug for ProviderHttpClient {
271    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
272        f.debug_struct("ProviderHttpClient")
273            .field("policy", &self.policy)
274            .finish_non_exhaustive()
275    }
276}
277
278impl ProviderHttpClient {
279    pub(crate) fn request_error_is_endpoint_policy(error: &reqwest::Error) -> bool {
280        reqwest_error_is_endpoint_policy(error)
281    }
282
283    pub fn new(
284        policy: ProviderEndpointPolicy,
285        timeout: Duration,
286    ) -> Result<Self, ProviderHttpClientError> {
287        Self::cached(policy, timeout, ProviderClientMode::Request)
288    }
289
290    pub fn streaming(policy: ProviderEndpointPolicy) -> Result<Self, ProviderHttpClientError> {
291        Self::cached(
292            policy,
293            Duration::from_secs(0),
294            ProviderClientMode::Streaming,
295        )
296    }
297
298    pub fn streaming_no_redirect(
299        policy: ProviderEndpointPolicy,
300    ) -> Result<Self, ProviderHttpClientError> {
301        Self::cached(
302            policy,
303            Duration::from_secs(0),
304            ProviderClientMode::StreamingNoRedirect,
305        )
306    }
307
308    pub fn no_redirect(
309        policy: ProviderEndpointPolicy,
310        timeout: Duration,
311    ) -> Result<Self, ProviderHttpClientError> {
312        Self::cached(policy, timeout, ProviderClientMode::NoRedirect)
313    }
314
315    fn cached(
316        policy: ProviderEndpointPolicy,
317        timeout: Duration,
318        mode: ProviderClientMode,
319    ) -> Result<Self, ProviderHttpClientError> {
320        let key = ProviderClientCacheKey {
321            policy: policy.clone(),
322            timeout,
323            mode,
324        };
325        let client = PROVIDER_CLIENT_CACHE
326            .get_or_init(DashMap::new)
327            .entry(key)
328            .or_try_insert_with(|| {
329                Self::build_client(&policy, timeout, mode, Arc::new(SystemHostResolver))
330                    .map(Arc::new)
331            })?
332            .clone();
333        Ok(Self { client, policy })
334    }
335
336    #[cfg(test)]
337    fn build(
338        policy: ProviderEndpointPolicy,
339        timeout: Duration,
340        mode: ProviderClientMode,
341        resolver: Arc<dyn HostResolver>,
342    ) -> Result<Self, ProviderHttpClientError> {
343        let client = Arc::new(Self::build_client(&policy, timeout, mode, resolver)?);
344        Ok(Self { client, policy })
345    }
346
347    fn build_client(
348        policy: &ProviderEndpointPolicy,
349        timeout: Duration,
350        mode: ProviderClientMode,
351        resolver: Arc<dyn HostResolver>,
352    ) -> Result<Client, ProviderHttpClientError> {
353        let resolver = Arc::new(PolicyDnsResolver {
354            access: policy.access(),
355            resolver,
356        });
357        Self::build_client_with_dns_resolver(policy, timeout, mode, resolver)
358    }
359
360    fn build_client_with_dns_resolver<R: reqwest::dns::Resolve + 'static>(
361        policy: &ProviderEndpointPolicy,
362        timeout: Duration,
363        mode: ProviderClientMode,
364        resolver: Arc<R>,
365    ) -> Result<Client, ProviderHttpClientError> {
366        let config = HttpClientPoolConfig::default();
367        let builder = match mode {
368            ProviderClientMode::Streaming | ProviderClientMode::StreamingNoRedirect => {
369                { ClientBuilder::new() }
370                    .pool_max_idle_per_host(config.pool_max_idle_per_host)
371                    .pool_idle_timeout(config.pool_idle_timeout)
372                    .connect_timeout(config.connect_timeout)
373                    .tcp_keepalive(config.tcp_keepalive)
374                    .tcp_nodelay(true)
375                    .user_agent(config.user_agent)
376            }
377            ProviderClientMode::Request | ProviderClientMode::NoRedirect => {
378                create_client_builder_with_config(timeout, &config)
379            }
380        };
381        let redirect_policy = if matches!(
382            mode,
383            ProviderClientMode::NoRedirect | ProviderClientMode::StreamingNoRedirect
384        ) {
385            redirect::Policy::none()
386        } else {
387            provider_redirect_policy(policy.clone())
388        };
389        let client = builder
390            .no_proxy()
391            .dns_resolver(resolver)
392            .redirect(redirect_policy)
393            .build()?;
394        Ok(client)
395    }
396
397    pub fn request<U: IntoUrl>(
398        &self,
399        method: Method,
400        url: U,
401    ) -> Result<ProviderRequestBuilder, ProviderHttpClientError> {
402        let url = url.into_url()?;
403        self.policy.validate_url_without_resolution(&url)?;
404        Ok(ProviderRequestBuilder {
405            inner: self.client.request(method, url),
406        })
407    }
408
409    pub fn get<U: IntoUrl>(
410        &self,
411        url: U,
412    ) -> Result<ProviderRequestBuilder, ProviderHttpClientError> {
413        self.request(Method::GET, url)
414    }
415
416    pub fn post<U: IntoUrl>(
417        &self,
418        url: U,
419    ) -> Result<ProviderRequestBuilder, ProviderHttpClientError> {
420        self.request(Method::POST, url)
421    }
422}
423
424/// Configuration for the HTTP client pool
425#[derive(Debug, Clone)]
426pub struct HttpClientPoolConfig {
427    /// Maximum idle connections per host
428    pub pool_max_idle_per_host: usize,
429    /// Idle connection timeout
430    pub pool_idle_timeout: Duration,
431    /// Connection timeout
432    pub connect_timeout: Duration,
433    /// TCP keepalive interval
434    pub tcp_keepalive: Duration,
435    /// User agent string
436    pub user_agent: &'static str,
437}
438
439impl Default for HttpClientPoolConfig {
440    fn default() -> Self {
441        Self {
442            pool_max_idle_per_host: 100, // Increased for high throughput
443            pool_idle_timeout: Duration::from_secs(90),
444            connect_timeout: Duration::from_secs(10),
445            tcp_keepalive: Duration::from_secs(60),
446            user_agent: "LiteLLM-RS/0.1.0",
447        }
448    }
449}
450
451/// Shared HTTP client instance with optimized settings
452static SHARED_HTTP_CLIENT: OnceLock<Client> = OnceLock::new();
453
454/// Timeout-specific client cache (keyed by milliseconds)
455static TIMEOUT_CLIENT_CACHE: OnceLock<DashMap<u64, Arc<Client>>> = OnceLock::new();
456
457/// Timeout-specific SSRF-safe client cache (keyed by milliseconds)
458static SSRF_SAFE_TIMEOUT_CLIENT_CACHE: OnceLock<DashMap<u64, Arc<Client>>> = OnceLock::new();
459
460/// Create a reqwest client builder with unified pool/timeout defaults.
461pub fn create_client_builder_with_config(
462    timeout: Duration,
463    config: &HttpClientPoolConfig,
464) -> ClientBuilder {
465    ClientBuilder::new()
466        // Connection pool settings
467        .pool_max_idle_per_host(config.pool_max_idle_per_host)
468        .pool_idle_timeout(config.pool_idle_timeout)
469        // Request timeouts
470        .timeout(timeout)
471        .connect_timeout(config.connect_timeout)
472        // TCP optimizations
473        .tcp_keepalive(config.tcp_keepalive)
474        .tcp_nodelay(true)
475        // User agent
476        .user_agent(config.user_agent)
477}
478
479/// Create a reqwest client builder with default pool configuration.
480pub fn create_client_builder(timeout: Duration) -> ClientBuilder {
481    create_client_builder_with_config(timeout, &HttpClientPoolConfig::default())
482}
483
484/// Get the shared HTTP client instance
485///
486/// This client uses a default timeout of 30 seconds. For custom timeouts,
487/// use `get_client_with_timeout`.
488pub fn get_shared_client() -> &'static Client {
489    SHARED_HTTP_CLIENT.get_or_init(|| {
490        debug!("Initializing shared HTTP client with optimized settings");
491        create_optimized_client(Duration::from_secs(30))
492    })
493}
494
495/// Get or create a client with a specific timeout
496///
497/// Clients are cached by timeout duration (in milliseconds) to avoid creating
498/// multiple clients with the same configuration.
499pub fn get_client_with_timeout(timeout: Duration) -> Arc<Client> {
500    let cache = TIMEOUT_CLIENT_CACHE.get_or_init(DashMap::new);
501    let timeout_millis = timeout.as_millis().min(u64::MAX as u128) as u64;
502
503    cache
504        .entry(timeout_millis)
505        .or_insert_with(|| {
506            debug!(timeout_millis, "Creating cached HTTP client for timeout");
507            Arc::new(create_optimized_client(timeout))
508        })
509        .clone()
510}
511
512/// Get or create a client with a specific timeout, returning errors on failure
513///
514/// This is useful when caller error semantics must be preserved.
515pub fn get_client_with_timeout_fallible(timeout: Duration) -> Result<Arc<Client>, reqwest::Error> {
516    let cache = TIMEOUT_CLIENT_CACHE.get_or_init(DashMap::new);
517    let timeout_millis = timeout.as_millis().min(u64::MAX as u128) as u64;
518
519    if let Some(existing) = cache.get(&timeout_millis) {
520        return Ok(existing.clone());
521    }
522
523    let client = Arc::new(create_custom_client(timeout)?);
524    cache.insert(timeout_millis, client.clone());
525    Ok(client)
526}
527
528/// Create an optimized HTTP client with the given timeout
529fn create_optimized_client(timeout: Duration) -> Client {
530    let config = HttpClientPoolConfig::default();
531
532    create_client_builder_with_config(timeout, &config)
533        .build()
534        .unwrap_or_else(|e| {
535            warn!(
536                "Failed to create optimized HTTP client, falling back to default: {}",
537                e
538            );
539            Client::builder()
540                .build()
541                .expect("default HTTP client should build")
542        })
543}
544
545/// Create a custom HTTP client with specific timeout and pool configuration.
546pub fn create_custom_client_with_config(
547    timeout: Duration,
548    config: &HttpClientPoolConfig,
549) -> Result<Client, reqwest::Error> {
550    create_client_builder_with_config(timeout, config).build()
551}
552
553/// Create a custom HTTP client with specific timeout
554///
555/// Use this when you need a one-off client that won't be reused.
556/// For reusable clients, prefer `get_client_with_timeout`.
557pub fn create_custom_client(timeout: Duration) -> Result<Client, reqwest::Error> {
558    create_custom_client_with_config(timeout, &HttpClientPoolConfig::default())
559}
560
561/// Create an HTTP client for long-running SSE streams.
562///
563/// Unlike `create_custom_client`, this client does not set a total request timeout
564/// so streams lasting longer than the configured timeout value won't be cut off.
565/// Only the initial TCP connection is time-bounded via `connect_timeout`.
566pub fn create_streaming_client() -> Result<Client, reqwest::Error> {
567    let config = HttpClientPoolConfig::default();
568    ClientBuilder::new()
569        .pool_max_idle_per_host(config.pool_max_idle_per_host)
570        .pool_idle_timeout(config.pool_idle_timeout)
571        .connect_timeout(config.connect_timeout)
572        .tcp_keepalive(config.tcp_keepalive)
573        .tcp_nodelay(true)
574        .user_agent(config.user_agent)
575        .build()
576}
577
578/// Get or create an HTTP client with SSRF-safe DNS resolution for the given timeout.
579///
580/// Unlike `get_client_with_timeout_fallible`, this client installs `SsrfSafeDnsResolver`
581/// so every request re-validates the resolved IP against private/reserved ranges.
582/// Use this for providers whose endpoint URL is user-controlled to prevent DNS-rebinding attacks.
583pub fn get_ssrf_safe_client_with_timeout_fallible(
584    timeout: Duration,
585) -> Result<Arc<Client>, reqwest::Error> {
586    let cache = SSRF_SAFE_TIMEOUT_CLIENT_CACHE.get_or_init(DashMap::new);
587    let timeout_millis = timeout.as_millis().min(u64::MAX as u128) as u64;
588
589    if let Some(existing) = cache.get(&timeout_millis) {
590        return Ok(existing.clone());
591    }
592
593    let client = Arc::new(create_ssrf_safe_client(
594        timeout,
595        ssrf_safe_redirect_policy(),
596    )?);
597    cache.insert(timeout_millis, client.clone());
598    Ok(client)
599}
600
601fn create_ssrf_safe_client(
602    timeout: Duration,
603    redirect_policy: redirect::Policy,
604) -> Result<Client, reqwest::Error> {
605    create_client_builder_with_config(timeout, &HttpClientPoolConfig::default())
606        .no_proxy()
607        .dns_resolver(Arc::new(PolicyDnsResolver {
608            access: ProviderEndpointAccess::PublicOnly,
609            resolver: Arc::new(SystemHostResolver),
610        }))
611        .redirect(redirect_policy)
612        .build()
613}
614
615/// Create a custom HTTP client with specific timeout and default headers
616pub fn create_custom_client_with_headers(
617    timeout: Duration,
618    default_headers: reqwest::header::HeaderMap,
619) -> Result<Client, reqwest::Error> {
620    create_client_builder(timeout)
621        .default_headers(default_headers)
622        .build()
623}
624
625/// Get statistics about the client cache
626pub fn get_cache_stats() -> HttpClientCacheStats {
627    let cache = TIMEOUT_CLIENT_CACHE.get_or_init(DashMap::new);
628    HttpClientCacheStats {
629        cached_clients: cache.len(),
630        timeout_configs: cache.iter().map(|e| *e.key()).collect(),
631    }
632}
633
634/// Statistics about the HTTP client cache
635#[derive(Debug, Clone)]
636pub struct HttpClientCacheStats {
637    /// Number of cached clients
638    pub cached_clients: usize,
639    /// List of cached timeout configurations (in milliseconds)
640    pub timeout_configs: Vec<u64>,
641}
642
643#[cfg(test)]
644mod tests {
645    use super::*;
646    use std::net::{Ipv4Addr, SocketAddr};
647    use tokio::io::{AsyncReadExt, AsyncWriteExt};
648
649    #[test]
650    fn test_shared_client_creation() {
651        let client = get_shared_client();
652        // Just verify we can get the client without panicking
653        assert!(std::ptr::addr_of!(*client) == std::ptr::addr_of!(*get_shared_client()));
654    }
655
656    #[test]
657    fn test_custom_client_creation() {
658        let client = create_custom_client(Duration::from_secs(15));
659        assert!(client.is_ok());
660    }
661
662    #[test]
663    fn test_client_with_timeout_caching() {
664        let client1 = get_client_with_timeout(Duration::from_secs(60));
665        let client2 = get_client_with_timeout(Duration::from_secs(60));
666
667        // Same timeout should return the same cached client
668        assert!(Arc::ptr_eq(&client1, &client2));
669
670        // Different timeout should return different client
671        let client3 = get_client_with_timeout(Duration::from_secs(120));
672        assert!(!Arc::ptr_eq(&client1, &client3));
673    }
674
675    #[test]
676    fn test_client_with_timeout_fallible_caching() {
677        let client1 = get_client_with_timeout_fallible(Duration::from_millis(1500)).unwrap();
678        let client2 = get_client_with_timeout_fallible(Duration::from_millis(1500)).unwrap();
679
680        assert!(Arc::ptr_eq(&client1, &client2));
681    }
682
683    #[test]
684    fn test_ssrf_safe_dns_filter_rejects_private_and_reserved_addresses() {
685        let addrs = vec![
686            SocketAddr::from((Ipv4Addr::new(93, 184, 216, 34), 443)),
687            SocketAddr::from((Ipv4Addr::new(198, 18, 0, 1), 443)),
688            SocketAddr::from((Ipv4Addr::new(224, 0, 0, 1), 443)),
689            SocketAddr::from((Ipv4Addr::new(240, 0, 0, 1), 443)),
690        ];
691
692        assert!(
693            validate_provider_addresses(ProviderEndpointAccess::PublicOnly, addrs).is_err(),
694            "a mixed DNS answer must fail as a whole"
695        );
696    }
697
698    #[test]
699    fn test_ssrf_safe_client_with_timeout_fallible_caching() {
700        let client1 = match get_ssrf_safe_client_with_timeout_fallible(Duration::from_millis(1500))
701        {
702            Ok(client) => client,
703            Err(error) => panic!("SSRF-safe client should build: {error}"),
704        };
705        let client2 = match get_ssrf_safe_client_with_timeout_fallible(Duration::from_millis(1500))
706        {
707            Ok(client) => client,
708            Err(error) => panic!("SSRF-safe client should build: {error}"),
709        };
710
711        assert!(Arc::ptr_eq(&client1, &client2));
712    }
713
714    #[tokio::test]
715    async fn test_ssrf_safe_redirect_policy_rejects_private_redirect_targets()
716    -> Result<(), Box<dyn std::error::Error>> {
717        let source = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
718        let source_address = source.local_addr()?;
719        let target = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
720        let target_address = target.local_addr()?;
721
722        let server = tokio::spawn(async move {
723            let (mut stream, _) = source.accept().await?;
724            let mut buffer = [0_u8; 1024];
725            let bytes_read = stream.read(&mut buffer).await?;
726            assert!(bytes_read > 0);
727
728            let location = format!("http://{target_address}/private");
729            let body = "redirect";
730            let response = format!(
731                "HTTP/1.1 302 Found\r\nLocation: {location}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
732                body.len()
733            );
734            stream.write_all(response.as_bytes()).await?;
735            Ok::<(), std::io::Error>(())
736        });
737
738        let client = ClientBuilder::new()
739            .redirect(ssrf_safe_redirect_policy())
740            .build()?;
741        let result = client
742            .get(format!("http://{source_address}/redirect"))
743            .send()
744            .await;
745
746        let error = match result {
747            Ok(response) => panic!(
748                "private redirect target should be rejected, got status {}",
749                response.status()
750            ),
751            Err(error) => error,
752        };
753        assert!(error.is_redirect(), "{error:?}");
754        assert!(reqwest_error_is_endpoint_policy(&error), "{error:?}");
755        assert!(
756            tokio::time::timeout(Duration::from_millis(100), target.accept())
757                .await
758                .is_err(),
759            "blocked redirect must not open the target socket"
760        );
761        server.await??;
762        Ok(())
763    }
764
765    #[test]
766    fn test_cache_stats() {
767        // Ensure some clients are cached
768        let _ = get_client_with_timeout(Duration::from_secs(30));
769        let _ = get_client_with_timeout(Duration::from_secs(45));
770
771        let stats = get_cache_stats();
772        assert!(stats.cached_clients >= 2);
773        assert!(stats.timeout_configs.contains(&30_000));
774        assert!(stats.timeout_configs.contains(&45_000));
775    }
776
777    #[test]
778    fn test_pool_config_defaults() {
779        let config = HttpClientPoolConfig::default();
780        assert_eq!(config.pool_max_idle_per_host, 100);
781        assert_eq!(config.pool_idle_timeout, Duration::from_secs(90));
782        assert_eq!(config.connect_timeout, Duration::from_secs(10));
783        assert_eq!(config.tcp_keepalive, Duration::from_secs(60));
784        assert_eq!(config.user_agent, "LiteLLM-RS/0.1.0");
785    }
786}
787
788#[cfg(test)]
789mod provider_tests;
litellm_rs/utils/net/http.rs

litellm_rs/utils/net/
http.rs
