wasmer_cli/commands/run/

wasi.rs

use std::{
    collections::{BTreeSet, HashMap},
    path::{Path, PathBuf},
    str::FromStr,
    sync::{Arc, mpsc::Sender},
    time::Duration,
};

use anyhow::{Context, Result, bail};
use bytes::Bytes;
use clap::Parser;
use itertools::Itertools;
use tokio::runtime::Handle;
use url::Url;
use virtual_fs::{
    ArcFileSystem, DeviceFile, FileSystem, MountFileSystem, OverlayFileSystem,
    RootFileSystemBuilder,
};
use virtual_net::ruleset::Ruleset;
use wasmer::{Engine, Function, Instance, Memory32, Memory64, Module, RuntimeError, Store, Value};
use wasmer_config::package::PackageSource as PackageSpecifier;
use wasmer_types::ModuleHash;
#[cfg(feature = "journal")]
use wasmer_wasix::journal::{LogFileJournal, SnapshotTrigger};
use wasmer_wasix::{
    PluggableRuntime, RewindState, Runtime, WasiEnv, WasiEnvBuilder, WasiError, WasiFunctionEnv,
    WasiVersion,
    bin_factory::BinaryPackage,
    capabilities::Capabilities,
    get_wasi_versions,
    http::HttpClient,
    journal::{CompactingLogFileJournal, DynJournal, DynReadableJournal},
    os::{TtyBridge, tty_sys::SysTty},
    rewind_ext,
    runners::MAPPED_CURRENT_DIR_DEFAULT_PATH,
    runners::{MappedCommand, MappedDirectory, MountedDirectory},
    runtime::{
        module_cache::{FileSystemCache, ModuleCache},
        package_loader::{BuiltinPackageLoader, PackageLoader},
        resolver::{
            BackendSource, FileSystemSource, InMemorySource, LocalRegistrySource, MultiSource,
            Source, WebSource,
        },
        task_manager::{
            VirtualTaskManagerExt,
            tokio::{RuntimeOrHandle, TokioTaskManager},
        },
    },
    types::__WASI_STDIN_FILENO,
    wasmer_wasix_types::wasi::Errno,
};

use crate::{
    config::{UserRegistry, WasmerEnv},
    utils::{
        WAPM_SOURCE_CACHE_TIMEOUT, parse_envvar, parse_mapdir, parse_volume,
        registry_query_cache_dir,
    },
};

use super::{
    CliPackageSource, ExecutableTarget,
    capabilities::{self, PkgCapabilityCache},
};

