Skip to main content

wasmer_wasix/runtime/package_loader/
builtin_loader.rs

1use std::{
2    collections::HashMap,
3    io::{ErrorKind, Read, Write as _},
4    path::PathBuf,
5    sync::{Arc, RwLock},
6};
7
8use anyhow::{Context, Error, bail};
9use bytes::Bytes;
10use http::{HeaderMap, Method};
11use tempfile::NamedTempFile;
12use url::Url;
13use wasmer_package::{
14    WasmerPackageError,
15    utils::{from_bytes, from_disk},
16};
17use webc::DetectError;
18use webc::{Container, ContainerError};
19
20use crate::{
21    bin_factory::BinaryPackage,
22    http::{HttpClient, HttpRequest, USER_AGENT},
23    runtime::{
24        package_loader::PackageLoader,
25        resolver::{DistributionInfo, PackageSummary, Resolution, WebcHash},
26    },
27};
28
29/// The builtin [`PackageLoader`] that is used by the `wasmer` CLI and
30/// respects `$WASMER_DIR`.
31#[derive(Debug)]
32pub struct BuiltinPackageLoader {
33    client: Arc<dyn HttpClient + Send + Sync>,
34    in_memory: Option<InMemoryCache>,
35    cache: Option<Arc<dyn PackageCache>>,
36    filesystem_cache: Option<Arc<FileSystemCache>>,
37    /// A mapping from hostnames to tokens
38    tokens: HashMap<String, String>,
39
40    hash_validation: HashIntegrityValidationMode,
41}
42
43/// Defines how to validate package hash integrity.
44#[derive(Clone, Copy, Debug, PartialEq, Eq)]
45pub enum HashIntegrityValidationMode {
46    /// Do not validate anything.
47    /// Best for performance.
48    NoValidate,
49    /// Compute the image hash and produce a trace warning on hash mismatches.
50    WarnOnHashMismatch,
51    /// Compute the image hash and fail on a mismatch.
52    FailOnHashMismatch,
53}
54
55impl BuiltinPackageLoader {
56    pub fn new() -> Self {
57        BuiltinPackageLoader {
58            in_memory: Some(InMemoryCache::default()),
59            client: Arc::new(crate::http::default_http_client().unwrap()),
60            cache: None,
61            filesystem_cache: None,
62            hash_validation: HashIntegrityValidationMode::NoValidate,
63            tokens: HashMap::new(),
64        }
65    }
66
67    /// Set the validation mode to apply after downloading an image.
68    ///
69    /// See [`HashIntegrityValidationMode`] for details.
70    pub fn with_hash_validation_mode(mut self, mode: HashIntegrityValidationMode) -> Self {
71        self.hash_validation = mode;
72        self
73    }
74
75    pub fn with_cache_dir(self, cache_dir: impl Into<PathBuf>) -> Self {
76        let cache = Arc::new(FileSystemCache {
77            cache_dir: cache_dir.into(),
78        });
79        BuiltinPackageLoader {
80            cache: Some(cache.clone()),
81            filesystem_cache: Some(cache),
82            ..self
83        }
84    }
85
86    /// Use a target-provided persistent package cache.
87    pub fn with_cache(self, cache: Arc<dyn PackageCache>) -> Self {
88        BuiltinPackageLoader {
89            cache: Some(cache),
90            filesystem_cache: None,
91            ..self
92        }
93    }
94
95    /// Disable promotion of loaded containers into the in-memory cache.
96    pub fn without_in_memory_cache(self) -> Self {
97        BuiltinPackageLoader {
98            in_memory: None,
99            ..self
100        }
101    }
102
103    pub fn cache(&self) -> Option<&FileSystemCache> {
104        self.filesystem_cache.as_deref()
105    }
106
107    pub fn validate_cache(
108        &self,
109        mode: CacheValidationMode,
110    ) -> Result<Vec<ImageHashMismatchError>, anyhow::Error> {
111        let cache = self
112            .filesystem_cache
113            .as_deref()
114            .context("can not validate cache - no filesystem cache configured")?;
115
116        let items = cache.validate_hashes()?;
117        let mut errors = Vec::new();
118        for (path, error) in items {
119            match mode {
120                CacheValidationMode::WarnOnMismatch => {
121                    tracing::warn!(?error, "hash mismatch in cached image file");
122                }
123                CacheValidationMode::PruneOnMismatch => {
124                    tracing::warn!(?error, "deleting cached image file due to hash mismatch");
125                    match std::fs::remove_file(&path) {
126                        Ok(()) => {}
127                        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
128                        Err(fs_err) => {
129                            tracing::error!(
130                                path=%error.source,
131                                ?fs_err,
132                                "could not delete cached image file with hash mismatch"
133                            );
134                        }
135                    }
136                }
137            }
138
139            errors.push(error);
140        }
141
142        Ok(errors)
143    }
144
145    pub fn with_http_client(self, client: impl HttpClient + Send + Sync + 'static) -> Self {
146        self.with_shared_http_client(Arc::new(client))
147    }
148
149    pub fn with_shared_http_client(self, client: Arc<dyn HttpClient + Send + Sync>) -> Self {
150        BuiltinPackageLoader { client, ..self }
151    }
152
153    pub fn with_tokens<I, K, V>(mut self, tokens: I) -> Self
154    where
155        I: IntoIterator<Item = (K, V)>,
156        K: Into<String>,
157        V: Into<String>,
158    {
159        for (hostname, token) in tokens {
160            self = self.with_token(hostname, token);
161        }
162
163        self
164    }
165
166    /// Add an API token that will be used whenever sending requests to a
167    /// particular hostname.
168    ///
169    /// Note that this uses [`Url::authority()`] when looking up tokens, so it
170    /// will match both plain hostnames (e.g. `registry.wasmer.io`) and hosts
171    /// with a port number (e.g. `localhost:8000`).
172    pub fn with_token(mut self, hostname: impl Into<String>, token: impl Into<String>) -> Self {
173        self.tokens.insert(hostname.into(), token.into());
174        self
175    }
176
177    /// Insert a container into the in-memory hash.
178    pub fn insert_cached(&self, hash: WebcHash, container: &Container) {
179        if let Some(in_memory) = &self.in_memory {
180            in_memory.save(container, hash);
181        }
182    }
183
184    /// Remove a container from the in-memory cache.
185    pub fn evict_cached(&self, hash: &WebcHash) -> Option<Container> {
186        self.in_memory
187            .as_ref()
188            .and_then(|in_memory| in_memory.remove(hash))
189    }
190
191    #[tracing::instrument(level = "debug", skip_all, fields(pkg.hash=%hash))]
192    async fn get_cached(&self, hash: &WebcHash) -> Result<Option<Container>, Error> {
193        if let Some(cached) = self
194            .in_memory
195            .as_ref()
196            .and_then(|in_memory| in_memory.lookup(hash))
197        {
198            return Ok(Some(cached));
199        }
200
201        if let Some(cache) = self.cache.as_ref() {
202            match cache.lookup(hash).await {
203                Ok(Some(cached)) => {
204                    if let Some(in_memory) = &self.in_memory {
205                        tracing::debug!("Copying from the persistent cache to the in-memory cache");
206                        in_memory.save(&cached, *hash);
207                    }
208                    return Ok(Some(cached));
209                }
210                Ok(None) => {}
211                Err(error) => {
212                    tracing::warn!(
213                        error = &*error,
214                        pkg.hash = %hash,
215                        "Unable to read a cached package; treating it as a cache miss",
216                    );
217                }
218            }
219        }
220
221        Ok(None)
222    }
223
224    /// Validate image contents with the specified validation mode.
225    async fn validate_hash(
226        image: &bytes::Bytes,
227        mode: HashIntegrityValidationMode,
228        info: &DistributionInfo,
229    ) -> Result<(), anyhow::Error> {
230        let info = info.clone();
231        let image = image.clone();
232        crate::spawn_blocking(move || Self::validate_hash_sync(&image, mode, &info))
233            .await
234            .context("tokio runtime failed")?
235    }
236
237    /// Validate image contents with the specified validation mode.
238    fn validate_hash_sync(
239        image: &[u8],
240        mode: HashIntegrityValidationMode,
241        info: &DistributionInfo,
242    ) -> Result<(), anyhow::Error> {
243        match mode {
244            HashIntegrityValidationMode::NoValidate => {
245                // Nothing to do.
246                Ok(())
247            }
248            HashIntegrityValidationMode::WarnOnHashMismatch => {
249                let actual_hash = WebcHash::sha256(image);
250                if actual_hash != info.webc_sha256 {
251                    tracing::warn!(%info.webc_sha256, %actual_hash, "image hash mismatch - actual image hash does not match the expected hash!");
252                }
253                Ok(())
254            }
255            HashIntegrityValidationMode::FailOnHashMismatch => {
256                let actual_hash = WebcHash::sha256(image);
257                if actual_hash != info.webc_sha256 {
258                    Err(ImageHashMismatchError {
259                        source: info.webc.to_string(),
260                        actual_hash,
261                        expected_hash: info.webc_sha256,
262                    }
263                    .into())
264                } else {
265                    Ok(())
266                }
267            }
268        }
269    }
270
271    #[tracing::instrument(level = "debug", skip_all, fields(%dist.webc, %dist.webc_sha256))]
272    async fn download(&self, dist: &DistributionInfo) -> Result<Bytes, Error> {
273        if dist.webc.scheme() == "file" {
274            match crate::runtime::resolver::utils::file_path_from_url(&dist.webc) {
275                Ok(path) => {
276                    let bytes = crate::spawn_blocking({
277                        let path = path.clone();
278                        move || std::fs::read(path)
279                    })
280                    .await?
281                    .with_context(|| format!("Unable to read \"{}\"", path.display()))?;
282
283                    let bytes = bytes::Bytes::from(bytes);
284
285                    Self::validate_hash(&bytes, self.hash_validation, dist).await?;
286
287                    return Ok(bytes);
288                }
289                Err(e) => {
290                    tracing::debug!(
291                        url=%dist.webc,
292                        error=&*e,
293                        "Unable to convert the file:// URL to a path",
294                    );
295                }
296            }
297        }
298
299        let request = HttpRequest {
300            headers: self.headers(&dist.webc),
301            url: dist.webc.clone(),
302            method: Method::GET,
303            body: None,
304            options: Default::default(),
305        };
306
307        tracing::debug!(%request.url, %request.method, "webc_package_download_start");
308        tracing::trace!(?request.headers);
309
310        let response = self.client.request(request).await?;
311
312        tracing::trace!(
313            %response.status,
314            %response.redirected,
315            ?response.headers,
316            response.len=response.body.as_ref().map(|body| body.len()),
317            "Received a response",
318        );
319
320        let url = &dist.webc;
321        if !response.is_ok() {
322            return Err(
323                crate::runtime::resolver::utils::http_error(&response).context(format!(
324                    "package download failed: GET request to \"{}\" failed with status {}",
325                    url, response.status
326                )),
327            );
328        }
329
330        let body = response.body.context("package download failed")?;
331        let body = Self::decode_response_body(&response.headers, body)
332            .context("package download failed: could not decode response body")?;
333        tracing::debug!(%url, "package_download_succeeded");
334
335        let body = bytes::Bytes::from(body);
336
337        Self::validate_hash(&body, self.hash_validation, dist).await?;
338
339        Ok(body)
340    }
341
342    fn headers(&self, url: &Url) -> HeaderMap {
343        let mut headers = HeaderMap::new();
344        headers.insert("Accept", "application/webc".parse().unwrap());
345        headers.insert("User-Agent", USER_AGENT.parse().unwrap());
346
347        // Accept compressed responses.
348        // NOTE: gzip and zstd decoding is available on native platforms.
349        // In browser platforms, the fetch implementation should automatically
350        // handle decoding of gzip/zstd responses transparently.
351        headers.insert(
352            http::header::ACCEPT_ENCODING,
353            "zstd;q=1.0, gzip;q=0.8".parse().unwrap(),
354        );
355
356        if url.has_authority()
357            && let Some(token) = self.tokens.get(url.authority())
358        {
359            let header = format!("Bearer {token}");
360            match header.parse() {
361                Ok(header) => {
362                    headers.insert(http::header::AUTHORIZATION, header);
363                }
364                Err(e) => {
365                    tracing::warn!(
366                        error = &e as &dyn std::error::Error,
367                        "An error occurred while parsing the authorization header",
368                    );
369                }
370            }
371        }
372
373        headers
374    }
375
376    /// Decode the response body according to the `Content-Encoding` header.
377    ///
378    /// * Supports `gzip` and `zstd` encodings
379    /// * Supports nested encodings (e.g. `gzip, zstd`)
380    /// * Passes through unencoded bodies or "identity" encoding unchanged
381    fn decode_response_body(headers: &HeaderMap, body: Vec<u8>) -> Result<Vec<u8>, anyhow::Error> {
382        let encodings = match headers.get(http::header::CONTENT_ENCODING) {
383            Some(header) => header
384                .to_str()
385                .context("non-utf8 content-encoding header")?
386                .split(',')
387                .map(|encoding| encoding.trim().to_ascii_lowercase())
388                .filter(|encoding| !encoding.is_empty())
389                .collect::<Vec<_>>(),
390            None => Vec::new(),
391        };
392
393        // Check if there is nothing to decode, return early.
394        // "identity" is the default encoding meaning "no encoding" (See RFC 2616 / RFC 7231)
395        if encodings.is_empty() || (encodings.len() == 1 && encodings[0] == "identity") {
396            return Ok(body);
397        }
398
399        let mut reader: Box<dyn Read> = Box::new(std::io::Cursor::new(body));
400        for encoding in encodings.iter().rev() {
401            match encoding.as_str() {
402                "gzip" => {
403                    reader = Box::new(flate2::read::GzDecoder::new(reader));
404                }
405                "zstd" => {
406                    #[cfg(not(target_arch = "wasm32"))]
407                    {
408                        reader = Box::new(
409                            zstd::stream::read::Decoder::new(reader)
410                                .context("failed to initialize zstd decoder")?,
411                        );
412                    }
413                    #[cfg(target_arch = "wasm32")]
414                    {
415                        // NOTE: in browsers this code will not be hit because
416                        // the fetch API automatically handles content decoding.
417                        bail!("zstd content-encoding is not supported on wasm32");
418                    }
419                }
420                "identity" => {}
421                other => bail!("unsupported content-encoding: {other}"),
422            }
423        }
424
425        let mut decoded = Vec::new();
426        reader
427            .read_to_end(&mut decoded)
428            .context("failed to decode response body")?;
429        Ok(decoded)
430    }
431}
432
433impl Default for BuiltinPackageLoader {
434    fn default() -> Self {
435        BuiltinPackageLoader::new()
436    }
437}
438
439#[async_trait::async_trait]
440impl PackageLoader for BuiltinPackageLoader {
441    #[tracing::instrument(
442        level="debug",
443        skip_all,
444        fields(
445            pkg=%summary.pkg.id,
446        ),
447    )]
448    async fn load(&self, summary: &PackageSummary) -> Result<Container, Error> {
449        if let Some(container) = self.get_cached(&summary.dist.webc_sha256).await? {
450            tracing::debug!("Cache hit!");
451            return Ok(container);
452        }
453
454        // looks like we had a cache miss and need to download it manually
455        let bytes = self
456            .download(&summary.dist)
457            .await
458            .with_context(|| format!("Unable to download \"{}\"", summary.dist.webc))?;
459
460        // We want to cache the container we downloaded, but we want to do it
461        // in a smart way to keep memory usage down.
462
463        if let Some(cache) = &self.cache {
464            match cache.save(bytes.clone(), &summary.dist).await {
465                Ok(container) => {
466                    tracing::debug!("Cached to disk");
467                    if let Some(in_memory) = &self.in_memory {
468                        in_memory.save(&container, summary.dist.webc_sha256);
469                    }
470                    // The happy path - we've saved to both caches and loaded the
471                    // container from disk (hopefully using mmap) so we're done.
472                    return Ok(container);
473                }
474                Err(e) => {
475                    tracing::warn!(
476                        error=&*e,
477                        pkg=%summary.pkg.id,
478                        pkg.hash=%summary.dist.webc_sha256,
479                        pkg.url=%summary.dist.webc,
480                        "Unable to save the downloaded package to disk",
481                    );
482                }
483            }
484        }
485
486        // The sad path - looks like we don't have a filesystem cache so we'll
487        // need to keep the whole thing in memory.
488        let container = crate::spawn_blocking(move || from_bytes(bytes)).await??;
489        if let Some(in_memory) = &self.in_memory {
490            // We still want to cache it in memory, of course
491            in_memory.save(&container, summary.dist.webc_sha256);
492        }
493        Ok(container)
494    }
495
496    async fn load_package_tree(
497        &self,
498        root: &Container,
499        resolution: &Resolution,
500        root_is_local_dir: bool,
501    ) -> Result<BinaryPackage, Error> {
502        super::load_package_tree(root, self, resolution, root_is_local_dir).await
503    }
504}
505
506#[derive(Clone, Debug)]
507pub struct ImageHashMismatchError {
508    source: String,
509    expected_hash: WebcHash,
510    actual_hash: WebcHash,
511}
512
513impl std::fmt::Display for ImageHashMismatchError {
514    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
515        write!(
516            f,
517            "image hash mismatch! expected hash '{}', but the computed hash is '{}' (source '{}')",
518            self.expected_hash, self.actual_hash, self.source,
519        )
520    }
521}
522
523impl std::error::Error for ImageHashMismatchError {}
524
525#[derive(Clone, Copy, Debug, PartialEq, Eq)]
526pub enum CacheValidationMode {
527    /// Just emit a warning for all images where the filename doesn't match
528    /// the expected hash.
529    WarnOnMismatch,
530    /// Remove images from the cache if the filename doesn't match the actual
531    /// hash.
532    PruneOnMismatch,
533}
534
535// FIXME: This implementation will block the async runtime and should use
536// some sort of spawn_blocking() call to run it in the background.
537#[derive(Debug)]
538pub struct FileSystemCache {
539    cache_dir: PathBuf,
540}
541
542/// Persistent storage for immutable WEBC package contents.
543#[async_trait::async_trait]
544pub trait PackageCache: Send + Sync + std::fmt::Debug {
545    /// Load and decode a package by its expected content hash.
546    async fn lookup(&self, hash: &WebcHash) -> Result<Option<Container>, Error>;
547
548    /// Persist downloaded bytes and return a decoded container.
549    async fn save(&self, webc: Bytes, dist: &DistributionInfo) -> Result<Container, Error>;
550}
551
552impl FileSystemCache {
553    const FILE_SUFFIX: &'static str = ".bin";
554
555    fn temp_dir(&self) -> PathBuf {
556        self.cache_dir.join("__temp__")
557    }
558
559    /// Validate that the cached image file names correspond to their actual
560    /// file content hashes.
561    fn validate_hashes(&self) -> Result<Vec<(PathBuf, ImageHashMismatchError)>, anyhow::Error> {
562        let mut items = Vec::<(PathBuf, ImageHashMismatchError)>::new();
563
564        let iter = match std::fs::read_dir(&self.cache_dir) {
565            Ok(v) => v,
566            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
567                // Cache dir does not exist, so nothing to validate.
568                return Ok(Vec::new());
569            }
570            Err(err) => {
571                return Err(err).with_context(|| {
572                    format!(
573                        "Could not read image cache dir: '{}'",
574                        self.cache_dir.display()
575                    )
576                });
577            }
578        };
579
580        for res in iter {
581            let entry = res?;
582            if !entry.file_type()?.is_file() {
583                continue;
584            }
585
586            // Extract the hash from the filename.
587
588            let hash_opt = entry
589                .file_name()
590                .to_str()
591                .and_then(|x| {
592                    let (raw_hash, _) = x.split_once(Self::FILE_SUFFIX)?;
593                    Some(raw_hash)
594                })
595                .and_then(|x| WebcHash::parse_hex(x).ok());
596            let Some(expected_hash) = hash_opt else {
597                continue;
598            };
599
600            // Compute the actual hash.
601            let path = entry.path();
602            let actual_hash = WebcHash::for_file(&path)?;
603
604            if actual_hash != expected_hash {
605                let err = ImageHashMismatchError {
606                    source: path.to_string_lossy().to_string(),
607                    actual_hash,
608                    expected_hash,
609                };
610                items.push((path, err));
611            }
612        }
613
614        Ok(items)
615    }
616
617    async fn lookup_impl(&self, hash: &WebcHash) -> Result<Option<Container>, Error> {
618        let path = self.path(hash);
619
620        let container = crate::spawn_blocking({
621            let path = path.clone();
622            move || from_disk(path)
623        })
624        .await?;
625        match container {
626            Ok(c) => Ok(Some(c)),
627            Err(WasmerPackageError::ContainerError(ContainerError::Open { error, .. }))
628            | Err(WasmerPackageError::ContainerError(ContainerError::Read { error, .. }))
629            | Err(WasmerPackageError::ContainerError(ContainerError::Detect(DetectError::Io(
630                error,
631            )))) if error.kind() == ErrorKind::NotFound => Ok(None),
632            Err(e) => {
633                let msg = format!("Unable to read \"{}\"", path.display());
634                Err(Error::new(e).context(msg))
635            }
636        }
637    }
638
639    async fn save_impl(&self, webc: Bytes, dist: &DistributionInfo) -> Result<PathBuf, Error> {
640        let path = self.path(&dist.webc_sha256);
641        let dist = dist.clone();
642        let temp_dir = self.temp_dir();
643
644        let path2 = path.clone();
645        crate::spawn_blocking(move || {
646            // Keep files in a temporary directory until they are fully written
647            // to prevent temp files being included in [`Self::scan`] or `[Self::retain]`.
648
649            std::fs::create_dir_all(&temp_dir)
650                .with_context(|| format!("Unable to create directory '{}'", temp_dir.display()))?;
651
652            let mut temp = NamedTempFile::new_in(&temp_dir)?;
653            temp.write_all(&webc)?;
654            temp.flush()?;
655            temp.as_file_mut().sync_all()?;
656
657            // Move the temporary file to the final location.
658            temp.persist(&path)?;
659
660            tracing::debug!(
661                pkg.hash=%dist.webc_sha256,
662                pkg.url=%dist.webc,
663                path=%path.display(),
664                num_bytes=webc.len(),
665                "Saved to disk",
666            );
667            Result::<_, Error>::Ok(())
668        })
669        .await??;
670
671        Ok(path2)
672    }
673
674    #[tracing::instrument(level = "debug", skip_all)]
675    async fn save_and_load_as_mmapped(
676        &self,
677        webc: Bytes,
678        dist: &DistributionInfo,
679    ) -> Result<Container, Error> {
680        // First, save it to disk
681        self.save_impl(webc, dist).await?;
682
683        // Now try to load it again. The resulting container should use
684        // a memory-mapped file rather than an in-memory buffer.
685        match self.lookup_impl(&dist.webc_sha256).await? {
686            Some(container) => Ok(container),
687            None => {
688                // Something really weird has occurred and we can't see the
689                // saved file. Just error out and let the fallback code do its
690                // thing.
691                Err(Error::msg("Unable to load the downloaded memory from disk"))
692            }
693        }
694    }
695
696    fn path(&self, hash: &WebcHash) -> PathBuf {
697        self.cache_dir.join(format!(
698            "{}{}",
699            hex::encode(hash.as_bytes()),
700            Self::FILE_SUFFIX
701        ))
702    }
703
704    /// Scan all the cached webc files and invoke the callback for each.
705    pub async fn scan<S, F>(&self, state: S, callback: F) -> Result<S, Error>
706    where
707        S: Send + 'static,
708        F: Fn(&mut S, &std::fs::DirEntry) -> Result<(), Error> + Send + 'static,
709    {
710        let cache_dir = self.cache_dir.clone();
711        tokio::task::spawn_blocking(move || -> Result<S, anyhow::Error> {
712            let mut state = state;
713
714            let iter = match std::fs::read_dir(&cache_dir) {
715                Ok(v) => v,
716                Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
717                    // path does not exist, so nothing to scan.
718                    return Ok(state);
719                }
720                Err(err) => {
721                    return Err(err).with_context(|| {
722                        format!("Could not read image cache dir: '{}'", cache_dir.display())
723                    });
724                }
725            };
726
727            for res in iter {
728                let entry = res?;
729                if !entry.file_type()?.is_file() {
730                    continue;
731                }
732
733                callback(&mut state, &entry)?;
734            }
735
736            Ok(state)
737        })
738        .await?
739        .context("tokio runtime failed")
740    }
741
742    /// Remove entries from the cache that do not pass the callback.
743    pub async fn retain<S, F>(&self, state: S, filter: F) -> Result<S, Error>
744    where
745        S: Send + 'static,
746        F: Fn(&mut S, &std::fs::DirEntry) -> Result<bool, anyhow::Error> + Send + 'static,
747    {
748        let cache_dir = self.cache_dir.clone();
749        tokio::task::spawn_blocking(move || {
750            let iter = match std::fs::read_dir(&cache_dir) {
751                Ok(v) => v,
752                Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
753                    // path does not exist, so nothing to scan.
754                    return Ok(state);
755                }
756                Err(err) => {
757                    return Err(err).with_context(|| {
758                        format!("Could not read image cache dir: '{}'", cache_dir.display())
759                    });
760                }
761            };
762
763            let mut state = state;
764            for res in iter {
765                let entry = res?;
766                if !entry.file_type()?.is_file() {
767                    continue;
768                }
769
770                if !filter(&mut state, &entry)? {
771                    tracing::debug!(
772                        path=%entry.path().display(),
773                        "Removing cached image file - does not pass the filter",
774                    );
775                    match std::fs::remove_file(entry.path()) {
776                        Ok(()) => {}
777                        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
778                        Err(fs_err) => {
779                            tracing::warn!(
780                                path=%entry.path().display(),
781                                ?fs_err,
782                                "Could not delete cached image file",
783                            );
784                        }
785                    }
786                }
787            }
788
789            Ok(state)
790        })
791        .await?
792        .context("tokio runtime failed")
793    }
794}
795
796#[async_trait::async_trait]
797impl PackageCache for FileSystemCache {
798    async fn lookup(&self, hash: &WebcHash) -> Result<Option<Container>, Error> {
799        self.lookup_impl(hash).await
800    }
801
802    async fn save(&self, webc: Bytes, dist: &DistributionInfo) -> Result<Container, Error> {
803        self.save_and_load_as_mmapped(webc, dist).await
804    }
805}
806
807#[derive(Debug, Default)]
808struct InMemoryCache(RwLock<HashMap<WebcHash, Container>>);
809
810impl InMemoryCache {
811    fn lookup(&self, hash: &WebcHash) -> Option<Container> {
812        self.0.read().unwrap().get(hash).cloned()
813    }
814
815    fn save(&self, container: &Container, hash: WebcHash) {
816        let mut cache = self.0.write().unwrap();
817        cache.entry(hash).or_insert_with(|| container.clone());
818    }
819
820    fn remove(&self, hash: &WebcHash) -> Option<Container> {
821        self.0.write().unwrap().remove(hash)
822    }
823}
824
825#[cfg(test)]
826mod tests {
827    use std::{collections::VecDeque, io::Write, sync::Mutex};
828
829    use futures::future::BoxFuture;
830    use http::{HeaderMap, HeaderValue, StatusCode};
831    use tempfile::TempDir;
832    use wasmer_config::package::PackageId;
833
834    use crate::{
835        http::{HttpRequest, HttpResponse},
836        runtime::resolver::PackageInfo,
837    };
838
839    use super::*;
840
841    const PYTHON: &[u8] =
842        include_bytes!("../../../../../wasmer-test-files/examples/python--python@3.13.5.webc");
843
844    #[derive(Debug)]
845    pub(crate) struct DummyClient {
846        requests: Mutex<Vec<HttpRequest>>,
847        responses: Mutex<VecDeque<HttpResponse>>,
848    }
849
850    impl DummyClient {
851        pub fn with_responses(responses: impl IntoIterator<Item = HttpResponse>) -> Self {
852            DummyClient {
853                requests: Mutex::new(Vec::new()),
854                responses: Mutex::new(responses.into_iter().collect()),
855            }
856        }
857    }
858
859    impl HttpClient for DummyClient {
860        fn request(
861            &self,
862            request: HttpRequest,
863        ) -> BoxFuture<'_, Result<HttpResponse, anyhow::Error>> {
864            let response = self.responses.lock().unwrap().pop_front().unwrap();
865            self.requests.lock().unwrap().push(request);
866            Box::pin(async { Ok(response) })
867        }
868    }
869
870    async fn cache_misses_will_trigger_a_download_internal() {
871        let temp = TempDir::new().unwrap();
872        let client = Arc::new(DummyClient::with_responses([HttpResponse {
873            body: Some(PYTHON.to_vec()),
874            redirected: false,
875            status: StatusCode::OK,
876            headers: HeaderMap::new(),
877        }]));
878        let loader = BuiltinPackageLoader::new()
879            .with_cache_dir(temp.path())
880            .with_shared_http_client(client.clone());
881        let summary = PackageSummary {
882            pkg: PackageInfo {
883                id: PackageId::new_named("python/python", "0.1.0".parse().unwrap()),
884                dependencies: Vec::new(),
885                commands: Vec::new(),
886                entrypoint: Some("asdf".to_string()),
887                filesystem: Vec::new(),
888            },
889            dist: DistributionInfo {
890                webc: "https://wasmer.io/python/python".parse().unwrap(),
891                webc_sha256: [0xaa; 32].into(),
892            },
893        };
894
895        let container = loader.load(&summary).await.unwrap();
896
897        // A HTTP request was sent
898        let requests = client.requests.lock().unwrap();
899        let request = &requests[0];
900        assert_eq!(request.url, summary.dist.webc);
901        assert_eq!(request.method, "GET");
902        #[cfg(not(target_arch = "wasm32"))]
903        {
904            assert_eq!(request.headers.len(), 3);
905            assert_eq!(request.headers["Accept-Encoding"], "zstd;q=1.0, gzip;q=0.8");
906        }
907        #[cfg(target_arch = "wasm32")]
908        {
909            assert_eq!(request.headers.len(), 2);
910            assert!(!request.headers.contains_key(http::header::ACCEPT_ENCODING));
911        }
912        assert_eq!(request.headers["Accept"], "application/webc");
913        assert_eq!(request.headers["User-Agent"], USER_AGENT);
914        // Make sure we got the right package
915        let manifest = container.manifest();
916        assert_eq!(manifest.entrypoint.as_deref(), Some("python"));
917        // it should have been automatically saved to disk
918        let path = loader
919            .filesystem_cache
920            .as_ref()
921            .unwrap()
922            .path(&summary.dist.webc_sha256);
923        assert!(path.exists());
924        assert_eq!(std::fs::read(&path).unwrap(), PYTHON);
925        // and cached in memory for next time
926        let in_memory = loader.in_memory.as_ref().unwrap().0.read().unwrap();
927        assert!(in_memory.contains_key(&summary.dist.webc_sha256));
928    }
929
930    #[cfg(not(target_arch = "wasm32"))]
931    #[tokio::test(flavor = "multi_thread")]
932    async fn cache_misses_will_trigger_a_download() {
933        cache_misses_will_trigger_a_download_internal().await
934    }
935
936    #[cfg(not(target_arch = "wasm32"))]
937    #[tokio::test]
938    async fn can_disable_in_memory_cache() {
939        let temp = TempDir::new().unwrap();
940        let client = Arc::new(DummyClient::with_responses([HttpResponse {
941            body: Some(PYTHON.to_vec()),
942            redirected: false,
943            status: StatusCode::OK,
944            headers: HeaderMap::new(),
945        }]));
946        let loader = BuiltinPackageLoader::new()
947            .with_cache_dir(temp.path())
948            .without_in_memory_cache()
949            .with_shared_http_client(client);
950        let summary = PackageSummary {
951            pkg: PackageInfo {
952                id: PackageId::new_named("python/python", "0.1.0".parse().unwrap()),
953                dependencies: Vec::new(),
954                commands: Vec::new(),
955                entrypoint: Some("asdf".to_string()),
956                filesystem: Vec::new(),
957            },
958            dist: DistributionInfo {
959                webc: "https://wasmer.io/python/python".parse().unwrap(),
960                webc_sha256: [0xbb; 32].into(),
961            },
962        };
963
964        loader.load(&summary).await.unwrap();
965
966        assert!(loader.in_memory.is_none());
967    }
968
969    #[cfg(target_arch = "wasm32")]
970    #[tokio::test()]
971    async fn cache_misses_will_trigger_a_download() {
972        cache_misses_will_trigger_a_download_internal().await
973    }
974
975    #[tokio::test]
976    async fn evict_cached_removes_in_memory_container() {
977        let loader = BuiltinPackageLoader::new();
978        let container = from_bytes(PYTHON).unwrap();
979        let hash: WebcHash = [0xaa; 32].into();
980        loader.insert_cached(hash, &container);
981        let evicted = loader.evict_cached(&hash);
982        assert!(evicted.is_some());
983        {
984            let in_memory = loader.in_memory.as_ref().unwrap().0.read().unwrap();
985            assert!(!in_memory.contains_key(&hash));
986        }
987        assert!(loader.evict_cached(&hash).is_none());
988    }
989
990    /// Small helper to construct headers with a given content-encoding.
991    fn headers_with_encoding(content_encoding: Option<&str>) -> HeaderMap {
992        let mut headers = HeaderMap::new();
993        if let Some(value) = content_encoding {
994            headers.insert(http::header::CONTENT_ENCODING, value.parse().unwrap());
995        }
996        headers
997    }
998
999    /// Small helper to construct headers with a raw content-encoding value.
1000    fn headers_with_raw_encoding(value: &[u8]) -> HeaderMap {
1001        let mut headers = HeaderMap::new();
1002        headers.insert(
1003            http::header::CONTENT_ENCODING,
1004            HeaderValue::from_bytes(value).unwrap(),
1005        );
1006        headers
1007    }
1008
1009    /// Confirm decode_response_body passes through unencoded bodies unchanged.
1010    #[test]
1011    fn decode_response_body_passthrough() {
1012        let body = b"plain-bytes".to_vec();
1013
1014        let decoded =
1015            BuiltinPackageLoader::decode_response_body(&headers_with_encoding(None), body.clone())
1016                .unwrap();
1017        assert_eq!(decoded, body);
1018
1019        let decoded = BuiltinPackageLoader::decode_response_body(
1020            &headers_with_encoding(Some("identity")),
1021            body.clone(),
1022        )
1023        .unwrap();
1024        assert_eq!(decoded, body);
1025    }
1026
1027    /// Confirm decode_response_body treats empty/whitespace encoding lists as no encoding.
1028    #[test]
1029    fn decode_response_body_empty_encoding_list() {
1030        let body = b"plain-bytes".to_vec();
1031        let decoded = BuiltinPackageLoader::decode_response_body(
1032            &headers_with_encoding(Some(" , , ")),
1033            body.clone(),
1034        )
1035        .unwrap();
1036        assert_eq!(decoded, body);
1037    }
1038
1039    /// Confirm decode_response_body errors on non-utf8 content-encoding headers.
1040    #[test]
1041    fn decode_response_body_non_utf8_encoding_header() {
1042        let body = b"bytes".to_vec();
1043        let err =
1044            BuiltinPackageLoader::decode_response_body(&headers_with_raw_encoding(&[0xff]), body)
1045                .unwrap_err();
1046        let msg = err.to_string();
1047        assert!(msg.contains("non-utf8 content-encoding"));
1048    }
1049
1050    /// Confirm decode_response_body decodes gzip-encoded bodies.
1051    #[test]
1052    fn decode_response_body_gzip() {
1053        let body = b"gzip-bytes".to_vec();
1054        let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
1055        encoder.write_all(&body).unwrap();
1056        let encoded = encoder.finish().unwrap();
1057
1058        let decoded = BuiltinPackageLoader::decode_response_body(
1059            &headers_with_encoding(Some("gzip")),
1060            encoded,
1061        )
1062        .unwrap();
1063        assert_eq!(decoded, body);
1064    }
1065
1066    /// Confirm decode_response_body ignores identity when combined with other encodings.
1067    #[test]
1068    fn decode_response_body_identity_and_gzip() {
1069        let body = b"gzip-bytes".to_vec();
1070        let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
1071        encoder.write_all(&body).unwrap();
1072        let encoded = encoder.finish().unwrap();
1073
1074        let decoded = BuiltinPackageLoader::decode_response_body(
1075            &headers_with_encoding(Some("identity, gzip")),
1076            encoded,
1077        )
1078        .unwrap();
1079        assert_eq!(decoded, body);
1080    }
1081
1082    /// Confirm decode_response_body errors on invalid gzip payloads.
1083    #[test]
1084    fn decode_response_body_gzip_invalid_payload() {
1085        let body = b"not-gzip".to_vec();
1086        let err =
1087            BuiltinPackageLoader::decode_response_body(&headers_with_encoding(Some("gzip")), body)
1088                .unwrap_err();
1089        let msg = err.to_string();
1090        assert!(msg.contains("failed to decode response body"));
1091    }
1092
1093    /// Confirm decode_response_body decodes zstd-encoded bodies.
1094    #[cfg(not(target_arch = "wasm32"))]
1095    #[test]
1096    fn decode_response_body_zstd() {
1097        let body = b"zstd-bytes".to_vec();
1098        let encoded = zstd::stream::encode_all(std::io::Cursor::new(&body), 0).unwrap();
1099
1100        let decoded = BuiltinPackageLoader::decode_response_body(
1101            &headers_with_encoding(Some("zstd")),
1102            encoded,
1103        )
1104        .unwrap();
1105        assert_eq!(decoded, body);
1106    }
1107
1108    /// Confirm decode_response_body errors on invalid zstd payloads.
1109    #[cfg(not(target_arch = "wasm32"))]
1110    #[test]
1111    fn decode_response_body_zstd_invalid_payload() {
1112        let body = b"not-zstd".to_vec();
1113        let err =
1114            BuiltinPackageLoader::decode_response_body(&headers_with_encoding(Some("zstd")), body)
1115                .unwrap_err();
1116        let msg = err.to_string();
1117        assert!(msg.contains("failed to decode response body"));
1118    }
1119
1120    /// Confirm decode_response_body decodes layered gzip+zstd-encoded bodies.
1121    #[cfg(not(target_arch = "wasm32"))]
1122    #[test]
1123    fn decode_response_body_zstd_and_gzip() {
1124        let body = b"layered-bytes".to_vec();
1125        let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
1126        encoder.write_all(&body).unwrap();
1127        let gzipped = encoder.finish().unwrap();
1128        let encoded = zstd::stream::encode_all(std::io::Cursor::new(gzipped), 0).unwrap();
1129
1130        let decoded = BuiltinPackageLoader::decode_response_body(
1131            &headers_with_encoding(Some("gzip, zstd")),
1132            encoded,
1133        )
1134        .unwrap();
1135        assert_eq!(decoded, body);
1136    }
1137
1138    /// Confirm decode_response_body errors on unknown encodings.
1139    #[test]
1140    fn decode_response_body_unknown_encoding() {
1141        let body = b"weird".to_vec();
1142        let err =
1143            BuiltinPackageLoader::decode_response_body(&headers_with_encoding(Some("br")), body)
1144                .unwrap_err();
1145        let msg = err.to_string();
1146        assert!(msg.contains("unsupported content-encoding"));
1147    }
1148
1149    // NOTE: must be a tokio test because the BuiltinPackageLoader::new()
1150    // constructor requires a runtime...
1151    #[tokio::test]
1152    async fn test_builtin_package_downloader_cache_validation() {
1153        let dir = tempfile::tempdir().unwrap();
1154        let path = dir.path();
1155
1156        let contents = "fail";
1157        let correct_hash = WebcHash::sha256(contents);
1158        let used_hash =
1159            WebcHash::parse_hex("0000a28ea38a000f3a3328cb7fabe330638d3258affe1a869e3f92986222d997")
1160                .unwrap();
1161        let filename = format!("{}{}", used_hash, FileSystemCache::FILE_SUFFIX);
1162        let file_path = path.join(filename);
1163        std::fs::write(&file_path, contents).unwrap();
1164
1165        let dl = BuiltinPackageLoader::new().with_cache_dir(path);
1166
1167        let errors = dl
1168            .validate_cache(CacheValidationMode::PruneOnMismatch)
1169            .unwrap();
1170        assert_eq!(errors.len(), 1);
1171        assert_eq!(errors[0].actual_hash, correct_hash);
1172        assert_eq!(errors[0].expected_hash, used_hash);
1173
1174        assert_eq!(file_path.exists(), false);
1175    }
1176
1177    #[tokio::test]
1178    async fn test_file_cache_scan_retain() {
1179        let dir = tempfile::tempdir().unwrap();
1180        let path = dir.path();
1181
1182        let cache = FileSystemCache {
1183            cache_dir: path.to_path_buf(),
1184        };
1185
1186        {
1187            let state = cache
1188                .scan(0u64, |state: &mut u64, _entry| {
1189                    *state += 1;
1190                    Ok(())
1191                })
1192                .await
1193                .unwrap();
1194
1195            assert_eq!(state, 0);
1196        }
1197
1198        let path1 = cache
1199            .save_impl(
1200                Bytes::from_static(b"test1"),
1201                &DistributionInfo {
1202                    webc: Url::parse("file:///test1.webc").unwrap(),
1203                    webc_sha256: WebcHash::sha256(b"test1"),
1204                },
1205            )
1206            .await
1207            .unwrap();
1208        let path2 = cache
1209            .save_impl(
1210                Bytes::from_static(b"test2"),
1211                &DistributionInfo {
1212                    webc: Url::parse("file:///test2.webc").unwrap(),
1213                    webc_sha256: WebcHash::sha256(b"test2"),
1214                },
1215            )
1216            .await
1217            .unwrap();
1218
1219        {
1220            let path1 = path1.clone();
1221            let path2 = path2.clone();
1222            let state = cache
1223                .scan(0u64, move |state: &mut u64, entry| {
1224                    *state += 1;
1225                    assert!(entry.path() == path1 || entry.path() == path2);
1226                    Ok(())
1227                })
1228                .await
1229                .unwrap();
1230
1231            assert_eq!(state, 2);
1232        }
1233
1234        {
1235            let path1 = path1.clone();
1236            let state = cache
1237                .retain(0u64, move |state: &mut u64, entry| {
1238                    *state += 1;
1239                    Ok(entry.path() == path1)
1240                })
1241                .await
1242                .unwrap();
1243            assert_eq!(state, 2);
1244        }
1245
1246        assert!(path1.exists());
1247        assert!(!path2.exists(), "Path 2 should have been deleted");
1248
1249        {
1250            let path1 = path1.clone();
1251            let state = cache
1252                .scan(0u64, move |state: &mut u64, entry| {
1253                    *state += 1;
1254                    assert!(entry.path() == path1);
1255                    Ok(())
1256                })
1257                .await
1258                .unwrap();
1259            assert_eq!(state, 1);
1260        }
1261    }
1262}