wasmer_wasix/runtime/package_loader/

load_package_tree.rs

use std::{
    collections::{BTreeMap, HashMap, HashSet},
    path::{Path, PathBuf},
    sync::Arc,
};

use anyhow::{Context, Error};
use futures::{StreamExt, TryStreamExt};
use once_cell::sync::OnceCell;
use petgraph::visit::EdgeRef;
use virtual_fs::{FileSystem, WebcVolumeFileSystem};
use wasmer_config::package::PackageId;
use wasmer_package::utils::wasm_annotations_to_features;
use webc::metadata::annotations::Atom as AtomAnnotation;
use webc::{Container, Volume};

use crate::{
    bin_factory::{BinaryPackage, BinaryPackageCommand, BinaryPackageMount, BinaryPackageMounts},
    runtime::{
        package_loader::PackageLoader,
        resolver::{
            DependencyGraph, ItemLocation, PackageSummary, Resolution, ResolvedFileSystemMapping,
            ResolvedPackage,
        },
    },
};

use super::to_module_hash;

/// Convert WebAssembly feature annotations to a Features object
fn wasm_annotation_to_features(
    wasm_annotation: &webc::metadata::annotations::Wasm,
) -> Option<wasmer_types::Features> {
    Some(wasm_annotations_to_features(&wasm_annotation.features))
}

/// Extract WebAssembly features from atom metadata if available
fn extract_features_from_atom_metadata(
    atom_metadata: &webc::metadata::Atom,
) -> Option<wasmer_types::Features> {
    if let Ok(Some(wasm_annotation)) = atom_metadata
        .annotation::<webc::metadata::annotations::Wasm>(webc::metadata::annotations::Wasm::KEY)
    {
        wasm_annotation_to_features(&wasm_annotation)
    } else {
        None
    }
}

/// The maximum number of packages that will be loaded in parallel.
const MAX_PARALLEL_DOWNLOADS: usize = 32;

/// Given a fully resolved package, load it into memory for execution.
#[tracing::instrument(level = "debug", skip_all)]
pub async fn load_package_tree(
    root: &Container,
    loader: &dyn PackageLoader,
    resolution: &Resolution,
    root_is_local_dir: bool,
) -> Result<BinaryPackage, Error> {
    let mut containers = fetch_dependencies(loader, &resolution.package, &resolution.graph).await?;
    containers.insert(resolution.package.root_package.clone(), root.clone());
    let package_ids = containers.keys().cloned().collect();
    let fs_opt = filesystem(&containers, &resolution.package, root_is_local_dir)?;

    let root = &resolution.package.root_package;
    let commands = commands(&resolution.package.commands, &containers, resolution)?;

    let file_system_memory_footprint = if let Some(fs) = &fs_opt {
        count_package_mounts(fs)
    } else {
        0
    };

    let loaded = BinaryPackage {
        id: root.clone(),
        package_ids,
        when_cached: crate::syscalls::platform_clock_time_get(
            wasmer_wasix_types::wasi::Snapshot0Clockid::Monotonic,
            1_000_000,
        )
        .ok()
        .map(|ts| ts as u128),
        hash: OnceCell::new(),
        entrypoint_cmd: resolution.package.entrypoint.clone(),
        package_mounts: fs_opt.map(Arc::new),
        commands,
        uses: Vec::new(),
        file_system_memory_footprint,

        additional_host_mapped_directories: vec![],
    };

    Ok(loaded)
}

fn commands(
    commands: &BTreeMap<String, ItemLocation>,
    containers: &HashMap<PackageId, Container>,
    resolution: &Resolution,
) -> Result<Vec<BinaryPackageCommand>, Error> {
    let mut pkg_commands = Vec::new();

    for (
        name,
        ItemLocation {
            name: original_name,
            package,
        },
    ) in commands
    {
        let webc = &containers[package];
        let manifest = webc.manifest();
        let command_metadata = &manifest.commands[original_name];

        if let Some(cmd) =
            load_binary_command(package, name, command_metadata, containers, resolution)?
        {
            pkg_commands.push(cmd);
        }
    }

    Ok(pkg_commands)
}