#[derive(Debug, Parser, Clone, Default)]
/// WASI Options
pub struct Wasi {
    /// Map a host directory to a different location for the Wasm module
    #[clap(
        long = "volume",
        name = "[HOST_DIR:]GUEST_DIR",
        value_parser = parse_volume,
    )]
    pub(crate) volumes: Vec<MappedDirectory>,

    // Legacy option
    #[clap(long = "dir", group = "wasi", hide = true)]
    pub(crate) pre_opened_directories: Vec<PathBuf>,

    // Legacy option
    #[clap(
        long = "mapdir",
        value_parser = parse_mapdir,
        hide = true
     )]
    pub(crate) mapped_dirs: Vec<MappedDirectory>,

    /// Set the module's initial CWD to this path; does not work with
    /// WASI preview 1 modules.
    #[clap(long = "cwd")]
    pub(crate) cwd: Option<PathBuf>,

    /// Pass custom environment variables
    #[clap(
        long = "env",
        name = "KEY=VALUE",
        value_parser=parse_envvar,
    )]
    pub(crate) env_vars: Vec<(String, String)>,

    /// Forward all host env variables to guest
    #[clap(long, env)]
    pub(crate) forward_host_env: bool,

    /// List of other containers this module depends on
    #[clap(long = "use", name = "USE")]
    pub(crate) uses: Vec<String>,

    /// Webc packages that are explicitly included for execution, taking
    /// precedence over the registry ones: either a `*.webc` file, or a
    /// directory laid out `<namespace>/<name>/<version>.webc` and queried on
    /// demand like a registry. Resolves named dependencies from local files,
    /// offline.
    #[clap(long = "include-webc", name = "WEBC")]
    pub(super) include_webcs: Vec<PathBuf>,

    /// Resolve only from local sources (`--include-webc`, filesystem paths), never
    /// the registry. Needed offline, where a registry query would fail resolution.
    #[clap(long = "offline")]
    pub(super) offline: bool,

    /// List of injected atoms
    #[clap(long = "map-command", name = "MAPCMD")]
    pub(super) map_commands: Vec<String>,

    /// Enable networking with the host network.
    ///
    /// Allows WASI modules to open TCP and UDP connections, create sockets, ...
    ///
    /// Optionally, a set of network filters could be defined which allows fine-grained
    /// control over the network sandbox.
    ///
    /// Rule Syntax:
    ///
    /// <rule-type>:<allow|deny>=<rule-expression>
    ///
    /// Examples:
    ///
    ///  - Allow a specific domain and port: dns:allow=example.com:80
    ///
    ///  - Deny a domain and all its subdomains on all ports: dns:deny=*danger.xyz:*
    ///
    ///  - Allow opening ipv4 sockets only on a specific IP and port: ipv4:allow=127.0.0.1:80/in.
    #[clap(long = "net", require_equals = true)]
    // Note that when --net is passed to the cli, the first Option will be initialized: Some(None)
    // and when --net=<ruleset> is specified, the inner Option will be initialized: Some(Some(ruleset))
    pub networking: Option<Option<String>>,

    /// Disables the TTY bridge
    #[clap(long = "no-tty")]
    pub no_tty: bool,

    /// Enables or disables asynchronous threading.
    ///
    /// If omitted, the runtime default is used.
    #[clap(
        long = "enable-async-threads",
        require_equals = true,
        default_missing_value = "true",
        num_args = 0..=1,
        action = clap::ArgAction::Set
    )]
    pub enable_async_threads: Option<bool>,

    /// Enables an exponential backoff (measured in milli-seconds) of
    /// the process CPU usage when there are no active run tokens (when set
    /// holds the maximum amount of time that it will pause the CPU)
    /// (default = off)
    #[clap(long = "enable-cpu-backoff")]
    pub enable_cpu_backoff: Option<u64>,

    /// Specifies one or more journal files that Wasmer will use to restore
    /// the state of the WASM process as it executes.
    ///
    /// The state of the WASM process and its sandbox will be reapplied using
    /// the journals in the order that you specify here.
    #[cfg(feature = "journal")]
    #[clap(long = "journal")]
    pub read_only_journals: Vec<PathBuf>,

    /// Specifies one or more journal files that Wasmer will use to restore
    /// and save the state of the WASM process as it executes.
    ///
    /// The state of the WASM process and its sandbox will be reapplied using
    /// the journals in the order that you specify here.
    ///
    /// The last journal file specified will be created if it does not exist
    /// and opened for read and write. New journal events will be written to this
    /// file
    #[cfg(feature = "journal")]
    #[clap(long = "journal-writable")]
    pub writable_journals: Vec<PathBuf>,

    /// Flag that indicates if the journal will be automatically compacted
    /// as it fills up and when the process exits
    #[cfg(feature = "journal")]
    #[clap(long = "enable-compaction")]
    pub enable_compaction: bool,

    /// Tells the compactor not to compact when the journal log file is closed
    #[cfg(feature = "journal")]
    #[clap(long = "without-compact-on-drop")]
    pub without_compact_on_drop: bool,

    /// Tells the compactor to compact when it grows by a certain factor of
    /// its original size. (i.e. '0.2' would be it compacts after the journal
    /// has grown by 20 percent)
    ///
    /// Default is to compact on growth that exceeds 15%
    #[cfg(feature = "journal")]
    #[clap(long = "with-compact-on-growth", default_value = "0.15")]
    pub with_compact_on_growth: f32,

    /// Indicates what events will cause a snapshot to be taken
    /// and written to the journal file.
    ///
    /// If not specified, the default is to snapshot when the process idles, when
    /// the process exits or periodically if an interval argument is also supplied,
    /// as well as when the process requests a snapshot explicitly.
    ///
    /// Additionally if the snapshot-on is not specified it will also take a snapshot
    /// on the first stdin, environ or socket listen - this can be used to accelerate
    /// the boot up time of WASM processes.
    #[cfg(feature = "journal")]
    #[clap(long = "snapshot-on")]
    pub snapshot_on: Vec<SnapshotTrigger>,

    /// Adds a periodic interval (measured in milli-seconds) that the runtime will automatically
    /// take snapshots of the running process and write them to the journal. When specifying
    /// this parameter it implies that `--snapshot-on interval` has also been specified.
    #[cfg(feature = "journal")]
    #[clap(long = "snapshot-period")]
    pub snapshot_interval: Option<u64>,

    /// If specified, the runtime will stop executing the WASM module after the first snapshot
    /// is taken.
    #[cfg(feature = "journal")]
    #[clap(long = "stop-after-snapshot")]
    pub stop_after_snapshot: bool,

    /// Skip writes to stdout and stderr when replying journal events to bootstrap a module.
    #[cfg(feature = "journal")]
    #[clap(long = "skip-journal-stdio")]
    pub skip_stdio_during_bootstrap: bool,

    /// Allow instances to send http requests.
    ///
    /// Access to domains is granted by default.
    #[clap(long)]
    pub http_client: bool,

    /// Require WASI modules to only import 1 version of WASI.
    #[clap(long = "deny-multiple-wasi-versions")]
    pub deny_multiple_wasi_versions: bool,

    /// Disable the cache for the compiled modules.
    ///
    /// Cache is used to speed up the loading of modules, as the
    /// generated artifacts are cached.
    #[clap(long = "disable-cache")]
    disable_cache: bool,
}

