1use crate::error::{Error, Result};
2use crate::runtime::{Language, Runtime};
3use sha2::{Digest, Sha256};
4use std::fs;
5use std::io::Read;
6use std::path::PathBuf;
7
8pub struct CacheManager {
9 cache_dir: PathBuf,
10}
11
12impl CacheManager {
13 pub fn new() -> Result<Self> {
14 let cache_dir = Self::default_cache_dir()?;
15 Ok(Self { cache_dir })
16 }
17
18 pub fn with_cache_dir(cache_dir: PathBuf) -> Self {
19 Self { cache_dir }
20 }
21
22 pub fn default_cache_dir() -> Result<PathBuf> {
23 let cache_dir = dirs::cache_dir()
24 .ok_or_else(|| Error::Other("Could not determine cache directory".to_string()))?
25 .join("wasm-runtime");
26 Ok(cache_dir)
27 }
28
29 pub fn get_path(&self, language: Language, version: &str) -> PathBuf {
30 self.cache_dir
31 .join(language.as_str())
32 .join(format!("{version}.wasm"))
33 }
34
35 pub fn get(&self, language: Language, version: &str) -> Option<Runtime> {
36 let path = self.get_path(language, version);
37 if !path.exists() {
38 return None;
39 }
40
41 let metadata = fs::metadata(&path).ok()?;
42 let size = metadata.len();
43
44 let sha256 = Self::compute_sha256(&path).ok()?;
45
46 Some(Runtime::new(
47 language,
48 version.to_string(),
49 path,
50 size,
51 sha256,
52 ))
53 }
54
55 pub fn store(&self, language: Language, version: &str, data: &[u8]) -> Result<Runtime> {
56 let path = self.get_path(language, version);
57
58 if let Some(parent) = path.parent() {
59 fs::create_dir_all(parent)?;
60 }
61
62 fs::write(&path, data)?;
63
64 let size = data.len() as u64;
65 let sha256 = Self::compute_sha256(&path)?;
66
67 Ok(Runtime::new(
68 language,
69 version.to_string(),
70 path,
71 size,
72 sha256,
73 ))
74 }
75
76 pub fn clear(&self, language: Language, version: &str) -> Result<()> {
77 let path = self.get_path(language, version);
78 if path.exists() {
79 fs::remove_file(path)?;
80 }
81 Ok(())
82 }
83
84 pub fn clear_all(&self) -> Result<()> {
85 if self.cache_dir.exists() {
86 fs::remove_dir_all(&self.cache_dir)?;
87 }
88 Ok(())
89 }
90
91 pub fn list(&self) -> Result<Vec<Runtime>> {
92 let mut runtimes = Vec::new();
93
94 if !self.cache_dir.exists() {
95 return Ok(runtimes);
96 }
97
98 for language in Language::all() {
99 let lang_dir = self.cache_dir.join(language.as_str());
100 if !lang_dir.exists() {
101 continue;
102 }
103
104 for entry in fs::read_dir(&lang_dir)? {
105 let entry = entry?;
106 let path = entry.path();
107
108 if !path.is_file() || path.extension().and_then(|s| s.to_str()) != Some("wasm") {
109 continue;
110 }
111
112 if let Some(version) = path
113 .file_stem()
114 .and_then(|s| s.to_str())
115 .map(|s| s.to_string())
116 {
117 if let Some(runtime) = self.get(*language, &version) {
118 runtimes.push(runtime);
119 }
120 }
121 }
122 }
123
124 Ok(runtimes)
125 }
126
127 pub fn compute_sha256(path: &PathBuf) -> Result<String> {
128 let mut file = fs::File::open(path)?;
129 let mut hasher = Sha256::new();
130 let mut buffer = [0; 8192];
131
132 loop {
133 let n = file.read(&mut buffer)?;
134 if n == 0 {
135 break;
136 }
137 hasher.update(&buffer[..n]);
138 }
139
140 let hash = hasher.finalize();
141 Ok(format!("{hash:x}"))
142 }
143
144 pub fn verify_integrity(&self, runtime: &Runtime, expected_sha256: &str) -> Result<()> {
145 let actual_sha256 = Self::compute_sha256(&runtime.path)?;
146 if actual_sha256 != expected_sha256 {
147 return Err(Error::IntegrityCheckFailed {
148 expected: expected_sha256.to_string(),
149 actual: actual_sha256,
150 });
151 }
152 Ok(())
153 }
154}
155
156impl Default for CacheManager {
157 fn default() -> Self {
158 Self::new().expect("Failed to create cache manager")
159 }
160}
161
162#[cfg(test)]
163mod tests {
164 use super::*;
165 use std::fs;
166 use tempfile::TempDir;
167
168 fn create_test_cache() -> (CacheManager, TempDir) {
169 let temp_dir = TempDir::new().unwrap();
170 let cache_manager = CacheManager::with_cache_dir(temp_dir.path().to_path_buf());
171 (cache_manager, temp_dir)
172 }
173
174 #[test]
175 fn test_get_path() {
176 let (cache, _temp) = create_test_cache();
177 let path = cache.get_path(Language::Python, "3.11.7");
178 assert!(path.to_string_lossy().contains("python"));
179 assert!(path.to_string_lossy().contains("3.11.7.wasm"));
180 }
181
182 #[test]
183 fn test_store_and_get() {
184 let (cache, _temp) = create_test_cache();
185 let data = b"test wasm data";
186 let runtime = cache
187 .store(Language::Python, "3.11.7", data)
188 .expect("Failed to store");
189
190 assert_eq!(runtime.language, Language::Python);
191 assert_eq!(runtime.version, "3.11.7");
192 assert_eq!(runtime.size, data.len() as u64);
193
194 let retrieved = cache.get(Language::Python, "3.11.7");
195 assert!(retrieved.is_some());
196 let retrieved = retrieved.unwrap();
197 assert_eq!(retrieved.version, "3.11.7");
198 assert_eq!(retrieved.sha256, runtime.sha256);
199 }
200
201 #[test]
202 fn test_get_nonexistent() {
203 let (cache, _temp) = create_test_cache();
204 let result = cache.get(Language::Python, "3.11.7");
205 assert!(result.is_none());
206 }
207
208 #[test]
209 fn test_clear() {
210 let (cache, _temp) = create_test_cache();
211 let data = b"test wasm data";
212 cache
213 .store(Language::Python, "3.11.7", data)
214 .expect("Failed to store");
215
216 assert!(cache.get(Language::Python, "3.11.7").is_some());
217
218 cache
219 .clear(Language::Python, "3.11.7")
220 .expect("Failed to clear");
221 assert!(cache.get(Language::Python, "3.11.7").is_none());
222 }
223
224 #[test]
225 fn test_clear_all() {
226 let (cache, _temp) = create_test_cache();
227 let data = b"test wasm data";
228 cache
229 .store(Language::Python, "3.11.7", data)
230 .expect("Failed to store");
231 cache
232 .store(Language::Ruby, "3.2.2", data)
233 .expect("Failed to store");
234
235 cache.clear_all().expect("Failed to clear all");
236
237 assert!(cache.get(Language::Python, "3.11.7").is_none());
238 assert!(cache.get(Language::Ruby, "3.2.2").is_none());
239 }
240
241 #[test]
242 fn test_list() {
243 let (cache, _temp) = create_test_cache();
244 let data = b"test wasm data";
245
246 cache
247 .store(Language::Python, "3.11.7", data)
248 .expect("Failed to store");
249 cache
250 .store(Language::Ruby, "3.2.2", data)
251 .expect("Failed to store");
252
253 let runtimes = cache.list().expect("Failed to list");
254 assert_eq!(runtimes.len(), 2);
255
256 let versions: Vec<String> = runtimes.iter().map(|r| r.version.clone()).collect();
257 assert!(versions.contains(&"3.11.7".to_string()));
258 assert!(versions.contains(&"3.2.2".to_string()));
259 }
260
261 #[test]
262 fn test_compute_sha256() {
263 let temp_dir = TempDir::new().unwrap();
264 let file_path = temp_dir.path().join("test.wasm");
265 let data = b"test data for hashing";
266 fs::write(&file_path, data).unwrap();
267
268 let hash1 = CacheManager::compute_sha256(&file_path).unwrap();
269 let hash2 = CacheManager::compute_sha256(&file_path).unwrap();
270
271 assert_eq!(hash1, hash2);
272 assert_eq!(hash1.len(), 64);
273 }
274
275 #[test]
276 fn test_verify_integrity() {
277 let (cache, _temp) = create_test_cache();
278 let data = b"test wasm data";
279 let runtime = cache
280 .store(Language::Python, "3.11.7", data)
281 .expect("Failed to store");
282
283 let result = cache.verify_integrity(&runtime, &runtime.sha256);
284 assert!(result.is_ok());
285
286 let result = cache.verify_integrity(&runtime, "invalid_hash");
287 assert!(result.is_err());
288 }
289}