/// Given a [`webc::metadata::Command`], figure out which atom it uses and load
/// that atom into a [`BinaryPackageCommand`].
#[tracing::instrument(skip_all, fields(%package_id, %command_name))]
fn load_binary_command(
    package_id: &PackageId,
    command_name: &str,
    cmd: &webc::metadata::Command,
    containers: &HashMap<PackageId, Container>,
    resolution: &Resolution,
) -> Result<Option<BinaryPackageCommand>, anyhow::Error> {
    let AtomAnnotation {
        name: atom_name,
        dependency,
        ..
    } = match atom_name_for_command(command_name, cmd)? {
        Some(name) => name,
        None => {
            tracing::warn!(
                cmd.name=command_name,
                cmd.runner=%cmd.runner,
                "Skipping unsupported command",
            );
            return Ok(None);
        }
    };

    let package = &containers[package_id];

    let (webc, resolved_package_id) = match dependency {
        Some(dep) => {
            let ix = resolution
                .graph
                .packages()
                .get(package_id)
                .copied()
                .unwrap();
            let graph = resolution.graph.graph();
            let edge_reference = graph
                .edges_directed(ix, petgraph::Direction::Outgoing)
                .find(|edge| edge.weight().alias == dep)
                .with_context(|| format!("Unable to find the \"{dep}\" dependency for the \"{command_name}\" command in \"{package_id}\""))?;

            let other_package = graph.node_weight(edge_reference.target()).unwrap();
            let id = &other_package.id;

            tracing::debug!(
                dependency=%dep,
                resolved_package_id=%id,
                "command atom resolution: resolved dependency",
            );
            (&containers[id], id)
        }
        None => (package, package_id),
    };

    let atom = webc.get_atom(&atom_name);

    if atom.is_none() && cmd.annotations.is_empty() {
        tracing::info!("applying legacy atom hack");
        return legacy_atom_hack(webc, package_id, command_name, cmd);
    }

    let hash = to_module_hash(webc.manifest().atom_signature(&atom_name)?);

    let atom = atom.with_context(|| {

        let available_atoms = webc.atoms().keys().map(|x| x.as_str()).collect::<Vec<_>>().join(",");

        tracing::warn!(
            %atom_name,
            %resolved_package_id,
            %available_atoms,
            "invalid command: could not find atom in package",
        );

        format!(
            "The '{command_name}' command uses the '{atom_name}' atom, but it isn't present in the package: {resolved_package_id})"
        )
    })?;

    // Get WebAssembly features from manifest atom annotations
    let features = if let Some(atom_metadata) = webc.manifest().atoms.get(&atom_name) {
        extract_features_from_atom_metadata(atom_metadata)
    } else {
        None
    };

    let cmd = BinaryPackageCommand::new(
        command_name.to_string(),
        cmd.clone(),
        atom,
        hash,
        features,
        package_id.clone(),
        resolved_package_id.clone(),
    );

    Ok(Some(cmd))
}

fn atom_name_for_command(
    command_name: &str,
    cmd: &webc::metadata::Command,
) -> Result<Option<AtomAnnotation>, anyhow::Error> {
    use webc::metadata::annotations::{WASI_RUNNER_URI, WCGI_RUNNER_URI};

    if let Some(atom) = cmd
        .atom()
        .context("Unable to deserialize atom annotations")?
    {
        return Ok(Some(atom));
    }

    if [WASI_RUNNER_URI, WCGI_RUNNER_URI]
        .iter()
        .any(|uri| cmd.runner.starts_with(uri))
    {
        // Note: We use the command name as the atom name as a special case
        // for known runner types because sometimes people will construct
        // a manifest by hand instead of using wapm2pirita.
        tracing::debug!(
            command = command_name,
            "No annotations specifying the atom name found. Falling back to the command name"
        );
        return Ok(Some(AtomAnnotation::new(command_name, None)));
    }

    Ok(None)
}

