use std::collections::HashSet;
use crate::range::RangeHandling;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[allow(clippy::manual_non_exhaustive)]
pub enum StreamingDecision {
Buffer,
SkipCache,
#[doc(hidden)]
StreamThrough,
StreamIfPossible,
}
#[derive(Debug, Clone)]
pub struct StreamingPolicy {
pub enabled: bool,
pub max_cacheable_size: Option<usize>,
pub excluded_content_types: HashSet<String>,
pub force_cache_content_types: HashSet<String>,
pub stream_threshold: usize,
pub range_handling: RangeHandling,
pub enable_chunk_cache: bool,
pub chunk_size: usize,
pub min_chunk_file_size: u64,
}
impl Default for StreamingPolicy {
fn default() -> Self {
Self {
enabled: true,
max_cacheable_size: Some(1024 * 1024), excluded_content_types: HashSet::from([
"application/pdf".to_string(),
"video/*".to_string(),
"audio/*".to_string(),
"application/zip".to_string(),
"application/x-rar".to_string(),
"application/x-tar".to_string(),
"application/gzip".to_string(),
"application/x-7z-compressed".to_string(),
"application/octet-stream".to_string(),
]),
force_cache_content_types: HashSet::from([
"application/json".to_string(),
"application/xml".to_string(),
"text/*".to_string(),
]),
stream_threshold: 512 * 1024, range_handling: RangeHandling::default(),
enable_chunk_cache: false, chunk_size: 1024 * 1024, min_chunk_file_size: 10 * 1024 * 1024, }
}
}
impl StreamingPolicy {
pub fn disabled() -> Self {
Self {
enabled: false,
max_cacheable_size: None,
excluded_content_types: HashSet::new(),
force_cache_content_types: HashSet::new(),
stream_threshold: usize::MAX,
range_handling: RangeHandling::PassThrough,
enable_chunk_cache: false,
chunk_size: 1024 * 1024,
min_chunk_file_size: 0,
}
}
pub fn size_only(max_size: usize) -> Self {
Self {
enabled: true,
max_cacheable_size: Some(max_size),
excluded_content_types: HashSet::new(),
force_cache_content_types: HashSet::new(),
stream_threshold: max_size,
range_handling: RangeHandling::PassThrough,
enable_chunk_cache: false,
chunk_size: 1024 * 1024,
min_chunk_file_size: 0,
}
}
pub fn content_type_only(excluded: HashSet<String>) -> Self {
Self {
enabled: true,
max_cacheable_size: None,
excluded_content_types: excluded,
force_cache_content_types: HashSet::new(),
stream_threshold: usize::MAX,
range_handling: RangeHandling::PassThrough,
enable_chunk_cache: false,
chunk_size: 1024 * 1024,
min_chunk_file_size: 0,
}
}
}
pub fn should_stream(
policy: &StreamingPolicy,
size_hint: &http_body::SizeHint,
content_type: Option<&str>,
content_length: Option<u64>,
) -> StreamingDecision {
if !policy.enabled {
return StreamingDecision::Buffer;
}
let is_forced = if let Some(ct) = content_type {
if is_excluded_content_type(ct, &policy.excluded_content_types) {
return StreamingDecision::SkipCache;
}
is_forced_content_type(ct, &policy.force_cache_content_types)
} else {
false
};
if let Some(exact_size) = size_hint.exact() {
return decide_by_size(exact_size as usize, policy, is_forced);
}
if let Some(upper_bound) = size_hint.upper() {
return decide_by_size(upper_bound as usize, policy, is_forced);
}
if let Some(content_len) = content_length {
return decide_by_size(content_len as usize, policy, is_forced);
}
StreamingDecision::StreamIfPossible
}
fn decide_by_size(size: usize, policy: &StreamingPolicy, _is_forced: bool) -> StreamingDecision {
if let Some(max_size) = policy.max_cacheable_size {
if size > max_size {
return StreamingDecision::SkipCache;
}
}
StreamingDecision::Buffer
}
fn is_excluded_content_type(content_type: &str, excluded: &HashSet<String>) -> bool {
let normalized = content_type.to_lowercase();
for pattern in excluded {
if matches_pattern(&normalized, pattern) {
return true;
}
}
false
}
fn is_forced_content_type(content_type: &str, forced: &HashSet<String>) -> bool {
let normalized = content_type.to_lowercase();
for pattern in forced {
if matches_pattern(&normalized, pattern) {
return true;
}
}
false
}
fn matches_pattern(content_type: &str, pattern: &str) -> bool {
let pattern_lower = pattern.to_lowercase();
if pattern_lower.ends_with("/*") {
let prefix = &pattern_lower[..pattern_lower.len() - 2];
content_type.starts_with(prefix)
} else {
content_type == pattern_lower
|| content_type.starts_with(&format!("{};", pattern_lower))
}
}
pub fn extract_size_info(
size_hint: &http_body::SizeHint,
content_length: Option<u64>,
) -> Option<u64> {
size_hint
.exact()
.or_else(|| size_hint.upper())
.or(content_length)
}
#[cfg(test)]
mod tests {
use super::*;
use http_body::SizeHint;
#[test]
fn test_default_policy_excludes_pdf() {
let policy = StreamingPolicy::default();
let size_hint = SizeHint::with_exact(5 * 1024 * 1024);
let decision = should_stream(
&policy,
&size_hint,
Some("application/pdf"),
Some(5 * 1024 * 1024),
);
assert_eq!(decision, StreamingDecision::SkipCache);
}
#[test]
fn test_default_policy_excludes_video() {
let policy = StreamingPolicy::default();
let size_hint = SizeHint::with_exact(10 * 1024 * 1024);
let decision = should_stream(&policy, &size_hint, Some("video/mp4"), None);
assert_eq!(decision, StreamingDecision::SkipCache);
}
#[test]
fn test_small_json_gets_buffered() {
let policy = StreamingPolicy::default();
let size_hint = SizeHint::with_exact(1024);
let decision = should_stream(&policy, &size_hint, Some("application/json"), Some(1024));
assert_eq!(decision, StreamingDecision::Buffer);
}
#[test]
fn test_large_json_skipped_by_size() {
let policy = StreamingPolicy::default();
let size_hint = SizeHint::with_exact(2 * 1024 * 1024);
let decision = should_stream(
&policy,
&size_hint,
Some("application/json"),
Some(2 * 1024 * 1024),
);
assert_eq!(decision, StreamingDecision::SkipCache);
}
#[test]
fn test_force_cache_respects_size_limits() {
let mut policy = StreamingPolicy::default();
policy
.force_cache_content_types
.insert("application/important".to_string());
let size_hint = SizeHint::with_exact(5 * 1024 * 1024);
let decision = should_stream(
&policy,
&size_hint,
Some("application/important"),
Some(5 * 1024 * 1024),
);
assert_eq!(decision, StreamingDecision::SkipCache);
let small_hint = SizeHint::with_exact(500 * 1024); let decision_small = should_stream(
&policy,
&small_hint,
Some("application/important"),
Some(500 * 1024),
);
assert_eq!(decision_small, StreamingDecision::Buffer);
}
#[test]
fn test_disabled_policy_always_buffers() {
let policy = StreamingPolicy::disabled();
let size_hint = SizeHint::with_exact(10 * 1024 * 1024);
let decision = should_stream(
&policy,
&size_hint,
Some("application/pdf"),
Some(10 * 1024 * 1024),
);
assert_eq!(decision, StreamingDecision::Buffer);
}
#[test]
fn test_wildcard_pattern_matching() {
assert!(matches_pattern("video/mp4", "video/*"));
assert!(matches_pattern("video/mpeg", "video/*"));
assert!(matches_pattern("audio/mp3", "audio/*"));
assert!(!matches_pattern("application/json", "video/*"));
}
#[test]
fn test_exact_pattern_matching() {
assert!(matches_pattern("application/pdf", "application/pdf"));
assert!(!matches_pattern("application/pdf", "pdf")); assert!(!matches_pattern("text/plain", "application/pdf"));
assert!(matches_pattern(
"application/json; charset=utf-8",
"application/json"
));
}
#[test]
fn test_size_hint_exact() {
let policy = StreamingPolicy::default();
let size_hint = SizeHint::with_exact(500 * 1024);
let decision = should_stream(&policy, &size_hint, None, None);
assert_eq!(decision, StreamingDecision::Buffer);
}
#[test]
fn test_size_hint_upper_bound() {
let policy = StreamingPolicy::default();
let mut size_hint = SizeHint::default();
size_hint.set_upper(500 * 1024);
let decision = should_stream(&policy, &size_hint, None, None);
assert_eq!(decision, StreamingDecision::Buffer);
}
#[test]
fn test_content_length_fallback() {
let policy = StreamingPolicy::default();
let size_hint = SizeHint::default();
let decision = should_stream(&policy, &size_hint, None, Some(500 * 1024));
assert_eq!(decision, StreamingDecision::Buffer);
}
#[test]
fn test_unknown_size_conservative() {
let policy = StreamingPolicy::default();
let size_hint = SizeHint::default();
let decision = should_stream(&policy, &size_hint, None, None);
assert_eq!(decision, StreamingDecision::StreamIfPossible);
}
#[test]
fn test_size_only_policy() {
let policy = StreamingPolicy::size_only(512 * 1024);
let size_hint = SizeHint::with_exact(1024 * 1024);
let decision = should_stream(&policy, &size_hint, Some("application/pdf"), None);
assert_eq!(decision, StreamingDecision::SkipCache);
let size_hint_small = SizeHint::with_exact(256 * 1024);
let decision_small =
should_stream(&policy, &size_hint_small, Some("application/pdf"), None);
assert_eq!(decision_small, StreamingDecision::Buffer); }
#[test]
fn test_content_type_only_policy() {
let mut excluded = HashSet::new();
excluded.insert("application/pdf".to_string());
let policy = StreamingPolicy::content_type_only(excluded);
let size_hint = SizeHint::with_exact(10 * 1024 * 1024);
let decision = should_stream(&policy, &size_hint, Some("application/json"), None);
assert_eq!(decision, StreamingDecision::Buffer);
let decision_pdf = should_stream(&policy, &size_hint, Some("application/pdf"), None);
assert_eq!(decision_pdf, StreamingDecision::SkipCache);
}
#[test]
fn test_extract_size_info() {
let size_hint = SizeHint::with_exact(1024);
assert_eq!(extract_size_info(&size_hint, None), Some(1024));
let mut size_hint_upper = SizeHint::default();
size_hint_upper.set_upper(2048);
assert_eq!(extract_size_info(&size_hint_upper, None), Some(2048));
let size_hint_none = SizeHint::default();
assert_eq!(extract_size_info(&size_hint_none, Some(4096)), Some(4096));
assert_eq!(extract_size_info(&size_hint_none, None), None);
}
#[test]
fn test_case_insensitive_content_type() {
let policy = StreamingPolicy::default();
let size_hint = SizeHint::with_exact(1024);
assert_eq!(
should_stream(&policy, &size_hint, Some("Application/PDF"), None),
StreamingDecision::SkipCache
);
assert_eq!(
should_stream(&policy, &size_hint, Some("APPLICATION/PDF"), None),
StreamingDecision::SkipCache
);
assert_eq!(
should_stream(&policy, &size_hint, Some("Video/MP4"), None),
StreamingDecision::SkipCache
);
}
#[test]
fn test_content_type_with_charset() {
let policy = StreamingPolicy::default();
let size_hint = SizeHint::with_exact(1024);
let decision = should_stream(
&policy,
&size_hint,
Some("application/json; charset=utf-8"),
None,
);
assert_eq!(decision, StreamingDecision::Buffer);
}
}