pub struct RunProperties {
    pub ctx: WasiFunctionEnv,
    pub path: PathBuf,
    pub invoke: Option<String>,
    pub args: Vec<String>,
}

#[allow(dead_code)]
impl Wasi {
    pub fn map_dir(&mut self, alias: &str, target_on_disk: PathBuf) {
        self.volumes.push(MappedDirectory {
            guest: alias.to_string(),
            host: target_on_disk,
        });
    }

    pub fn set_env(&mut self, key: &str, value: &str) {
        self.env_vars.push((key.to_string(), value.to_string()));
    }

    /// Gets the WASI version (if any) for the provided module
    pub fn get_versions(module: &Module) -> Option<BTreeSet<WasiVersion>> {
        // Get the wasi version in non-strict mode, so multiple wasi versions
        // are potentially allowed.
        //
        // Checking for multiple wasi versions is handled outside this function.
        get_wasi_versions(module, false)
    }

    /// Checks if a given module has any WASI imports at all.
    pub fn has_wasi_imports(module: &Module) -> bool {
        // Get the wasi version in non-strict mode, so no other imports
        // are allowed
        get_wasi_versions(module, false).is_some()
    }

    pub(crate) fn all_volumes(&self) -> Vec<MappedDirectory> {
        self.volumes
            .iter()
            .cloned()
            .chain(self.pre_opened_directories.iter().map(|d| MappedDirectory {
                host: d.clone(),
                guest: d.to_str().expect("must be a valid path string").to_string(),
            }))
            .chain(self.mapped_dirs.iter().cloned())
            .collect_vec()
    }