/// HACK: Some older packages like `sharrattj/bash` and `sharrattj/coreutils`
/// contain commands with no annotations. When this happens, you can just assume
/// it wants to use the first atom in the WEBC file.
///
/// That works because most of these packages only have a single atom (e.g. in
/// `sharrattj/coreutils` there are commands for `ls`, `pwd`, and so on, but
/// under the hood they all use the `coreutils` atom).
///
/// See <https://github.com/wasmerio/wasmer/commit/258903140680716da1431d92bced67d486865aeb>
/// for more.
fn legacy_atom_hack(
    webc: &Container,
    package_id: &PackageId,
    command_name: &str,
    metadata: &webc::metadata::Command,
) -> Result<Option<BinaryPackageCommand>, anyhow::Error> {
    let (name, atom) = webc
        .atoms()
        .into_iter()
        .next()
        .ok_or_else(|| anyhow::Error::msg("container does not have any atom"))?;

    tracing::debug!(
        command_name,
        atom.name = name.as_str(),
        atom.len = atom.len(),
        "(hack) The command metadata is malformed. Falling back to the first atom in the WEBC file",
    );

    let hash = to_module_hash(webc.manifest().atom_signature(&name)?);

    // Get WebAssembly features from manifest atom annotations
    let features = if let Some(atom_metadata) = webc.manifest().atoms.get(&name) {
        extract_features_from_atom_metadata(atom_metadata)
    } else {
        None
    };

    Ok(Some(BinaryPackageCommand::new(
        command_name.to_string(),
        metadata.clone(),
        atom,
        hash,
        features,
        package_id.clone(),
        package_id.clone(),
    )))
}

async fn fetch_dependencies(
    loader: &dyn PackageLoader,
    pkg: &ResolvedPackage,
    graph: &DependencyGraph,
) -> Result<HashMap<PackageId, Container>, Error> {
    let mut packages = HashSet::new();

    for loc in pkg.commands.values() {
        packages.insert(loc.package.clone());
    }

    for mapping in &pkg.filesystem {
        packages.insert(mapping.package.clone());
    }

    // We don't need to download the root package
    packages.remove(&pkg.root_package);

    let packages = packages.into_iter().filter_map(|id| {
        let crate::runtime::resolver::Node { pkg, dist, .. } = &graph[&id];
        let summary = PackageSummary {
            pkg: pkg.clone(),
            dist: dist.clone()?,
        };
        Some((id, summary))
    });
    let packages: HashMap<PackageId, Container> = futures::stream::iter(packages)
        .map(|(id, s)| async move {
            match loader.load(&s).await {
                Ok(webc) => Ok((id, webc)),
                Err(e) => Err(e),
            }
        })
        .buffer_unordered(MAX_PARALLEL_DOWNLOADS)
        .try_collect()
        .await?;

    Ok(packages)
}

/// How many bytes worth of files does a directory contain?
fn count_file_system(fs: &dyn FileSystem, path: &Path) -> u64 {
    let mut total = 0;

    let dir = match fs.read_dir(path) {
        Ok(d) => d,
        Err(_err) => {
            return 0;
        }
    };

    for entry in dir.flatten() {
        if let Ok(meta) = entry.metadata() {
            total += meta.len();
            if meta.is_dir() {
                total += count_file_system(fs, entry.path.as_path());
            }
        }
    }

    total
}

fn count_package_mounts(mounts: &BinaryPackageMounts) -> u64 {
    let mut total = 0;

    if let Some(root_layer) = &mounts.root_layer {
        total += count_file_system(root_layer.as_ref(), Path::new("/"));
    }

    for mount in &mounts.mounts {
        total += count_file_system(mount.fs.as_ref(), Path::new("/"));
    }

    total
}

