a3s_code_core/capability/
set.rs

1use std::collections::BTreeMap;
2use std::io::Write;
3use std::sync::Arc;
4
5use serde::Serialize;
6use sha2::{Digest, Sha256};
7
8use super::{
9    CapabilityContribution, CapabilityDescriptor, CapabilityId, CapabilityKind, CapabilitySetError,
10    CapabilitySource, CapabilitySourceClass, CapabilitySourceId, CodeCatalogGeneration,
11    Sha256Digest,
12};
13
14pub const CAPABILITY_SET_SCHEMA: &str = "a3s.code.capability-set.v1";
15pub const CAPABILITY_SET_DIGEST_DOMAIN: &str = "a3s.code.capability-set.v1";
16pub const MAX_CAPABILITY_SOURCES: usize = 1_024;
17pub const MAX_CAPABILITIES: usize = 4_096;
18pub const MAX_CAPABILITY_DEPENDENCY_EDGES: usize = 32_768;
19pub const MAX_CAPABILITY_CANONICAL_BYTES: u64 = 16 * 1024 * 1024;
20
21const CAPABILITY_DIGEST_PREFIX: &[u8] = b"a3s-code-capability\0";
22
23/// Canonically ordered immutable capability identity set.
24///
25/// Construction returns an [`Arc`] so readers retain the exact generation
26/// without consulting a mutable latest-value registry.
27#[derive(Debug)]
28pub struct CapabilitySet {
29    generation: CodeCatalogGeneration,
30    digest: Sha256Digest,
31    sources: BTreeMap<CapabilitySourceId, CapabilitySource>,
32    descriptors: BTreeMap<CapabilityId, CapabilityDescriptor>,
33    use_generation: Option<super::UseCapabilityGeneration>,
34}
35
36impl CapabilitySet {
37    pub const fn schema(&self) -> &'static str {
38        CAPABILITY_SET_SCHEMA
39    }
40
41    pub fn empty() -> Result<Arc<Self>, CapabilitySetError> {
42        Self::from_contributions(
43            CodeCatalogGeneration::INITIAL,
44            Vec::<CapabilityContribution>::new(),
45        )
46    }
47
48    pub fn from_contributions(
49        generation: CodeCatalogGeneration,
50        contributions: impl IntoIterator<Item = CapabilityContribution>,
51    ) -> Result<Arc<Self>, CapabilitySetError> {
52        Self::build(generation, None, contributions)
53    }
54
55    /// Freeze a complete product projection while retaining its upstream Use
56    /// cursor even when product filtering yields no package descriptors.
57    pub fn from_use_projection(
58        generation: CodeCatalogGeneration,
59        use_generation: super::UseCapabilityGeneration,
60        contributions: impl IntoIterator<Item = CapabilityContribution>,
61    ) -> Result<Arc<Self>, CapabilitySetError> {
62        Self::build(generation, Some(use_generation), contributions)
63    }
64
65    fn build(
66        generation: CodeCatalogGeneration,
67        expected_use_generation: Option<super::UseCapabilityGeneration>,
68        contributions: impl IntoIterator<Item = CapabilityContribution>,
69    ) -> Result<Arc<Self>, CapabilitySetError> {
70        let mut sources = BTreeMap::new();
71        let mut descriptors = BTreeMap::new();
72        let mut public_names = BTreeMap::<(CapabilityKind, Box<str>), CapabilitySourceClass>::new();
73        let mut dependency_edges = 0_usize;
74        let mut use_generation = expected_use_generation;
75
76        for contribution in contributions {
77            if sources.len() >= MAX_CAPABILITY_SOURCES {
78                return Err(CapabilitySetError::BoundExceeded {
79                    field: "sources",
80                    max: MAX_CAPABILITY_SOURCES,
81                });
82            }
83            let (source, contributed) = contribution.into_parts();
84            let source_id = source.id().clone();
85            if sources.contains_key(&source_id) {
86                return Err(CapabilitySetError::DuplicateSource {
87                    source_id: source_id.to_string(),
88                });
89            }
90            if let Some(observed) = source.use_capability_generation() {
91                match &use_generation {
92                    Some(expected) if expected != observed => {
93                        return Err(CapabilitySetError::MixedUseGeneration {
94                            expected_generation: expected.generation(),
95                            actual_generation: observed.generation(),
96                            revision_mismatch: expected.revision() != observed.revision(),
97                            registry_revision_mismatch: expected.registry_revision()
98                                != observed.registry_revision(),
99                        });
100                    }
101                    None => use_generation = Some(observed.clone()),
102                    Some(_) => {}
103                }
104            }
105
106            for (id, descriptor) in contributed {
107                if descriptors.len() >= MAX_CAPABILITIES {
108                    return Err(CapabilitySetError::BoundExceeded {
109                        field: "capabilities",
110                        max: MAX_CAPABILITIES,
111                    });
112                }
113                dependency_edges = dependency_edges
114                    .checked_add(descriptor.dependencies().len())
115                    .ok_or(CapabilitySetError::BoundExceeded {
116                        field: "dependency_edges",
117                        max: MAX_CAPABILITY_DEPENDENCY_EDGES,
118                    })?;
119                if dependency_edges > MAX_CAPABILITY_DEPENDENCY_EDGES {
120                    return Err(CapabilitySetError::BoundExceeded {
121                        field: "dependency_edges",
122                        max: MAX_CAPABILITY_DEPENDENCY_EDGES,
123                    });
124                }
125                let public_key = (
126                    id.kind(),
127                    descriptor.public_name().to_owned().into_boxed_str(),
128                );
129                if let Some(existing_class) = public_names.get(&public_key) {
130                    let error = if *existing_class == CapabilitySourceClass::BuiltIn
131                        || source.class() == CapabilitySourceClass::BuiltIn
132                    {
133                        CapabilitySetError::BuiltinShadow {
134                            kind: id.kind(),
135                            public_name: descriptor.public_name().to_owned(),
136                        }
137                    } else {
138                        CapabilitySetError::PublicNameConflict {
139                            kind: id.kind(),
140                            public_name: descriptor.public_name().to_owned(),
141                        }
142                    };
143                    return Err(error);
144                }
145                public_names.insert(public_key, source.class());
146                if descriptors.insert(id.clone(), descriptor).is_some() {
147                    return Err(CapabilitySetError::DuplicateCapability {
148                        capability: id.to_string(),
149                    });
150                }
151            }
152            sources.insert(source_id, source);
153        }
154
155        for (id, descriptor) in &descriptors {
156            for dependency in descriptor.dependencies() {
157                if !descriptors.contains_key(dependency) {
158                    return Err(CapabilitySetError::MissingDependency {
159                        capability: id.to_string(),
160                        dependency: dependency.to_string(),
161                    });
162                }
163            }
164        }
165
166        let digest = canonical_digest(generation, use_generation.as_ref(), &sources, &descriptors)?;
167        Ok(Arc::new(Self {
168            generation,
169            digest,
170            sources,
171            descriptors,
172            use_generation,
173        }))
174    }
175
176    pub const fn generation(&self) -> CodeCatalogGeneration {
177        self.generation
178    }
179
180    pub fn digest(&self) -> &Sha256Digest {
181        &self.digest
182    }
183
184    pub fn len(&self) -> usize {
185        self.descriptors.len()
186    }
187
188    pub fn is_empty(&self) -> bool {
189        self.descriptors.is_empty()
190    }
191
192    pub fn source_count(&self) -> usize {
193        self.sources.len()
194    }
195
196    pub fn use_capability_generation(&self) -> Option<&super::UseCapabilityGeneration> {
197        self.use_generation.as_ref()
198    }
199
200    pub fn source(&self, id: &CapabilitySourceId) -> Option<&CapabilitySource> {
201        self.sources.get(id)
202    }
203
204    pub fn get(&self, id: &CapabilityId) -> Option<&CapabilityDescriptor> {
205        self.descriptors.get(id)
206    }
207
208    pub fn contains(&self, id: &CapabilityId) -> bool {
209        self.descriptors.contains_key(id)
210    }
211
212    pub fn sources(
213        &self,
214    ) -> impl ExactSizeIterator<Item = (&CapabilitySourceId, &CapabilitySource)> {
215        self.sources.iter()
216    }
217
218    pub fn iter(&self) -> impl ExactSizeIterator<Item = (&CapabilityId, &CapabilityDescriptor)> {
219        self.descriptors.iter()
220    }
221}
222
223fn canonical_digest(
224    generation: CodeCatalogGeneration,
225    use_generation: Option<&super::UseCapabilityGeneration>,
226    sources: &BTreeMap<CapabilitySourceId, CapabilitySource>,
227    descriptors: &BTreeMap<CapabilityId, CapabilityDescriptor>,
228) -> Result<Sha256Digest, CapabilitySetError> {
229    #[derive(Serialize)]
230    #[serde(rename_all = "camelCase")]
231    struct Identity<'a> {
232        schema: &'static str,
233        generation: CodeCatalogGeneration,
234        use_generation: Option<&'a super::UseCapabilityGeneration>,
235        sources: Vec<&'a CapabilitySource>,
236        capabilities: Vec<&'a CapabilityDescriptor>,
237    }
238
239    let identity = Identity {
240        schema: CAPABILITY_SET_SCHEMA,
241        generation,
242        use_generation,
243        sources: sources.values().collect(),
244        capabilities: descriptors.values().collect(),
245    };
246    let mut writer = DigestWriter::new(CAPABILITY_SET_DIGEST_DOMAIN);
247    let encoded = serde_json::to_writer(&mut writer, &identity);
248    if writer.exceeded {
249        return Err(CapabilitySetError::BoundExceeded {
250            field: "canonical_bytes",
251            max: usize::try_from(MAX_CAPABILITY_CANONICAL_BYTES).unwrap_or(usize::MAX),
252        });
253    }
254    encoded.map_err(|error| CapabilitySetError::CanonicalEncoding(error.to_string()))?;
255    Sha256Digest::new(writer.finish())
256}
257
258struct DigestWriter {
259    hasher: Sha256,
260    bytes: u64,
261    exceeded: bool,
262}
263
264impl DigestWriter {
265    fn new(domain: &str) -> Self {
266        let mut hasher = Sha256::new();
267        hasher.update(CAPABILITY_DIGEST_PREFIX);
268        hasher.update((domain.len() as u64).to_be_bytes());
269        hasher.update(domain.as_bytes());
270        Self {
271            hasher,
272            bytes: 0,
273            exceeded: false,
274        }
275    }
276
277    fn finish(mut self) -> String {
278        self.hasher.update(self.bytes.to_be_bytes());
279        format!("sha256:{:x}", self.hasher.finalize())
280    }
281}
282
283impl Write for DigestWriter {
284    fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
285        let buffer_len = u64::try_from(buffer.len()).unwrap_or(u64::MAX);
286        if self.bytes.saturating_add(buffer_len) > MAX_CAPABILITY_CANONICAL_BYTES {
287            self.exceeded = true;
288            return Err(std::io::Error::other(
289                "canonical capability set exceeds its byte bound",
290            ));
291        }
292        self.hasher.update(buffer);
293        self.bytes += buffer_len;
294        Ok(buffer.len())
295    }
296
297    fn flush(&mut self) -> std::io::Result<()> {
298        Ok(())
299    }
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305
306    fn digest(byte: char) -> Sha256Digest {
307        Sha256Digest::new(format!("sha256:{}", byte.to_string().repeat(64))).unwrap()
308    }
309
310    #[test]
311    fn external_sources_cannot_shadow_a_sealed_builtin() {
312        let builtin = CapabilitySource::builtin("a3s-code", digest('a')).unwrap();
313        let external = CapabilitySource::host("desktop", digest('b')).unwrap();
314        let builtin_read = CapabilityDescriptor::new(
315            &builtin,
316            CapabilityKind::Tool,
317            "read",
318            "read",
319            digest('c'),
320            [],
321        )
322        .unwrap();
323        let external_read = CapabilityDescriptor::new(
324            &external,
325            CapabilityKind::Tool,
326            "workspace-read",
327            "read",
328            digest('d'),
329            [],
330        )
331        .unwrap();
332
333        for contributions in [
334            vec![
335                CapabilityContribution::new(builtin.clone(), vec![builtin_read.clone()]).unwrap(),
336                CapabilityContribution::new(external.clone(), vec![external_read.clone()]).unwrap(),
337            ],
338            vec![
339                CapabilityContribution::new(external.clone(), vec![external_read.clone()]).unwrap(),
340                CapabilityContribution::new(builtin.clone(), vec![builtin_read.clone()]).unwrap(),
341            ],
342        ] {
343            let error =
344                CapabilitySet::from_contributions(CodeCatalogGeneration::new(1), contributions)
345                    .unwrap_err();
346            assert!(matches!(error, CapabilitySetError::BuiltinShadow { .. }));
347        }
348    }
349}
a3s_code_core/capability/set.rs

a3s_code_core/capability/
set.rs