    pub fn prepare(
        &self,
        module: &Module,
        program_name: String,
        args: Vec<String>,
        rt: Arc<dyn Runtime + Send + Sync>,
    ) -> Result<WasiEnvBuilder> {
        let args = args.into_iter().map(|arg| arg.into_bytes());

        let map_commands = self
            .map_commands
            .iter()
            .map(|map| map.split_once('=').unwrap())
            .map(|(a, b)| (a.to_string(), b.to_string()))
            .collect::<HashMap<_, _>>();

        let mut uses = Vec::new();
        for name in &self.uses {
            let specifier = PackageSpecifier::from_str(name)
                .with_context(|| format!("Unable to parse \"{name}\" as a package specifier"))?;
            let pkg = {
                let inner_rt = rt.clone();
                rt.task_manager()
                    .spawn_and_block_on(async move {
                        BinaryPackage::from_registry(&specifier, &*inner_rt).await
                    })
                    .with_context(|| format!("Unable to load \"{name}\""))??
            };
            uses.push(pkg);
        }

        let mut builder = WasiEnv::builder(program_name)
            .runtime(Arc::clone(&rt))
            .args(args)
            .envs(self.env_vars.clone())
            .uses(uses)
            .map_commands(map_commands);

        let mut builder = {
            let mount_fs = RootFileSystemBuilder::new()
                .with_tty(Box::new(DeviceFile::new(__WASI_STDIN_FILENO)))
                .build();
            let (have_current_dir, mapped_dirs) = self.build_mapped_directories(false)?;
            let mut root_layers: Vec<Arc<dyn FileSystem + Send + Sync>> = Vec::new();

            for mapped in mapped_dirs {
                let MountedDirectory { guest, fs } = MountedDirectory::from(mapped);
                if guest == "/" {
                    root_layers.push(fs);
                } else {
                    mount_fs.mount(&guest, Arc::new(fs))?;
                }
            }

            if !root_layers.is_empty() {
                let existing_root = mount_fs
                    .filesystem_at(Path::new("/"))
                    .expect("root fs builder should always mount /");
                mount_fs.set_mount(
                    Path::new("/"),
                    Arc::new(OverlayFileSystem::new(
                        ArcFileSystem::new(existing_root),
                        root_layers,
                    )),
                )?;
            };

            if let Some(cwd) = self.cwd.as_ref() {
                if !cwd.starts_with("/") {
                    bail!("The argument to --cwd must be an absolute path");
                }
                builder = builder.current_dir(cwd.clone());
            }

            // Open the root of the new filesystem
            builder = builder
                .mount_fs(mount_fs)
                .preopen_dir(Path::new("/"))
                .unwrap();

            let dot_path = if have_current_dir {
                PathBuf::from(MAPPED_CURRENT_DIR_DEFAULT_PATH)
            } else {
                PathBuf::from("/")
            };

            builder.add_preopen_build(|p| {
                p.directory(&dot_path)
                    .alias(".")
                    .read(true)
                    .write(true)
                    .create(true)
            })?;

            builder
        };

        *builder.capabilities_mut() = self.capabilities();

        #[cfg(feature = "journal")]
        {
            for trigger in self.snapshot_on.iter().cloned() {
                builder.add_snapshot_trigger(trigger);
            }
            if let Some(interval) = self.snapshot_interval {
                builder.with_snapshot_interval(std::time::Duration::from_millis(interval));
            }
            if self.stop_after_snapshot {
                builder.with_stop_running_after_snapshot(true);
            }
            let (r, w) = self.build_journals()?;
            for journal in r {
                builder.add_read_only_journal(journal);
            }
            for journal in w {
                builder.add_writable_journal(journal);
            }
            builder.with_skip_stdio_during_bootstrap(self.skip_stdio_during_bootstrap);
        }

        Ok(builder)
    }

    #[cfg(feature = "journal")]
    #[allow(clippy::type_complexity)]
    pub fn build_journals(
        &self,
    ) -> anyhow::Result<(Vec<Arc<DynReadableJournal>>, Vec<Arc<DynJournal>>)> {
        let mut readable = Vec::new();
        for journal in self.read_only_journals.clone() {
            if matches!(std::fs::metadata(&journal), Err(e) if e.kind() == std::io::ErrorKind::NotFound)
            {
                bail!("Read-only journal file does not exist: {journal:?}");
            }

            readable
                .push(Arc::new(LogFileJournal::new_readonly(journal)?) as Arc<DynReadableJournal>);
        }

        let mut writable = Vec::new();
        for journal in self.writable_journals.clone() {
            if self.enable_compaction {
                let mut journal = CompactingLogFileJournal::new(journal)?;
                if !self.without_compact_on_drop {
                    journal = journal.with_compact_on_drop()
                }
                if self.with_compact_on_growth.is_normal() && self.with_compact_on_growth != 0f32 {
                    journal = journal.with_compact_on_factor_size(self.with_compact_on_growth);
                }
                writable.push(Arc::new(journal) as Arc<DynJournal>);
            } else {
                writable.push(Arc::new(LogFileJournal::new(journal)?));
            }
        }
        Ok((readable, writable))
    }

