kget/
optimization.rs

1use std::error::Error;
2use std::fs::{self, File};
3use std::io::{Read, Write};
4use std::path::PathBuf;
5use flate2::write::{GzEncoder, GzDecoder};
6use lz4::block::{compress, CompressionMode};
7use crate::config::OptimizationConfig;
8
9/// Structure responsible for optimizing download operations through compression and caching
10#[derive(Clone)]
11pub struct Optimizer {
12    config: OptimizationConfig,
13    pub speed_limit: Option<u64>,
14}
15
16impl Optimizer {
17    /// Make a new Optimizer instance with the provided configuration
18    pub fn new(config: OptimizationConfig) -> Self {
19        let speed_limit = config.speed_limit;
20        Self { 
21            config,
22            speed_limit,
23        }
24    }
25
26    /// Compress data using different algorithms based on the configured compression level
27    ///
28    /// Levels 1-3: Gzip
29    /// Levels 4-6: LZ4
30    /// Levels 7-9: Brotli
31    #[allow(dead_code)]
32    pub fn compress(&self, data: &[u8]) -> Result<Vec<u8>, Box<dyn Error>> {
33        if !self.config.compression {
34            return Ok(data.to_vec());
35        }
36        let compressed = match self.config.compression_level {
37            1..=3 => {
38                let mut encoder = GzEncoder::new(Vec::new(), flate2::Compression::fast());
39                encoder.write_all(data)?;
40                encoder.finish()?
41            }
42            4..=6 => {
43                compress(data, Some(CompressionMode::FAST(0)), true)?
44            }
45            7..=9 => {
46                let mut encoder = brotli::CompressorWriter::new(
47                    Vec::new(),
48                    self.config.compression_level as usize,
49                    4096,
50                    22,
51                );
52                encoder.write_all(data)?;
53                encoder.into_inner()
54            }
55            _ => return Ok(data.to_vec()),
56        };
57        Ok(compressed)
58    }
59
60    /// Decompress data using the appropriate algorithm based on the file header
61    ///
62    /// Supports Gzip, Brotli, and LZ4
63    pub fn decompress(&self, data: &[u8]) -> Result<Vec<u8>, Box<dyn Error>> {
64        if !self.config.compression {
65            return Ok(data.to_vec());
66        }
67        let mut decompressed = Vec::new();
68        if data.starts_with(&[0x1f, 0x8b]) {
69            let mut decoder = GzDecoder::new(Vec::new());
70            decoder.write_all(data)?;
71            decompressed = decoder.finish()?;
72        } else if data.starts_with(&[0x28, 0xb5, 0x2f, 0xfd]) {
73            let mut decoder = brotli::Decompressor::new(data, 4096);
74            decoder.read_to_end(&mut decompressed)?;
75        } else {
76            let mut decoder = lz4::Decoder::new(data)?;
77            decoder.read_to_end(&mut decompressed)?;
78        }
79        Ok(decompressed)
80    }
81
82    /// Retrieve a file from the cache if it exists
83    ///
84    /// Returns None if caching is disabled or the file does not exist
85    #[allow(dead_code)]
86    pub fn get_cached_file(&self, url: &str) -> Result<Option<Vec<u8>>, Box<dyn Error>> {
87        if !self.config.cache_enabled {
88            return Ok(None);
89        }
90
91        let _cache_dir = self.config.cache_dir.as_str();
92        let cache_path = self.get_cache_path(url)?;
93        if cache_path.exists() {
94            let mut file = File::open(cache_path)?;
95            let mut contents = Vec::new();
96            file.read_to_end(&mut contents)?;
97            return Ok(Some(contents));
98        }
99        Ok(None)
100    }
101
102    /// Store a file in the cache
103    ///
104    /// Does nothing if caching is disabled
105    #[allow(dead_code)]
106    pub fn cache_file(&self, url: &str, data: &[u8]) -> Result<(), Box<dyn Error>> {
107        if !self.config.cache_enabled {
108            return Ok(());
109        }
110        let cache_path = self.get_cache_path(url)?;
111        if let Some(parent) = cache_path.parent() {
112            fs::create_dir_all(parent)?;
113        }
114        let mut file = File::create(cache_path)?;
115        file.write_all(data)?;
116        Ok(())
117    }
118
119    /// Generate the cache file path based on the URL
120    ///
121    /// Uses a simple hash to generate a unique filename
122    #[allow(dead_code)]
123    fn get_cache_path(&self, url: &str) -> Result<PathBuf, Box<dyn Error>> {
124        let mut cache_dir = PathBuf::from(
125            if self.config.cache_dir.is_empty() {
126                "~/.cache/kget".to_string()
127            } else {
128                self.config.cache_dir.clone()
129            }
130        );
131        
132        if cache_dir.starts_with("~") {
133            if let Some(home) = dirs::home_dir() {
134                cache_dir = home.join(cache_dir.strip_prefix("~").unwrap());
135            }
136        }
137        
138        // Simple hash function to generate a unique filename
139        let mut hash = 0u64;
140        for byte in url.bytes() {
141            hash = hash.wrapping_mul(31).wrapping_add(byte as u64);
142        }
143        
144        cache_dir.push(format!("{:x}", hash));
145        Ok(cache_dir)
146    }
147
148    pub fn get_peer_limit(&self) -> usize {
149        self.speed_limit.unwrap_or(50) as usize
150    }
151    
152    pub fn is_compression_enabled(&self) -> bool {
153        self.config.compression
154    }
155}
kget/optimization.rs

kget/
optimization.rs