/// Given a set of [`ResolvedFileSystemMapping`]s and the [`Container`] for each
/// package in a dependency tree, construct the resulting filesystem.
///
/// Returns `Ok(None)` if no filesystem mappings were specified.
fn filesystem(
    packages: &HashMap<PackageId, Container>,
    pkg: &ResolvedPackage,
    root_is_local_dir: bool,
) -> Result<Option<BinaryPackageMounts>, Error> {
    if pkg.filesystem.is_empty() {
        return Ok(None);
    }

    let mut found_v2 = None;
    let mut found_v3 = None;

    for ResolvedFileSystemMapping { package, .. } in &pkg.filesystem {
        let container = packages.get(package).with_context(|| {
            format!(
                "\"{}\" wants to use the \"{}\" package, but it isn't in the dependency tree",
                pkg.root_package, package,
            )
        })?;

        match container.version() {
            webc::Version::V1 => {
                anyhow::bail!(
                    "the package '{package}' is a webc v1 package, but webc v1 support was removed"
                );
            }
            webc::Version::V2 => {
                if found_v2.is_none() {
                    found_v2 = Some(package.clone());
                }
            }
            webc::Version::V3 => {
                if found_v3.is_none() {
                    found_v3 = Some(package.clone());
                }
            }
            other => {
                anyhow::bail!("the package '{package}' has an unknown webc version: {other}");
            }
        }
    }

    match (found_v2, found_v3) {
        (None, Some(_)) => filesystem_v3(packages, pkg, root_is_local_dir).map(Some),
        (Some(_), None) => filesystem_v2(packages, pkg, root_is_local_dir).map(Some),
        (Some(v2), Some(v3)) => {
            anyhow::bail!(
                "Mix of webc v2 and v3 in the same dependency tree is not supported; v2: {v2}, v3: {v3}"
            )
        }
        (None, None) => anyhow::bail!("Internal error: no packages found in tree"),
    }
}

/// Build the filesystem for webc v3 packages.
fn filesystem_v3(
    packages: &HashMap<PackageId, Container>,
    pkg: &ResolvedPackage,
    root_is_local_dir: bool,
) -> Result<BinaryPackageMounts, Error> {
    let mut volumes: HashMap<&PackageId, BTreeMap<String, Volume>> = HashMap::new();
    let mut root_layer = None;
    let mut mounts = Vec::new();

    for ResolvedFileSystemMapping {
        mount_path,
        volume_name,
        package,
        ..
    } in &pkg.filesystem
    {
        if *package == pkg.root_package && root_is_local_dir {
            continue;
        }

        if mount_path.as_path() == Path::new("/") {
            tracing::warn!(
                "The \"{package}\" package wants to mount a volume at \"/\", which breaks WASIX modules' filesystems",
            );
        }

        // Note: We want to reuse existing Volume instances if we can. That way
        // we can keep the memory usage down. A webc::compat::Volume is
        // reference-counted, anyway.
        // looks like we need to insert it
        let container = packages.get(package).with_context(|| {
            format!(
                "\"{}\" wants to use the \"{}\" package, but it isn't in the dependency tree",
                pkg.root_package, package,
            )
        })?;
        let container_volumes = match volumes.entry(package) {
            std::collections::hash_map::Entry::Occupied(entry) => &*entry.into_mut(),
            std::collections::hash_map::Entry::Vacant(entry) => &*entry.insert(container.volumes()),
        };

        let volume = container_volumes.get(volume_name).with_context(|| {
            format!("The \"{package}\" package doesn't have a \"{volume_name}\" volume")
        })?;

        let webc_vol = WebcVolumeFileSystem::new(volume.clone());
        if mount_path.as_path() == Path::new("/") {
            root_layer = Some(Arc::new(webc_vol) as Arc<dyn FileSystem + Send + Sync>);
        } else {
            mounts.push(BinaryPackageMount {
                guest_path: mount_path.clone(),
                fs: Arc::new(webc_vol),
                source_path: PathBuf::from("/"),
            });
        }
    }

    Ok(BinaryPackageMounts { root_layer, mounts })
}

