1use bytes::Bytes;
49use reqwest::Response;
50use std::path::PathBuf;
51use tokio::io::AsyncRead;
52use tracing::debug;
53
54use ngdp_cdn::CdnClient;
55
56use crate::{CdnCache, Result};
57
58#[derive(Debug, Clone, Copy, PartialEq)]
60enum ContentType {
61 Config,
62 Data,
63 Patch,
64}
65
66impl ContentType {
67 fn from_path(path: &str) -> Self {
69 let path_lower = path.to_lowercase();
70 if path_lower.contains("/config") || path_lower.ends_with("config") {
71 Self::Config
72 } else if path_lower.contains("/patch") || path_lower.ends_with("patch") {
73 Self::Patch
74 } else {
75 Self::Data
76 }
77 }
78}
79
80pub struct CachedCdnClient {
82 client: CdnClient,
84 cache_base_dir: PathBuf,
86 enabled: bool,
88}
89
90impl CachedCdnClient {
91 pub async fn new() -> Result<Self> {
93 let client = CdnClient::new()?;
94 let cache_base_dir = crate::get_cache_dir()?.join("cdn");
95 crate::ensure_dir(&cache_base_dir).await?;
96
97 debug!("Initialized cached CDN client");
98
99 Ok(Self {
100 client,
101 cache_base_dir,
102 enabled: true,
103 })
104 }
105
106 pub async fn for_product(product: &str) -> Result<Self> {
108 let client = CdnClient::new()?;
109 let cache_base_dir = crate::get_cache_dir()?.join("cdn").join(product);
110 crate::ensure_dir(&cache_base_dir).await?;
111
112 debug!("Initialized cached CDN client for product '{}'", product);
113
114 Ok(Self {
115 client,
116 cache_base_dir,
117 enabled: true,
118 })
119 }
120
121 pub async fn with_cache_dir(cache_dir: PathBuf) -> Result<Self> {
123 let client = CdnClient::new()?;
124 crate::ensure_dir(&cache_dir).await?;
125
126 Ok(Self {
127 client,
128 cache_base_dir: cache_dir,
129 enabled: true,
130 })
131 }
132
133 pub async fn with_client(client: CdnClient) -> Result<Self> {
135 let cache_base_dir = crate::get_cache_dir()?.join("cdn");
136 crate::ensure_dir(&cache_base_dir).await?;
137
138 Ok(Self {
139 client,
140 cache_base_dir,
141 enabled: true,
142 })
143 }
144
145 pub fn set_caching_enabled(&mut self, enabled: bool) {
147 self.enabled = enabled;
148 }
149
150 pub fn cache_dir(&self) -> &PathBuf {
152 &self.cache_base_dir
153 }
154
155 async fn get_cache_for_path(&self, cdn_path: &str) -> Result<CdnCache> {
157 let mut cache = CdnCache::with_base_dir(self.cache_base_dir.clone()).await?;
159 cache.set_cdn_path(Some(cdn_path.to_string()));
160 Ok(cache)
161 }
162
163 async fn is_cached(&self, path: &str, hash: &str) -> Result<bool> {
165 let cache = self.get_cache_for_path(path).await?;
166 let content_type = ContentType::from_path(path);
167 Ok(match content_type {
168 ContentType::Config => cache.has_config(hash).await,
169 ContentType::Data => cache.has_data(hash).await,
170 ContentType::Patch => cache.has_patch(hash).await,
171 })
172 }
173
174 async fn read_from_cache(&self, path: &str, hash: &str) -> Result<Bytes> {
176 let cache = self.get_cache_for_path(path).await?;
177 let content_type = ContentType::from_path(path);
178 let data = match content_type {
179 ContentType::Config => cache.read_config(hash).await?,
180 ContentType::Data => cache.read_data(hash).await?,
181 ContentType::Patch => cache.read_patch(hash).await?,
182 };
183 Ok(Bytes::from(data))
184 }
185
186 async fn write_to_cache(&self, path: &str, hash: &str, data: &[u8]) -> Result<()> {
188 let cache = self.get_cache_for_path(path).await?;
189 let content_type = ContentType::from_path(path);
190 match content_type {
191 ContentType::Config => cache.write_config(hash, data).await?,
192 ContentType::Data => cache.write_data(hash, data).await?,
193 ContentType::Patch => cache.write_patch(hash, data).await?,
194 };
195 Ok(())
196 }
197
198 pub async fn request(&self, url: &str) -> Result<Response> {
203 Ok(self.client.request(url).await?)
204 }
205
206 pub async fn download(&self, cdn_host: &str, path: &str, hash: &str) -> Result<CachedResponse> {
212 if self.enabled && self.is_cached(path, hash).await? {
214 debug!("Cache hit for CDN {}/{}", path, hash);
215 let data = self.read_from_cache(path, hash).await?;
216 return Ok(CachedResponse::from_cache(data));
217 }
218
219 debug!("Cache miss for CDN {}/{}, fetching from server", path, hash);
221 let response = self.client.download(cdn_host, path, hash).await?;
222
223 let data = response.bytes().await?;
225
226 if self.enabled {
228 if let Err(e) = self.write_to_cache(path, hash, &data).await {
229 debug!("Failed to write to CDN cache: {}", e);
230 }
231 }
232
233 Ok(CachedResponse::from_network(data))
234 }
235
236 pub async fn download_stream(
242 &self,
243 cdn_host: &str,
244 path: &str,
245 hash: &str,
246 ) -> Result<Box<dyn AsyncRead + Unpin + Send>> {
247 if ContentType::from_path(path) == ContentType::Data {
249 let cache = self.get_cache_for_path(path).await?;
251 if self.enabled && cache.has_data(hash).await {
252 debug!("Cache hit for CDN {}/{} (streaming)", path, hash);
253 let file = cache.open_data(hash).await?;
254 return Ok(Box::new(file));
255 }
256 }
257
258 let response = self.download(cdn_host, path, hash).await?;
260 let data = response.bytes().await?;
261
262 Ok(Box::new(std::io::Cursor::new(data.to_vec())))
264 }
265
266 pub async fn cached_size(&self, path: &str, hash: &str) -> Result<Option<u64>> {
268 if !self.enabled || !self.is_cached(path, hash).await? {
269 return Ok(None);
270 }
271
272 if ContentType::from_path(path) == ContentType::Data {
274 let cache = self.get_cache_for_path(path).await?;
275 Ok(Some(cache.data_size(hash).await?))
276 } else {
277 let data = self.read_from_cache(path, hash).await?;
279 Ok(Some(data.len() as u64))
280 }
281 }
282
283 pub async fn clear_cache(&self) -> Result<()> {
288 if tokio::fs::metadata(&self.cache_base_dir).await.is_ok() {
289 tokio::fs::remove_dir_all(&self.cache_base_dir).await?;
290 }
291 Ok(())
292 }
293
294 pub async fn cache_stats(&self) -> Result<CacheStats> {
296 let mut stats = CacheStats::default();
297
298 for entry in walkdir::WalkDir::new(&self.cache_base_dir)
300 .into_iter()
301 .flatten()
302 {
303 if entry.file_type().is_file() {
304 if let Ok(metadata) = entry.metadata() {
305 stats.total_files += 1;
306 stats.total_size += metadata.len();
307
308 let path = entry.path();
309 if path.to_string_lossy().contains("config") {
310 stats.config_files += 1;
311 stats.config_size += metadata.len();
312 } else if path.to_string_lossy().contains("patch") {
313 stats.patch_files += 1;
314 stats.patch_size += metadata.len();
315 } else if path.to_string_lossy().contains("data") {
316 stats.data_files += 1;
317 stats.data_size += metadata.len();
318 }
319 }
320 }
321 }
322
323 Ok(stats)
324 }
325
326 pub async fn download_build_config(
330 &self,
331 cdn_host: &str,
332 path: &str,
333 hash: &str,
334 ) -> Result<CachedResponse> {
335 let config_path = format!("{}/config", path.trim_end_matches('/'));
336 self.download(cdn_host, &config_path, hash).await
337 }
338
339 pub async fn download_cdn_config(
343 &self,
344 cdn_host: &str,
345 path: &str,
346 hash: &str,
347 ) -> Result<CachedResponse> {
348 let config_path = format!("{}/config", path.trim_end_matches('/'));
349 self.download(cdn_host, &config_path, hash).await
350 }
351
352 pub async fn download_product_config(
357 &self,
358 cdn_host: &str,
359 config_path: &str,
360 hash: &str,
361 ) -> Result<CachedResponse> {
362 self.download(cdn_host, config_path, hash).await
363 }
364
365 pub async fn download_key_ring(
369 &self,
370 cdn_host: &str,
371 path: &str,
372 hash: &str,
373 ) -> Result<CachedResponse> {
374 let config_path = format!("{}/config", path.trim_end_matches('/'));
375 self.download(cdn_host, &config_path, hash).await
376 }
377
378 pub async fn download_data(
382 &self,
383 cdn_host: &str,
384 path: &str,
385 hash: &str,
386 ) -> Result<CachedResponse> {
387 let data_path = format!("{}/data", path.trim_end_matches('/'));
388 self.download(cdn_host, &data_path, hash).await
389 }
390
391 pub async fn download_patch(
395 &self,
396 cdn_host: &str,
397 path: &str,
398 hash: &str,
399 ) -> Result<CachedResponse> {
400 let patch_path = format!("{}/patch", path.trim_end_matches('/'));
401 self.download(cdn_host, &patch_path, hash).await
402 }
403}
404
405pub struct CachedResponse {
407 data: Bytes,
409 from_cache: bool,
411}
412
413impl CachedResponse {
414 fn from_cache(data: Bytes) -> Self {
416 Self {
417 data,
418 from_cache: true,
419 }
420 }
421
422 fn from_network(data: Bytes) -> Self {
424 Self {
425 data,
426 from_cache: false,
427 }
428 }
429
430 pub fn is_from_cache(&self) -> bool {
432 self.from_cache
433 }
434
435 pub async fn bytes(self) -> Result<Bytes> {
437 Ok(self.data)
438 }
439
440 pub async fn text(self) -> Result<String> {
442 Ok(String::from_utf8(self.data.to_vec())?)
443 }
444
445 pub fn content_length(&self) -> usize {
447 self.data.len()
448 }
449}
450
451#[derive(Debug, Default, Clone)]
453pub struct CacheStats {
454 pub total_files: u64,
456 pub total_size: u64,
458 pub config_files: u64,
460 pub config_size: u64,
462 pub data_files: u64,
464 pub data_size: u64,
466 pub patch_files: u64,
468 pub patch_size: u64,
470}
471
472impl CacheStats {
473 pub fn total_size_human(&self) -> String {
475 format_bytes(self.total_size)
476 }
477
478 pub fn config_size_human(&self) -> String {
480 format_bytes(self.config_size)
481 }
482
483 pub fn data_size_human(&self) -> String {
485 format_bytes(self.data_size)
486 }
487
488 pub fn patch_size_human(&self) -> String {
490 format_bytes(self.patch_size)
491 }
492}
493
494fn format_bytes(bytes: u64) -> String {
496 const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
497 let mut size = bytes as f64;
498 let mut unit_idx = 0;
499
500 while size >= 1024.0 && unit_idx < UNITS.len() - 1 {
501 size /= 1024.0;
502 unit_idx += 1;
503 }
504
505 if unit_idx == 0 {
506 format!("{} {}", size as u64, UNITS[unit_idx])
507 } else {
508 format!("{:.2} {}", size, UNITS[unit_idx])
509 }
510}
511
512#[cfg(test)]
513mod tests {
514 use super::*;
515 use tempfile::TempDir;
516
517 #[tokio::test]
518 async fn test_cached_cdn_client_creation() {
519 let client = CachedCdnClient::new().await.unwrap();
520 assert!(client.enabled);
521 }
522
523 #[tokio::test]
524 async fn test_content_type_detection() {
525 assert_eq!(
526 ContentType::from_path("tpr/configs/data/config"),
527 ContentType::Config
528 );
529 assert_eq!(
530 ContentType::from_path("tpr/wow/config"),
531 ContentType::Config
532 );
533 assert_eq!(ContentType::from_path("config"), ContentType::Config);
534 assert_eq!(ContentType::from_path("tpr/wow/data"), ContentType::Data);
535 assert_eq!(ContentType::from_path("tpr/wow/patch"), ContentType::Patch);
536 assert_eq!(ContentType::from_path("tpr/wow"), ContentType::Data);
537 }
538
539 #[tokio::test]
540 async fn test_cache_enabling() {
541 let mut client = CachedCdnClient::new().await.unwrap();
542
543 client.set_caching_enabled(false);
544 assert!(!client.enabled);
545
546 client.set_caching_enabled(true);
547 assert!(client.enabled);
548 }
549
550 #[tokio::test]
551 async fn test_format_bytes() {
552 assert_eq!(format_bytes(0), "0 B");
553 assert_eq!(format_bytes(1023), "1023 B");
554 assert_eq!(format_bytes(1024), "1.00 KB");
555 assert_eq!(format_bytes(1536), "1.50 KB");
556 assert_eq!(format_bytes(1048576), "1.00 MB");
557 assert_eq!(format_bytes(1073741824), "1.00 GB");
558 }
559
560 #[tokio::test]
561 async fn test_cache_with_temp_dir() {
562 let temp_dir = TempDir::new().unwrap();
563 let client = CachedCdnClient::with_cache_dir(temp_dir.path().to_path_buf())
564 .await
565 .unwrap();
566
567 assert_eq!(client.cache_dir(), &temp_dir.path().to_path_buf());
568 }
569}