use std::{cmp::min, collections::HashMap, fmt};
use bitflags::bitflags;
use bon::Builder;
use kithara_bufpool::BytePool;
use kithara_platform::{time::Duration, traits::FromWithParams};
bitflags! {
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Compression: u8 {
const GZIP = 1 << 0;
const DEFLATE = 1 << 1;
const BROTLI = 1 << 2;
const ZSTD = 1 << 3;
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub enum ImpersonatePreset {
#[default]
Safari,
Chrome,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Headers {
inner: HashMap<String, String>,
}
impl Headers {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn get(&self, key: &str) -> Option<&str> {
self.inner.get(key).map(String::as_str)
}
pub fn insert<K: Into<String>, V: Into<String>>(&mut self, key: K, value: V) {
self.inner.insert(key.into(), value.into());
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
self.inner.iter().map(|(k, v)| (k.as_str(), v.as_str()))
}
}
impl From<HashMap<String, String>> for Headers {
fn from(map: HashMap<String, String>) -> Self {
Self { inner: map }
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RangeSpec {
pub end: Option<u64>,
pub start: u64,
}
impl RangeSpec {
#[must_use]
pub fn new(start: u64, end: Option<u64>) -> Self {
Self { end, start }
}
}
impl fmt::Display for RangeSpec {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.end {
Some(end) => write!(f, "bytes={}-{}", self.start, end),
None => write!(f, "bytes={}-", self.start),
}
}
}
#[derive(Clone, Debug, Builder)]
#[non_exhaustive]
pub struct RetryPolicy {
#[builder(default = Duration::from_millis(100))]
pub base_delay: Duration,
#[builder(default = Duration::from_secs(5))]
pub max_delay: Duration,
#[builder(default = 3)]
pub max_retries: u32,
}
impl Default for RetryPolicy {
fn default() -> Self {
Self::builder().build()
}
}
impl RetryPolicy {
#[must_use]
pub fn new(max_retries: u32, base_delay: Duration, max_delay: Duration) -> Self {
Self {
base_delay,
max_delay,
max_retries,
}
}
#[must_use]
pub fn delay_for_attempt(&self, attempt: u32) -> Duration {
const BACKOFF_BASE: u32 = 2;
if attempt == 0 {
return Duration::ZERO;
}
let exponential_delay = self.base_delay * BACKOFF_BASE.pow(attempt.saturating_sub(1));
min(exponential_delay, self.max_delay)
}
}
#[derive(Clone, Debug, Builder)]
#[non_exhaustive]
pub struct NetOptions {
#[builder(default = Compression::all())]
pub compression: Compression,
#[builder(default = Duration::from_secs(30))]
pub inactivity_timeout: Duration,
pub total_timeout: Option<Duration>,
#[builder(default)]
pub retry_policy: RetryPolicy,
#[builder(default)]
pub is_insecure: bool,
#[builder(default = 8)]
pub pool_max_idle_per_host: usize,
#[builder(default = 32)]
pub body_queue_capacity: usize,
#[builder(default = 16)]
pub body_queue_resume_at: usize,
pub byte_pool: Option<BytePool>,
#[builder(default)]
pub impersonate: ImpersonatePreset,
}
impl Default for NetOptions {
fn default() -> Self {
Self::builder()
.total_timeout(Duration::from_secs(120))
.build()
}
}
impl FromWithParams<Self, Option<BytePool>> for NetOptions {
fn build(options: Self, byte_pool: Option<BytePool>) -> Self {
Self::builder()
.compression(options.compression)
.inactivity_timeout(options.inactivity_timeout)
.maybe_total_timeout(options.total_timeout)
.retry_policy(options.retry_policy)
.is_insecure(options.is_insecure)
.pool_max_idle_per_host(options.pool_max_idle_per_host)
.body_queue_capacity(options.body_queue_capacity)
.body_queue_resume_at(options.body_queue_resume_at)
.maybe_byte_pool(options.byte_pool.or(byte_pool))
.impersonate(options.impersonate)
.build()
}
}
#[cfg(test)]
mod tests {
mod kithara {
pub(crate) use kithara_test_macros::test;
}
use super::*;
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
#[case::empty_headers(Headers::new(), true)]
#[case::headers_with_values({
let mut h = Headers::new();
h.insert("key1", "value1");
h.insert("key2", "value2");
h
}, false)]
async fn test_headers_is_empty(#[case] headers: Headers, #[case] expected_empty: bool) {
assert_eq!(headers.is_empty(), expected_empty);
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
#[case::simple_key("key1", "value1")]
#[case::content_type("Content-Type", "application/json")]
#[case::custom_header("X-Custom-Header", "custom-value")]
async fn test_headers_insert_and_get(#[case] key: &str, #[case] value: &str) {
let mut headers = Headers::new();
headers.insert(key, value);
assert_eq!(headers.get(key), Some(value));
assert_eq!(headers.get("non-existent"), None);
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
async fn test_headers_iter() {
let mut headers = Headers::new();
headers.insert("key1", "value1");
headers.insert("key2", "value2");
headers.insert("key3", "value3");
let mut iterated = HashMap::new();
for (k, v) in headers.iter() {
iterated.insert(k.to_string(), v.to_string());
}
assert_eq!(iterated.len(), 3);
assert_eq!(iterated.get("key1"), Some(&"value1".to_string()));
assert_eq!(iterated.get("key2"), Some(&"value2".to_string()));
assert_eq!(iterated.get("key3"), Some(&"value3".to_string()));
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
async fn test_headers_from_hashmap() {
let mut map = HashMap::new();
map.insert("key1".to_string(), "value1".to_string());
map.insert("key2".to_string(), "value2".to_string());
let headers: Headers = map.into();
assert!(!headers.is_empty());
assert_eq!(headers.get("key1"), Some("value1"));
assert_eq!(headers.get("key2"), Some("value2"));
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
async fn test_headers_default() {
let headers = Headers::default();
assert!(headers.is_empty());
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
#[case::full_range(0, Some(100), "bytes=0-100")]
#[case::open_ended(50, None, "bytes=50-")]
#[case::single_byte(10, Some(10), "bytes=10-10")]
#[case::zero_length(0, Some(0), "bytes=0-0")]
async fn test_range_spec_to_header_value(
#[case] start: u64,
#[case] end: Option<u64>,
#[case] expected_header: &str,
) {
let range = RangeSpec::new(start, end);
assert_eq!(range.to_string(), expected_header);
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
#[case::equal_ranges(RangeSpec::new(0, Some(100)), RangeSpec::new(0, Some(100)), true)]
#[case::different_starts(RangeSpec::new(0, Some(100)), RangeSpec::new(1, Some(100)), false)]
#[case::different_ends(RangeSpec::new(0, Some(100)), RangeSpec::new(0, Some(99)), false)]
#[case::one_open_ended(RangeSpec::new(0, None), RangeSpec::new(0, None), true)]
#[case::mixed_ends(RangeSpec::new(0, Some(100)), RangeSpec::new(0, None), false)]
async fn test_range_spec_partial_eq(
#[case] range1: RangeSpec,
#[case] range2: RangeSpec,
#[case] expected_equal: bool,
) {
assert_eq!(range1 == range2, expected_equal);
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
async fn test_range_spec_debug() {
let range = RangeSpec::new(10, Some(20));
let debug_output = format!("{:?}", range);
assert!(debug_output.contains("RangeSpec"));
assert!(debug_output.contains("start: 10"));
assert!(debug_output.contains("end: Some(20)"));
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
async fn test_range_spec_clone() {
let range1 = RangeSpec::new(10, Some(20));
let range2 = range1.clone();
assert_eq!(range1, range2);
assert_eq!(range1.start, range2.start);
assert_eq!(range1.end, range2.end);
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
async fn test_retry_policy_default() {
let policy = RetryPolicy::default();
assert_eq!(policy.max_retries, 3);
assert_eq!(policy.base_delay, Duration::from_millis(100));
assert_eq!(policy.max_delay, Duration::from_secs(5));
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
#[case(1, Duration::from_millis(50), Duration::from_secs(1))]
#[case(5, Duration::from_millis(100), Duration::from_secs(2))]
#[case(10, Duration::from_millis(200), Duration::from_secs(10))]
async fn test_retry_policy_new(
#[case] max_retries: u32,
#[case] base_delay: Duration,
#[case] max_delay: Duration,
) {
let policy = RetryPolicy::new(max_retries, base_delay, max_delay);
assert_eq!(policy.max_retries, max_retries);
assert_eq!(policy.base_delay, base_delay);
assert_eq!(policy.max_delay, max_delay);
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
#[case(0, Duration::ZERO)]
#[case(1, Duration::from_millis(100))]
#[case(2, Duration::from_millis(200))]
#[case(3, Duration::from_millis(400))]
#[case(4, Duration::from_millis(800))]
#[case(5, Duration::from_millis(1600))]
#[case(10, Duration::from_secs(5))]
#[case(20, Duration::from_secs(5))]
async fn test_retry_policy_delay_for_attempt_default(
#[case] attempt: u32,
#[case] expected_delay: Duration,
) {
let policy = RetryPolicy::default();
let delay = policy.delay_for_attempt(attempt);
assert_eq!(delay, expected_delay);
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
#[case(
1,
Duration::from_millis(50),
Duration::from_millis(200),
0,
Duration::ZERO
)]
#[case(
1,
Duration::from_millis(50),
Duration::from_millis(200),
1,
Duration::from_millis(50)
)]
#[case(
1,
Duration::from_millis(50),
Duration::from_millis(200),
2,
Duration::from_millis(100)
)]
#[case(
1,
Duration::from_millis(50),
Duration::from_millis(200),
3,
Duration::from_millis(200)
)]
#[case(
1,
Duration::from_millis(50),
Duration::from_millis(200),
4,
Duration::from_millis(200)
)]
async fn test_retry_policy_delay_for_attempt_custom(
#[case] max_retries: u32,
#[case] base_delay: Duration,
#[case] max_delay: Duration,
#[case] attempt: u32,
#[case] expected_delay: Duration,
) {
let policy = RetryPolicy::new(max_retries, base_delay, max_delay);
let delay = policy.delay_for_attempt(attempt);
assert_eq!(delay, expected_delay);
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
async fn test_retry_policy_debug() {
let policy = RetryPolicy::default();
let debug_output = format!("{:?}", policy);
assert!(debug_output.contains("RetryPolicy"));
assert!(debug_output.contains("max_retries: 3"));
assert!(debug_output.contains("base_delay"));
assert!(debug_output.contains("max_delay"));
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
async fn test_retry_policy_clone() {
let policy1 = RetryPolicy::new(5, Duration::from_millis(100), Duration::from_secs(2));
let policy2 = policy1.clone();
assert_eq!(policy1.max_retries, policy2.max_retries);
assert_eq!(policy1.base_delay, policy2.base_delay);
assert_eq!(policy1.max_delay, policy2.max_delay);
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
#[case::start_equals_end(10, Some(10), "bytes=10-10")]
#[case::start_greater_than_end(20, Some(10), "bytes=20-10")]
#[case::max_values(u64::MAX, Some(u64::MAX), &format!("bytes={}-{}", u64::MAX, u64::MAX))]
async fn test_range_spec_edge_cases(
#[case] start: u64,
#[case] end: Option<u64>,
#[case] expected_header: &str,
) {
let range = RangeSpec::new(start, end);
assert_eq!(range.to_string(), expected_header);
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
#[case::zero_max_retries(0, Duration::from_millis(100), Duration::from_secs(1))]
#[case::large_max_retries(100, Duration::from_millis(10), Duration::from_secs(10))]
#[case::zero_base_delay(3, Duration::ZERO, Duration::from_secs(1))]
#[case::zero_max_delay(3, Duration::from_millis(100), Duration::ZERO)]
async fn test_retry_policy_edge_cases(
#[case] max_retries: u32,
#[case] base_delay: Duration,
#[case] max_delay: Duration,
) {
let policy = RetryPolicy::new(max_retries, base_delay, max_delay);
for attempt in 0..=5 {
let delay = policy.delay_for_attempt(attempt);
assert!(delay >= Duration::ZERO);
assert!(delay <= max_delay);
if base_delay == Duration::ZERO {
assert_eq!(delay, Duration::ZERO);
}
}
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
#[case(10)]
#[case(20)]
async fn test_retry_policy_large_attempts(#[case] attempt: u32) {
let policy = RetryPolicy::default();
let delay = policy.delay_for_attempt(attempt);
assert!(delay <= policy.max_delay);
assert!(delay >= Duration::ZERO);
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
#[case::with_spaces("X-Custom Header", "value with spaces")]
#[case::with_unicode("X-Emoji", "🎉")]
#[case::with_special_chars("X-Special", "a\tb\nc")]
#[case::empty_value("X-Empty", "")]
async fn test_headers_special_characters(#[case] key: &str, #[case] value: &str) {
let mut headers = Headers::new();
headers.insert(key, value);
assert_eq!(headers.get(key), Some(value));
}
#[kithara::test(tokio, timeout(Duration::from_secs(5)))]
async fn test_headers_case_sensitive() {
let mut headers = Headers::new();
headers.insert("Content-Type", "application/json");
headers.insert("content-type", "text/plain");
assert_eq!(headers.get("Content-Type"), Some("application/json"));
assert_eq!(headers.get("content-type"), Some("text/plain"));
assert_ne!(headers.get("Content-Type"), headers.get("content-type"));
}
}