use bytes::Bytes;
use reqwest::Response;
use std::path::PathBuf;
use tokio::io::AsyncRead;
use tracing::debug;
use ngdp_cdn::CdnClient;
use crate::{CdnCache, Result};
#[derive(Debug, Clone, Copy, PartialEq)]
enum ContentType {
Config,
Data,
Patch,
}
impl ContentType {
fn from_path(path: &str) -> Self {
let path_lower = path.to_lowercase();
if path_lower.contains("/config") || path_lower.ends_with("config") {
Self::Config
} else if path_lower.contains("/patch") || path_lower.ends_with("patch") {
Self::Patch
} else {
Self::Data
}
}
}
pub struct CachedCdnClient {
client: CdnClient,
cache_base_dir: PathBuf,
enabled: bool,
}
impl CachedCdnClient {
pub async fn new() -> Result<Self> {
let client = CdnClient::new()?;
let cache_base_dir = crate::get_cache_dir()?.join("cdn");
crate::ensure_dir(&cache_base_dir).await?;
debug!("Initialized cached CDN client");
Ok(Self {
client,
cache_base_dir,
enabled: true,
})
}
pub async fn for_product(product: &str) -> Result<Self> {
let client = CdnClient::new()?;
let cache_base_dir = crate::get_cache_dir()?.join("cdn").join(product);
crate::ensure_dir(&cache_base_dir).await?;
debug!("Initialized cached CDN client for product '{}'", product);
Ok(Self {
client,
cache_base_dir,
enabled: true,
})
}
pub async fn with_cache_dir(cache_dir: PathBuf) -> Result<Self> {
let client = CdnClient::new()?;
crate::ensure_dir(&cache_dir).await?;
Ok(Self {
client,
cache_base_dir: cache_dir,
enabled: true,
})
}
pub async fn with_client(client: CdnClient) -> Result<Self> {
let cache_base_dir = crate::get_cache_dir()?.join("cdn");
crate::ensure_dir(&cache_base_dir).await?;
Ok(Self {
client,
cache_base_dir,
enabled: true,
})
}
pub fn add_primary_host(&self, host: impl Into<String>) {
self.client.add_primary_host(host);
}
pub fn add_primary_hosts(&self, hosts: impl IntoIterator<Item = impl Into<String>>) {
self.client.add_primary_hosts(hosts);
}
pub fn add_fallback_host(&self, host: impl Into<String>) {
self.client.add_fallback_host(host);
}
pub fn add_fallback_hosts(&self, hosts: impl IntoIterator<Item = impl Into<String>>) {
self.client.add_fallback_hosts(hosts);
}
pub fn set_primary_hosts(&self, hosts: impl IntoIterator<Item = impl Into<String>>) {
self.client.set_primary_hosts(hosts);
}
pub fn get_all_hosts(&self) -> Vec<String> {
self.client.get_all_hosts()
}
pub fn set_caching_enabled(&mut self, enabled: bool) {
self.enabled = enabled;
}
pub fn cache_dir(&self) -> &PathBuf {
&self.cache_base_dir
}
async fn get_cache_for_path(&self, cdn_path: &str) -> Result<CdnCache> {
let mut cache = CdnCache::with_base_dir(self.cache_base_dir.clone()).await?;
cache.set_cdn_path(Some(cdn_path.to_string()));
Ok(cache)
}
async fn is_cached(&self, path: &str, hash: &str) -> Result<bool> {
let cache = self.get_cache_for_path(path).await?;
let content_type = ContentType::from_path(path);
Ok(match content_type {
ContentType::Config => cache.has_config(hash).await,
ContentType::Data => cache.has_data(hash).await,
ContentType::Patch => cache.has_patch(hash).await,
})
}
async fn read_from_cache(&self, path: &str, hash: &str) -> Result<Bytes> {
let cache = self.get_cache_for_path(path).await?;
let content_type = ContentType::from_path(path);
let data = match content_type {
ContentType::Config => cache.read_config(hash).await?,
ContentType::Data => cache.read_data(hash).await?,
ContentType::Patch => cache.read_patch(hash).await?,
};
Ok(Bytes::from(data))
}
async fn write_to_cache(&self, path: &str, hash: &str, data: &[u8]) -> Result<()> {
let cache = self.get_cache_for_path(path).await?;
let content_type = ContentType::from_path(path);
match content_type {
ContentType::Config => cache.write_config(hash, data).await?,
ContentType::Data => cache.write_data(hash, data).await?,
ContentType::Patch => cache.write_patch(hash, data).await?,
};
Ok(())
}
pub async fn request(&self, url: &str) -> Result<Response> {
Ok(self.client.request(url).await?)
}
pub async fn download(&self, cdn_host: &str, path: &str, hash: &str) -> Result<CachedResponse> {
if self.enabled && self.is_cached(path, hash).await? {
debug!("Cache hit for CDN {}/{}", path, hash);
let data = self.read_from_cache(path, hash).await?;
return Ok(CachedResponse::from_cache(data));
}
debug!("Cache miss for CDN {}/{}, fetching from server", path, hash);
let response = self.client.download(cdn_host, path, hash).await?;
let data = response.bytes().await?;
if self.enabled {
if let Err(e) = self.write_to_cache(path, hash, &data).await {
debug!("Failed to write to CDN cache: {}", e);
}
}
Ok(CachedResponse::from_network(data))
}
pub async fn download_stream(
&self,
cdn_host: &str,
path: &str,
hash: &str,
) -> Result<Box<dyn AsyncRead + Unpin + Send>> {
if ContentType::from_path(path) == ContentType::Data {
let cache = self.get_cache_for_path(path).await?;
if self.enabled && cache.has_data(hash).await {
debug!("Cache hit for CDN {}/{} (streaming)", path, hash);
let file = cache.open_data(hash).await?;
return Ok(Box::new(file));
}
}
let response = self.download(cdn_host, path, hash).await?;
let data = response.bytes().await?;
Ok(Box::new(std::io::Cursor::new(data.to_vec())))
}
pub async fn cached_size(&self, path: &str, hash: &str) -> Result<Option<u64>> {
if !self.enabled || !self.is_cached(path, hash).await? {
return Ok(None);
}
if ContentType::from_path(path) == ContentType::Data {
let cache = self.get_cache_for_path(path).await?;
Ok(Some(cache.data_size(hash).await?))
} else {
let data = self.read_from_cache(path, hash).await?;
Ok(Some(data.len() as u64))
}
}
pub async fn clear_cache(&self) -> Result<()> {
if tokio::fs::metadata(&self.cache_base_dir).await.is_ok() {
tokio::fs::remove_dir_all(&self.cache_base_dir).await?;
}
Ok(())
}
pub async fn cache_stats(&self) -> Result<CacheStats> {
let mut stats = CacheStats::default();
for entry in walkdir::WalkDir::new(&self.cache_base_dir)
.into_iter()
.flatten()
{
if entry.file_type().is_file() {
if let Ok(metadata) = entry.metadata() {
stats.total_files += 1;
stats.total_size += metadata.len();
let path = entry.path();
if path.to_string_lossy().contains("config") {
stats.config_files += 1;
stats.config_size += metadata.len();
} else if path.to_string_lossy().contains("patch") {
stats.patch_files += 1;
stats.patch_size += metadata.len();
} else if path.to_string_lossy().contains("data") {
stats.data_files += 1;
stats.data_size += metadata.len();
}
}
}
}
Ok(stats)
}
pub async fn download_build_config(
&self,
cdn_host: &str,
path: &str,
hash: &str,
) -> Result<CachedResponse> {
let config_path = format!("{}/config", path.trim_end_matches('/'));
self.download(cdn_host, &config_path, hash).await
}
pub async fn download_cdn_config(
&self,
cdn_host: &str,
path: &str,
hash: &str,
) -> Result<CachedResponse> {
let config_path = format!("{}/config", path.trim_end_matches('/'));
self.download(cdn_host, &config_path, hash).await
}
pub async fn download_product_config(
&self,
cdn_host: &str,
config_path: &str,
hash: &str,
) -> Result<CachedResponse> {
self.download(cdn_host, config_path, hash).await
}
pub async fn download_key_ring(
&self,
cdn_host: &str,
path: &str,
hash: &str,
) -> Result<CachedResponse> {
let config_path = format!("{}/config", path.trim_end_matches('/'));
self.download(cdn_host, &config_path, hash).await
}
pub async fn download_data(
&self,
cdn_host: &str,
path: &str,
hash: &str,
) -> Result<CachedResponse> {
let data_path = format!("{}/data", path.trim_end_matches('/'));
self.download(cdn_host, &data_path, hash).await
}
pub async fn download_patch(
&self,
cdn_host: &str,
path: &str,
hash: &str,
) -> Result<CachedResponse> {
let patch_path = format!("{}/patch", path.trim_end_matches('/'));
self.download(cdn_host, &patch_path, hash).await
}
}
pub struct CachedResponse {
data: Bytes,
from_cache: bool,
}
impl CachedResponse {
fn from_cache(data: Bytes) -> Self {
Self {
data,
from_cache: true,
}
}
fn from_network(data: Bytes) -> Self {
Self {
data,
from_cache: false,
}
}
pub fn is_from_cache(&self) -> bool {
self.from_cache
}
pub async fn bytes(self) -> Result<Bytes> {
Ok(self.data)
}
pub async fn text(self) -> Result<String> {
Ok(String::from_utf8(self.data.to_vec())?)
}
pub fn content_length(&self) -> usize {
self.data.len()
}
}
#[derive(Debug, Default, Clone)]
pub struct CacheStats {
pub total_files: u64,
pub total_size: u64,
pub config_files: u64,
pub config_size: u64,
pub data_files: u64,
pub data_size: u64,
pub patch_files: u64,
pub patch_size: u64,
}
impl CacheStats {
pub fn total_size_human(&self) -> String {
format_bytes(self.total_size)
}
pub fn config_size_human(&self) -> String {
format_bytes(self.config_size)
}
pub fn data_size_human(&self) -> String {
format_bytes(self.data_size)
}
pub fn patch_size_human(&self) -> String {
format_bytes(self.patch_size)
}
}
fn format_bytes(bytes: u64) -> String {
const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
let mut size = bytes as f64;
let mut unit_idx = 0;
while size >= 1024.0 && unit_idx < UNITS.len() - 1 {
size /= 1024.0;
unit_idx += 1;
}
if unit_idx == 0 {
format!("{} {}", size as u64, UNITS[unit_idx])
} else {
format!("{:.2} {}", size, UNITS[unit_idx])
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[tokio::test]
async fn test_cached_cdn_client_creation() {
let client = CachedCdnClient::new().await.unwrap();
assert!(client.enabled);
}
#[tokio::test]
async fn test_content_type_detection() {
assert_eq!(
ContentType::from_path("tpr/configs/data/config"),
ContentType::Config
);
assert_eq!(
ContentType::from_path("tpr/wow/config"),
ContentType::Config
);
assert_eq!(ContentType::from_path("config"), ContentType::Config);
assert_eq!(ContentType::from_path("tpr/wow/data"), ContentType::Data);
assert_eq!(ContentType::from_path("tpr/wow/patch"), ContentType::Patch);
assert_eq!(ContentType::from_path("tpr/wow"), ContentType::Data);
}
#[tokio::test]
async fn test_cache_enabling() {
let mut client = CachedCdnClient::new().await.unwrap();
client.set_caching_enabled(false);
assert!(!client.enabled);
client.set_caching_enabled(true);
assert!(client.enabled);
}
#[tokio::test]
async fn test_format_bytes() {
assert_eq!(format_bytes(0), "0 B");
assert_eq!(format_bytes(1023), "1023 B");
assert_eq!(format_bytes(1024), "1.00 KB");
assert_eq!(format_bytes(1536), "1.50 KB");
assert_eq!(format_bytes(1048576), "1.00 MB");
assert_eq!(format_bytes(1073741824), "1.00 GB");
}
#[tokio::test]
async fn test_cache_with_temp_dir() {
let temp_dir = TempDir::new().unwrap();
let client = CachedCdnClient::with_cache_dir(temp_dir.path().to_path_buf())
.await
.unwrap();
assert_eq!(client.cache_dir(), &temp_dir.path().to_path_buf());
}
}