/// Build the filesystem for webc v2 packages.
///
// # Note to future readers
//
// Sooo... this code is a bit convoluted because we're constrained by the
// filesystem implementations we've got available.
//
// Ideally, we would create a WebcVolumeFileSystem for each volume we're
// using, then we'd have a single mount filesystem which lets you mount
// filesystem objects under various paths.
//
// The OverlayFileSystem lets us make files from multiple filesystem
// implementations available at the same time, however all of the
// filesystems will be mounted at "/", when the user wants to mount volumes
// at arbitrary locations.
//
// The TmpFileSystem *does* allow mounting at non-root paths, however it can't
// handle nested paths (e.g. mounting to "/lib" and "/lib/python3.10" - see
// <https://github.com/wasmerio/wasmer/issues/3678> for more) and you aren't
// allowed to mount to "/" because it's a special directory that already
// exists.
//
// As a result, we'll duct-tape things together and hope for the best 🤞
fn filesystem_v2(
    packages: &HashMap<PackageId, Container>,
    pkg: &ResolvedPackage,
    root_is_local_dir: bool,
) -> Result<BinaryPackageMounts, Error> {
    let mut volumes: HashMap<&PackageId, BTreeMap<String, Volume>> = HashMap::new();
    let mut root_layer = None;
    let mut mounts = Vec::new();

    for ResolvedFileSystemMapping {
        mount_path,
        volume_name,
        package,
        original_path,
    } in &pkg.filesystem
    {
        if *package == pkg.root_package && root_is_local_dir {
            continue;
        }

        if mount_path.as_path() == Path::new("/") {
            tracing::warn!(
                "The \"{package}\" package wants to mount a volume at \"/\", which breaks WASIX modules' filesystems",
            );
        }

        // Note: We want to reuse existing Volume instances if we can. That way
        // we can keep the memory usage down. A webc::compat::Volume is
        // reference-counted, anyway.
        let container_volumes = match volumes.entry(package) {
            std::collections::hash_map::Entry::Occupied(entry) => &*entry.into_mut(),
            std::collections::hash_map::Entry::Vacant(entry) => {
                // looks like we need to insert it
                let container = packages.get(package)
                    .with_context(|| format!(
                        "\"{}\" wants to use the \"{}\" package, but it isn't in the dependency tree",
                        pkg.root_package,
                        package,
                    ))?;
                &*entry.insert(container.volumes())
            }
        };

        let volume = container_volumes.get(volume_name).with_context(|| {
            format!("The \"{package}\" package doesn't have a \"{volume_name}\" volume")
        })?;

        let mounted_fs = Arc::new(WebcVolumeFileSystem::new(volume.clone()))
            as Arc<dyn FileSystem + Send + Sync>;
        let source_path = original_path
            .as_deref()
            .map(PathBuf::from)
            .unwrap_or_else(|| PathBuf::from("/"));

        if mount_path.as_path() == Path::new("/") {
            root_layer = Some(mounted_fs);
        } else {
            mounts.push(BinaryPackageMount {
                guest_path: mount_path.clone(),
                fs: mounted_fs,
                source_path,
            });
        }
    }

    Ok(BinaryPackageMounts { root_layer, mounts })
}

#[cfg(test)]
mod tests {
    use std::{
        collections::{BTreeMap, HashMap},
        path::{Path, PathBuf},
    };

    use ciborium::value::Value;
    use virtual_fs::FileSystem;
    use wasmer_config::package::PackageId;
    use webc::{
        Container,
        indexmap::IndexMap,
        metadata::{
            Manifest,
            annotations::{FileSystemMapping, FileSystemMappings},
        },
        v2::{
            SignatureAlgorithm,
            read::OwnedReader,
            write::{DirEntry, Directory, FileEntry, Writer},
        },
    };

    use super::{ResolvedFileSystemMapping, ResolvedPackage, filesystem_v2};