    #[cfg(not(feature = "journal"))]
    pub fn build_journals(&self) -> anyhow::Result<Vec<Arc<DynJournal>>> {
        Ok(Vec::new())
    }

    pub fn build_mapped_directories(
        &self,
        is_wasix: bool,
    ) -> Result<(bool, Vec<MappedDirectory>), anyhow::Error> {
        let mut mapped_dirs = Vec::new();

        // Process the --volume flag.
        let mut have_current_dir = false;
        for MappedDirectory { host, guest } in &self.all_volumes() {
            let resolved_host = host.canonicalize().with_context(|| {
                format!(
                    "could not canonicalize path for argument '--volume {}:{}'",
                    host.display(),
                    guest,
                )
            })?;

            if guest == "/" && is_wasix {
                // Note: it appears we canonicalize the path before this point and showing the value of
                // `host` in the error message may throw users off, so we use a placeholder.
                tracing::warn!(
                    "Mounting on the guest's virtual root with --volume <HOST_DIR>:/ breaks WASIX modules' filesystems"
                );
            }

            let mapping = if guest == "." {
                if have_current_dir {
                    bail!(
                        "Cannot pre-open the current directory twice: '--volume=.' must only be specified once"
                    );
                }
                have_current_dir = true;

                let host = if host == Path::new(".") {
                    std::env::current_dir().context("could not determine current directory")?
                } else {
                    host.clone()
                };
                MappedDirectory {
                    host: resolved_host,
                    guest: if is_wasix {
                        MAPPED_CURRENT_DIR_DEFAULT_PATH.to_string()
                    } else {
                        "/".to_string()
                    },
                }
            } else {
                MappedDirectory {
                    host: resolved_host,
                    guest: guest.clone(),
                }
            };
            mapped_dirs.push(mapping);
        }

        Ok((have_current_dir, mapped_dirs))
    }

    pub fn build_mapped_commands(&self) -> Result<Vec<MappedCommand>, anyhow::Error> {
        self.map_commands
            .iter()
            .map(|item| {
                let (a, b) = item.split_once('=').with_context(|| {
                    format!(
                        "Invalid --map-command flag: expected <ALIAS>=<HOST_PATH>, got '{item}'"
                    )
                })?;

                let a = a.trim();
                let b = b.trim();

                if a.is_empty() {
                    bail!("Invalid --map-command flag - alias cannot be empty: '{item}'");
                }
                // TODO(theduke): check if host command exists, and canonicalize PathBuf.
                if b.is_empty() {
                    bail!("Invalid --map-command flag - host path cannot be empty: '{item}'");
                }

                Ok(MappedCommand {
                    alias: a.to_string(),
                    target: b.to_string(),
                })
            })
            .collect::<Result<Vec<_>, anyhow::Error>>()
    }

    pub fn capabilities(&self) -> Capabilities {
        let mut caps = Capabilities::default();

        if self.http_client {
            caps.http_client = wasmer_wasix::http::HttpClientCapabilityV1::new_allow_all();
        }

        if let Some(enable_async_threads) = self.enable_async_threads {
            caps.threading.enable_asynchronous_threading = enable_async_threads;
        }
        caps.threading.enable_exponential_cpu_backoff =
            self.enable_cpu_backoff.map(Duration::from_millis);

        caps
    }

