spin_capabilities/
collect.rs

1use std::collections::BTreeSet;
2
3use wac_graph::types::are_semver_compatible;
4use wasmparser::{Parser, Payload};
5
6use crate::CAPABILITY_SETS;
7
8/// Infer the Spin capability sets a component requires, by inspecting its
9/// top-level (component) imports.
10///
11/// Each import is matched against the known capability sets (`ai_models`,
12/// `allowed_outbound_hosts`, etc.); every set that matches contributes its name
13/// to the result. The returned set is deduplicated and sorted. An empty set
14/// means the component imports nothing that maps to a Spin capability.
15pub fn required_capabilities(source: &[u8]) -> anyhow::Result<BTreeSet<&'static str>> {
16    let mut capabilities = BTreeSet::new();
17    let mut depth: u32 = 0;
18
19    for payload in Parser::new(0).parse_all(source) {
20        match payload? {
21            Payload::ModuleSection { .. } | Payload::ComponentSection { .. } => {
22                depth += 1;
23            }
24            Payload::End(_) if depth > 0 => {
25                depth -= 1;
26            }
27            Payload::ComponentImportSection(reader) if depth == 0 => {
28                for import in reader {
29                    let name = import?.name.name;
30                    for &(capability, set) in CAPABILITY_SETS {
31                        if set.iter().any(|s| are_semver_compatible(name, s)) {
32                            capabilities.insert(capability);
33                        }
34                    }
35                }
36            }
37            _ => {}
38        }
39    }
40
41    Ok(capabilities)
42}
43
44#[cfg(test)]
45mod tests {
46    use super::*;
47
48    /// Build a minimal Wasm component that imports the given interface names.
49    fn build_component(import_names: &[&str]) -> Vec<u8> {
50        use wasm_encoder::{
51            Component, ComponentImportSection, ComponentTypeRef, ComponentTypeSection, InstanceType,
52        };
53
54        let mut component = Component::new();
55
56        // Define one empty instance type to reference from all imports.
57        let mut types = ComponentTypeSection::new();
58        types.instance(&InstanceType::new());
59        component.section(&types);
60
61        let mut imports = ComponentImportSection::new();
62        for name in import_names {
63            imports.import(*name, ComponentTypeRef::Instance(0));
64        }
65        component.section(&imports);
66
67        component.finish()
68    }
69
70    #[test]
71    fn no_matching_imports_returns_empty() {
72        let bytes = build_component(&["some:unknown/interface@1.0.0"]);
73        assert!(required_capabilities(&bytes).unwrap().is_empty());
74    }
75
76    #[test]
77    fn empty_component_returns_empty() {
78        let bytes = wasm_encoder::Component::new().finish();
79        assert!(required_capabilities(&bytes).unwrap().is_empty());
80    }
81
82    #[test]
83    fn single_ai_models_import() {
84        let bytes = build_component(&["fermyon:spin/llm@2.0.0"]);
85        assert_eq!(
86            required_capabilities(&bytes).unwrap(),
87            BTreeSet::from(["ai_models"])
88        );
89    }
90
91    #[test]
92    fn single_allowed_outbound_hosts_import() {
93        let bytes = build_component(&["wasi:http/outgoing-handler@0.2.6"]);
94        assert_eq!(
95            required_capabilities(&bytes).unwrap(),
96            BTreeSet::from(["allowed_outbound_hosts"])
97        );
98    }
99
100    #[test]
101    fn multiple_capabilities_deduped() {
102        let bytes = build_component(&[
103            "fermyon:spin/llm@2.0.0",
104            "wasi:http/outgoing-handler@0.2.6",
105            "wasi:sockets/tcp@0.2.6",
106            "fermyon:spin/variables@2.0.0",
107        ]);
108        assert_eq!(
109            required_capabilities(&bytes).unwrap(),
110            BTreeSet::from(["ai_models", "allowed_outbound_hosts", "variables"])
111        );
112    }
113
114    #[test]
115    fn all_capability_sets_detected() {
116        let bytes = build_component(&[
117            "fermyon:spin/llm@2.0.0",           // ai_models
118            "wasi:http/outgoing-handler@0.2.6", // allowed_outbound_hosts
119            "wasi:cli/environment@0.2.6",       // environment
120            "wasi:filesystem/preopens@0.2.6",   // files
121            "fermyon:spin/key-value@2.0.0",     // key_value_stores
122            "fermyon:spin/sqlite@2.0.0",        // sqlite_databases
123            "fermyon:spin/variables@2.0.0",     // variables
124        ]);
125        assert_eq!(
126            required_capabilities(&bytes).unwrap(),
127            BTreeSet::from([
128                "ai_models",
129                "allowed_outbound_hosts",
130                "environment",
131                "files",
132                "key_value_stores",
133                "sqlite_databases",
134                "variables",
135            ])
136        );
137    }
138
139    #[test]
140    fn duplicate_set_entries_are_deduped() {
141        let bytes =
142            build_component(&["wasi:http/outgoing-handler@0.2.6", "wasi:sockets/tcp@0.2.6"]);
143        // Both map to allowed_outbound_hosts — should appear once.
144        assert_eq!(
145            required_capabilities(&bytes).unwrap(),
146            BTreeSet::from(["allowed_outbound_hosts"])
147        );
148    }
149
150    #[test]
151    fn mixed_known_and_unknown_imports() {
152        let bytes = build_component(&[
153            "fermyon:spin/llm@2.0.0",
154            "some:unknown/thing@1.0.0",
155            "wasi:cli/environment@0.2.6",
156        ]);
157        assert_eq!(
158            required_capabilities(&bytes).unwrap(),
159            BTreeSet::from(["ai_models", "environment"])
160        );
161    }
162}
spin_capabilities/collect.rs

spin_capabilities/
collect.rs