    #[test]
    fn v2_filesystem_mapping_resolves_mount_paths() {
        // Regression test: v2 fs mounts are already relative to the mount root,
        // so stripping the mount path again breaks lookups like /public.
        let mut manifest = Manifest::default();
        let fs = FileSystemMappings(vec![FileSystemMapping {
            from: None,
            volume_name: "atom".to_string(),
            host_path: Some("/public".to_string()),
            mount_path: "/public".to_string(),
        }]);
        let mut package = IndexMap::new();
        package.insert(
            FileSystemMappings::KEY.to_string(),
            Value::serialized(&fs).unwrap(),
        );
        manifest.package = package;

        let mut public_children = BTreeMap::new();
        public_children.insert(
            "index.html".parse().unwrap(),
            DirEntry::File(FileEntry::from(b"ok".as_slice())),
        );
        let public_mount_dir = Directory {
            children: public_children,
        };
        let mut root_children = BTreeMap::new();
        root_children.insert("public".parse().unwrap(), DirEntry::Dir(public_mount_dir));
        let atom_dir = Directory {
            children: root_children,
        };

        let writer = Writer::default().write_manifest(&manifest).unwrap();
        let writer = writer.write_atoms(BTreeMap::new()).unwrap();
        let writer = writer.with_volume("atom", atom_dir).unwrap();
        let bytes = writer.finish(SignatureAlgorithm::None).unwrap();

        let reader = OwnedReader::parse(bytes).unwrap();
        let container = Container::from(reader);

        let pkg_id = PackageId::new_named("ns/pkg", "0.1.0".parse().unwrap());
        let mut packages = HashMap::new();
        packages.insert(pkg_id.clone(), container);

        let pkg = ResolvedPackage {
            root_package: pkg_id.clone(),
            commands: BTreeMap::new(),
            entrypoint: None,
            filesystem: vec![ResolvedFileSystemMapping {
                mount_path: PathBuf::from("/public"),
                volume_name: "atom".to_string(),
                original_path: Some("/public".to_string()),
                package: pkg_id,
            }],
        };

        let mounts = filesystem_v2(&packages, &pkg, false).unwrap();
        let mount_fs = mounts.to_mount_fs().unwrap();
        assert!(mount_fs.metadata(Path::new("/public")).unwrap().is_dir());
        assert!(
            mount_fs
                .metadata(Path::new("/public/index.html"))
                .unwrap()
                .is_file()
        );
    }

    #[test]
    fn v2_filesystem_mapping_preserves_root_and_nested_mounts() {
        let mut manifest = Manifest::default();
        let fs = FileSystemMappings(vec![
            FileSystemMapping {
                from: None,
                volume_name: "root".to_string(),
                host_path: Some("/".to_string()),
                mount_path: "/".to_string(),
            },
            FileSystemMapping {
                from: None,
                volume_name: "public".to_string(),
                host_path: Some("/public".to_string()),
                mount_path: "/public".to_string(),
            },
        ]);
        let mut package = IndexMap::new();
        package.insert(
            FileSystemMappings::KEY.to_string(),
            Value::serialized(&fs).unwrap(),
        );
        manifest.package = package;

        let mut root_children = BTreeMap::new();
        root_children.insert(
            "root.txt".parse().unwrap(),
            DirEntry::File(FileEntry::from(b"root".as_slice())),
        );
        let root_dir = Directory {
            children: root_children,
        };

        let mut public_children = BTreeMap::new();
        public_children.insert(
            "index.html".parse().unwrap(),
            DirEntry::File(FileEntry::from(b"ok".as_slice())),
        );
        let public_dir = Directory {
            children: public_children,
        };

        let writer = Writer::default().write_manifest(&manifest).unwrap();
        let writer = writer.write_atoms(BTreeMap::new()).unwrap();
        let writer = writer.with_volume("root", root_dir).unwrap();
        let writer = writer.with_volume("public", public_dir).unwrap();
        let bytes = writer.finish(SignatureAlgorithm::None).unwrap();

        let reader = OwnedReader::parse(bytes).unwrap();
        let container = Container::from(reader);

        let pkg_id = PackageId::new_named("ns/pkg", "0.1.0".parse().unwrap());
        let mut packages = HashMap::new();
        packages.insert(pkg_id.clone(), container);

        let pkg = ResolvedPackage {
            root_package: pkg_id.clone(),
            commands: BTreeMap::new(),
            entrypoint: None,
            filesystem: vec![
                ResolvedFileSystemMapping {
                    mount_path: PathBuf::from("/"),
                    volume_name: "root".to_string(),
                    original_path: Some("/".to_string()),
                    package: pkg_id.clone(),
                },
                ResolvedFileSystemMapping {
                    mount_path: PathBuf::from("/public"),
                    volume_name: "public".to_string(),
                    original_path: Some("/public".to_string()),
                    package: pkg_id,
                },
            ],
        };

        let mounts = filesystem_v2(&packages, &pkg, false).unwrap();
        let root_layer = mounts
            .root_layer
            .as_ref()
            .expect("expected root layer mount");
        assert!(
            root_layer
                .metadata(Path::new("/root.txt"))
                .unwrap()
                .is_file()
        );
        assert_eq!(mounts.mounts.len(), 1);
        assert_eq!(mounts.mounts[0].guest_path, Path::new("/public"));
    }
}