    pub fn prepare_runtime<I>(
        &self,
        engine: Engine,
        env: &WasmerEnv,
        pkg_cache_path: &Path,
        rt_or_handle: I,
        preferred_webc_version: webc::Version,
        compiler_debug_dir_used: bool,
    ) -> Result<impl Runtime + Send + Sync + use<I>>
    where
        I: Into<RuntimeOrHandle>,
    {
        let tokio_task_manager = Arc::new(TokioTaskManager::new(rt_or_handle.into()));
        let mut rt = PluggableRuntime::new(tokio_task_manager.clone());

        let has_networking = self.networking.is_some()
            || capabilities::get_cached_capability(pkg_cache_path)
                .ok()
                .is_some_and(|v| v.enable_networking);

        let ruleset = self
            .networking
            .clone()
            .flatten()
            .map(|ruleset| Ruleset::from_str(&ruleset))
            .transpose()?;

        let network = if let Some(ruleset) = ruleset {
            virtual_net::host::LocalNetworking::with_ruleset(ruleset)
        } else {
            virtual_net::host::LocalNetworking::default()
        };

        if has_networking {
            rt.set_networking_implementation(network);
        } else {
            let net = super::capabilities::net::AskingNetworking::new(
                pkg_cache_path.to_path_buf(),
                Arc::new(network),
            );

            rt.set_networking_implementation(net);
        }

        #[cfg(feature = "journal")]
        {
            let (r, w) = self.build_journals()?;
            for journal in r {
                rt.add_read_only_journal(journal);
            }
            for journal in w {
                rt.add_writable_journal(journal);
            }
        }

        if !self.no_tty {
            let tty = Arc::new(SysTty);
            tty.reset();
            rt.set_tty(tty);
        }

        let client =
            wasmer_wasix::http::default_http_client().context("No HTTP client available")?;
        let client = Arc::new(client);

        let package_loader = self
            .prepare_package_loader(env, client.clone())
            .context("Unable to prepare the package loader")?;

        let registry = self.prepare_source(env, client, preferred_webc_version)?;

        if !self.disable_cache && !compiler_debug_dir_used {
            let cache_dir = env.cache_dir().join("compiled");
            let module_cache = wasmer_wasix::runtime::module_cache::in_memory()
                .with_fallback(FileSystemCache::new(cache_dir, tokio_task_manager));
            rt.set_module_cache(module_cache);
        }

        rt.set_package_loader(package_loader)
            .set_source(registry)
            .set_engine(engine);

        Ok(rt)
    }

    /// Helper function for instantiating a module with Wasi imports for the `Run` command.
    pub fn instantiate(
        &self,
        module: &Module,
        module_hash: ModuleHash,
        program_name: String,
        args: Vec<String>,
        runtime: Arc<dyn Runtime + Send + Sync>,
        store: &mut Store,
    ) -> Result<(WasiFunctionEnv, Instance)> {
        let builder = self.prepare(module, program_name, args, runtime)?;
        let (instance, wasi_env) = builder.instantiate_ext(module.clone(), module_hash, store)?;

        Ok((wasi_env, instance))
    }

    pub fn for_binfmt_interpreter() -> Result<Self> {
        let dir = std::env::var_os("WASMER_BINFMT_MISC_PREOPEN")
            .map(Into::into)
            .unwrap_or_else(|| PathBuf::from("."));
        Ok(Self {
            deny_multiple_wasi_versions: true,
            env_vars: std::env::vars().collect(),
            volumes: vec![MappedDirectory {
                host: dir.clone(),
                guest: dir
                    .to_str()
                    .expect("dir must be a valid string")
                    .to_string(),
            }],
            ..Self::default()
        })
    }

    fn prepare_package_loader(
        &self,
        env: &WasmerEnv,
        client: Arc<dyn HttpClient + Send + Sync>,
    ) -> Result<BuiltinPackageLoader> {
        let checkout_dir = env.cache_dir().join("checkouts");
        let tokens = tokens_by_authority(env)?;

        let loader = BuiltinPackageLoader::new()
            .with_cache_dir(checkout_dir)
            .with_shared_http_client(client)
            .with_tokens(tokens);

        Ok(loader)
    }

    fn prepare_source(
        &self,
        env: &WasmerEnv,
        client: Arc<dyn HttpClient + Send + Sync>,
        preferred_webc_version: webc::Version,
    ) -> Result<MultiSource> {
        let mut source = MultiSource::default();

        // Local packages go first so they override the registry. A directory
        // is served on demand as a local registry; a single file is loaded
        // eagerly under its manifest id.
        let mut preloaded = InMemorySource::new();
        for path in &self.include_webcs {
            if path.is_dir() {
                source.add_source(LocalRegistrySource::new(path)?);
            } else {
                preloaded
                    .add_webc(path)
                    .with_context(|| format!("Unable to load \"{}\"", path.display()))?;
            }
        }
        source.add_source(preloaded);

        // Drop the registry/web sources offline. Their network errors bubble out
        // of the merging MultiSource and abort resolution even when a local
        // source already matched.
        if !self.offline {
            let graphql_endpoint = self.graphql_endpoint(env)?;
            let cache_dir = registry_query_cache_dir(env.cache_dir(), &graphql_endpoint);
            let mut wapm_source = BackendSource::new(graphql_endpoint, Arc::clone(&client))
                .with_local_cache(cache_dir, WAPM_SOURCE_CACHE_TIMEOUT)
                .with_preferred_webc_version(preferred_webc_version);
            if let Some(token) = env
                .config()?
                .registry
                .get_login_token_for_registry(wapm_source.registry_endpoint().as_str())
            {
                wapm_source = wapm_source.with_auth_token(token);
            }
            source.add_source(wapm_source);

            let cache_dir = env.cache_dir().join("downloads");
            source.add_source(WebSource::new(cache_dir, client));
        }

        source.add_source(FileSystemSource::default());

        Ok(source)
    }

    fn graphql_endpoint(&self, env: &WasmerEnv) -> Result<Url> {
        if let Ok(endpoint) = env.registry_endpoint() {
            return Ok(endpoint);
        }

        let config = env.config()?;
        let graphql_endpoint = config.registry.get_graphql_url();
        let graphql_endpoint = graphql_endpoint
            .parse()
            .with_context(|| format!("Unable to parse \"{graphql_endpoint}\" as a URL"))?;

        Ok(graphql_endpoint)
    }
}

fn parse_registry(r: &str) -> Result<Url> {
    UserRegistry::from(r).graphql_endpoint()
}

fn tokens_by_authority(env: &WasmerEnv) -> Result<HashMap<String, String>> {
    let mut tokens = HashMap::new();
    let config = env.config()?;

    for credentials in config.registry.tokens {
        if let Ok(url) = Url::parse(&credentials.registry)
            && url.has_authority()
        {
            tokens.insert(url.authority().to_string(), credentials.token);
        }
    }

    if let (Ok(current_registry), Some(token)) = (env.registry_endpoint(), env.token())
        && current_registry.has_authority()
    {
        tokens.insert(current_registry.authority().to_string(), token);
    }

    // Note: The global wasmer.toml config file stores URLs for the GraphQL
    // endpoint, however that's often on the backend (i.e.
    // https://registry.wasmer.io/graphql) and we also want to use the same API
    // token when sending requests to the frontend (e.g. downloading a package
    // using the `Accept: application/webc` header).
    //
    // As a workaround to avoid needing to query *all* backends to find out
    // their frontend URL every time the `wasmer` CLI runs, we'll assume that
    // when a backend is called something like `registry.wasmer.io`, the
    // frontend will be at `wasmer.io`. This works everywhere except for people
    // developing the backend locally... Sorry, Ayush.

    let mut frontend_tokens = HashMap::new();
    for (hostname, token) in &tokens {
        if let Some(frontend_url) = hostname.strip_prefix("registry.")
            && !tokens.contains_key(frontend_url)
        {
            frontend_tokens.insert(frontend_url.to_string(), token.clone());
        }
    }
    tokens.extend(frontend_tokens);

    Ok(tokens)
}