Skip to main content

dcp/capability/
manifest.rs

1//! Pre-computed capability manifest for instant capability negotiation.
2//!
3//! The CapabilityManifest uses bitsets to represent supported tools, resources,
4//! and prompts, enabling O(1) capability intersection using bitwise AND operations.
5
6use crate::{DCPError, SecurityError};
7
8/// Pre-computed capability manifest
9///
10/// Uses bitsets for efficient capability representation and intersection.
11/// - tools: 8192 bits (128 × u64) = supports up to 8192 tools
12/// - resources: 1024 bits (16 × u64) = supports up to 1024 resources
13/// - prompts: 512 bits (8 × u64) = supports up to 512 prompts
14#[repr(C)]
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct CapabilityManifest {
17    /// Protocol version
18    pub version: u16,
19    /// Reserved for alignment
20    _reserved: u16,
21    /// Reserved for future use
22    _reserved2: u32,
23    /// Supported tools bitset (8192 bits = 1024 bytes)
24    pub tools: [u64; 128],
25    /// Supported resources bitset (1024 bits = 128 bytes)
26    pub resources: [u64; 16],
27    /// Supported prompts bitset (512 bits = 64 bytes)
28    pub prompts: [u64; 8],
29    /// Extension flags
30    pub extensions: u64,
31    /// Ed25519 signature
32    pub signature: [u8; 64],
33}
34
35impl CapabilityManifest {
36    /// Size of the manifest in bytes
37    pub const SIZE: usize = 8 + 1024 + 128 + 64 + 8 + 64; // 1296 bytes
38
39    /// Maximum number of tools supported
40    pub const MAX_TOOLS: usize = 8192;
41    /// Maximum number of resources supported
42    pub const MAX_RESOURCES: usize = 1024;
43    /// Maximum number of prompts supported
44    pub const MAX_PROMPTS: usize = 512;
45
46    /// Create a new empty capability manifest
47    pub fn new(version: u16) -> Self {
48        Self {
49            version,
50            _reserved: 0,
51            _reserved2: 0,
52            tools: [0u64; 128],
53            resources: [0u64; 16],
54            prompts: [0u64; 8],
55            extensions: 0,
56            signature: [0u8; 64],
57        }
58    }
59
60    /// Parse manifest from bytes
61    #[inline(always)]
62    pub fn from_bytes(bytes: &[u8]) -> Result<&Self, DCPError> {
63        if bytes.len() < Self::SIZE {
64            return Err(DCPError::InsufficientData);
65        }
66        // SAFETY: We've verified the slice is at least SIZE bytes
67        Ok(unsafe { &*(bytes.as_ptr() as *const Self) })
68    }
69
70    /// Serialize manifest to bytes
71    #[inline(always)]
72    pub fn as_bytes(&self) -> &[u8] {
73        // SAFETY: CapabilityManifest is repr(C) with predictable layout
74        unsafe { std::slice::from_raw_parts(self as *const Self as *const u8, Self::SIZE) }
75    }
76
77    /// Get the bytes that are signed (everything before the signature)
78    pub fn signed_bytes(&self) -> &[u8] {
79        &self.as_bytes()[..Self::SIZE - 64]
80    }
81
82    /// Compute capability intersection - single CPU instruction per word
83    ///
84    /// Returns a new manifest containing only capabilities present in both manifests.
85    #[inline]
86    pub fn intersect(&self, other: &Self) -> Self {
87        let mut result = Self::new(self.version.min(other.version));
88
89        // Intersect tools bitset
90        for i in 0..128 {
91            result.tools[i] = self.tools[i] & other.tools[i];
92        }
93
94        // Intersect resources bitset
95        for i in 0..16 {
96            result.resources[i] = self.resources[i] & other.resources[i];
97        }
98
99        // Intersect prompts bitset
100        for i in 0..8 {
101            result.prompts[i] = self.prompts[i] & other.prompts[i];
102        }
103
104        // Intersect extensions
105        result.extensions = self.extensions & other.extensions;
106
107        result
108    }
109
110    /// Negotiate a deny-by-default manifest from client and server offers.
111    #[inline]
112    pub fn negotiate(client: &Self, server: &Self) -> Self {
113        client.intersect(server)
114    }
115
116    /// Require that a negotiated tool capability is present.
117    pub fn require_tool(&self, tool_id: u16) -> Result<(), SecurityError> {
118        if self.has_tool(tool_id) {
119            Ok(())
120        } else {
121            Err(SecurityError::InsufficientCapabilities)
122        }
123    }
124
125    /// Require that a negotiated resource capability is present.
126    pub fn require_resource(&self, resource_id: u16) -> Result<(), SecurityError> {
127        if self.has_resource(resource_id) {
128            Ok(())
129        } else {
130            Err(SecurityError::InsufficientCapabilities)
131        }
132    }
133
134    /// Require that a negotiated prompt capability is present.
135    pub fn require_prompt(&self, prompt_id: u16) -> Result<(), SecurityError> {
136        if self.has_prompt(prompt_id) {
137            Ok(())
138        } else {
139            Err(SecurityError::InsufficientCapabilities)
140        }
141    }
142
143    /// Require that a negotiated extension capability is present.
144    pub fn require_extension(&self, bit: u8) -> Result<(), SecurityError> {
145        if self.has_extension(bit) {
146            Ok(())
147        } else {
148            Err(SecurityError::InsufficientCapabilities)
149        }
150    }
151
152    /// Set a tool capability
153    #[inline]
154    pub fn set_tool(&mut self, tool_id: u16) {
155        let id = tool_id as usize;
156        if id < Self::MAX_TOOLS {
157            let word = id / 64;
158            let bit = id % 64;
159            self.tools[word] |= 1u64 << bit;
160        }
161    }
162
163    /// Clear a tool capability
164    #[inline]
165    pub fn clear_tool(&mut self, tool_id: u16) {
166        let id = tool_id as usize;
167        if id < Self::MAX_TOOLS {
168            let word = id / 64;
169            let bit = id % 64;
170            self.tools[word] &= !(1u64 << bit);
171        }
172    }
173
174    /// Check if a tool is supported
175    #[inline]
176    pub fn has_tool(&self, tool_id: u16) -> bool {
177        let id = tool_id as usize;
178        if id >= Self::MAX_TOOLS {
179            return false;
180        }
181        let word = id / 64;
182        let bit = id % 64;
183        self.tools[word] & (1u64 << bit) != 0
184    }
185
186    /// Set a resource capability
187    #[inline]
188    pub fn set_resource(&mut self, resource_id: u16) {
189        let id = resource_id as usize;
190        if id < Self::MAX_RESOURCES {
191            let word = id / 64;
192            let bit = id % 64;
193            self.resources[word] |= 1u64 << bit;
194        }
195    }
196
197    /// Clear a resource capability
198    #[inline]
199    pub fn clear_resource(&mut self, resource_id: u16) {
200        let id = resource_id as usize;
201        if id < Self::MAX_RESOURCES {
202            let word = id / 64;
203            let bit = id % 64;
204            self.resources[word] &= !(1u64 << bit);
205        }
206    }
207
208    /// Check if a resource is supported
209    #[inline]
210    pub fn has_resource(&self, resource_id: u16) -> bool {
211        let id = resource_id as usize;
212        if id >= Self::MAX_RESOURCES {
213            return false;
214        }
215        let word = id / 64;
216        let bit = id % 64;
217        self.resources[word] & (1u64 << bit) != 0
218    }
219
220    /// Set a prompt capability
221    #[inline]
222    pub fn set_prompt(&mut self, prompt_id: u16) {
223        let id = prompt_id as usize;
224        if id < Self::MAX_PROMPTS {
225            let word = id / 64;
226            let bit = id % 64;
227            self.prompts[word] |= 1u64 << bit;
228        }
229    }
230
231    /// Clear a prompt capability
232    #[inline]
233    pub fn clear_prompt(&mut self, prompt_id: u16) {
234        let id = prompt_id as usize;
235        if id < Self::MAX_PROMPTS {
236            let word = id / 64;
237            let bit = id % 64;
238            self.prompts[word] &= !(1u64 << bit);
239        }
240    }
241
242    /// Check if a prompt is supported
243    #[inline]
244    pub fn has_prompt(&self, prompt_id: u16) -> bool {
245        let id = prompt_id as usize;
246        if id >= Self::MAX_PROMPTS {
247            return false;
248        }
249        let word = id / 64;
250        let bit = id % 64;
251        self.prompts[word] & (1u64 << bit) != 0
252    }
253
254    /// Set an extension flag
255    #[inline]
256    pub fn set_extension(&mut self, bit: u8) {
257        if bit < 64 {
258            self.extensions |= 1u64 << bit;
259        }
260    }
261
262    /// Clear an extension flag
263    #[inline]
264    pub fn clear_extension(&mut self, bit: u8) {
265        if bit < 64 {
266            self.extensions &= !(1u64 << bit);
267        }
268    }
269
270    /// Check if an extension is supported
271    #[inline]
272    pub fn has_extension(&self, bit: u8) -> bool {
273        if bit >= 64 {
274            return false;
275        }
276        self.extensions & (1u64 << bit) != 0
277    }
278
279    /// Count the number of supported tools
280    pub fn tool_count(&self) -> u32 {
281        self.tools.iter().map(|w| w.count_ones()).sum()
282    }
283
284    /// Count the number of supported resources
285    pub fn resource_count(&self) -> u32 {
286        self.resources.iter().map(|w| w.count_ones()).sum()
287    }
288
289    /// Count the number of supported prompts
290    pub fn prompt_count(&self) -> u32 {
291        self.prompts.iter().map(|w| w.count_ones()).sum()
292    }
293
294    /// Count the number of enabled extensions
295    pub fn extension_count(&self) -> u32 {
296        self.extensions.count_ones()
297    }
298
299    /// Get an iterator over all supported tool IDs
300    pub fn tool_ids(&self) -> impl Iterator<Item = u16> + '_ {
301        self.tools.iter().enumerate().flat_map(|(word_idx, &word)| {
302            (0..64).filter_map(move |bit| {
303                if word & (1u64 << bit) != 0 {
304                    Some((word_idx * 64 + bit) as u16)
305                } else {
306                    None
307                }
308            })
309        })
310    }
311
312    /// Get an iterator over all supported resource IDs
313    pub fn resource_ids(&self) -> impl Iterator<Item = u16> + '_ {
314        self.resources
315            .iter()
316            .enumerate()
317            .flat_map(|(word_idx, &word)| {
318                (0..64).filter_map(move |bit| {
319                    if word & (1u64 << bit) != 0 {
320                        Some((word_idx * 64 + bit) as u16)
321                    } else {
322                        None
323                    }
324                })
325            })
326    }
327
328    /// Get an iterator over all supported prompt IDs
329    pub fn prompt_ids(&self) -> impl Iterator<Item = u16> + '_ {
330        self.prompts
331            .iter()
332            .enumerate()
333            .flat_map(|(word_idx, &word)| {
334                (0..64).filter_map(move |bit| {
335                    if word & (1u64 << bit) != 0 {
336                        Some((word_idx * 64 + bit) as u16)
337                    } else {
338                        None
339                    }
340                })
341            })
342    }
343}
344
345impl Default for CapabilityManifest {
346    fn default() -> Self {
347        Self::new(1)
348    }
349}
350
351#[cfg(test)]
352mod tests {
353    use super::*;
354
355    #[test]
356    fn test_manifest_size() {
357        assert_eq!(
358            std::mem::size_of::<CapabilityManifest>(),
359            CapabilityManifest::SIZE
360        );
361    }
362
363    #[test]
364    fn test_tool_operations() {
365        let mut manifest = CapabilityManifest::new(1);
366
367        assert!(!manifest.has_tool(42));
368        manifest.set_tool(42);
369        assert!(manifest.has_tool(42));
370        manifest.clear_tool(42);
371        assert!(!manifest.has_tool(42));
372    }
373
374    #[test]
375    fn test_resource_operations() {
376        let mut manifest = CapabilityManifest::new(1);
377
378        assert!(!manifest.has_resource(100));
379        manifest.set_resource(100);
380        assert!(manifest.has_resource(100));
381        manifest.clear_resource(100);
382        assert!(!manifest.has_resource(100));
383    }
384
385    #[test]
386    fn test_prompt_operations() {
387        let mut manifest = CapabilityManifest::new(1);
388
389        assert!(!manifest.has_prompt(50));
390        manifest.set_prompt(50);
391        assert!(manifest.has_prompt(50));
392        manifest.clear_prompt(50);
393        assert!(!manifest.has_prompt(50));
394    }
395
396    #[test]
397    fn test_extension_operations() {
398        let mut manifest = CapabilityManifest::new(1);
399
400        assert!(!manifest.has_extension(5));
401        manifest.set_extension(5);
402        assert!(manifest.has_extension(5));
403        manifest.clear_extension(5);
404        assert!(!manifest.has_extension(5));
405    }
406
407    #[test]
408    fn test_intersection() {
409        let mut m1 = CapabilityManifest::new(1);
410        let mut m2 = CapabilityManifest::new(2);
411
412        // Set some tools in m1
413        m1.set_tool(1);
414        m1.set_tool(2);
415        m1.set_tool(3);
416
417        // Set some tools in m2
418        m2.set_tool(2);
419        m2.set_tool(3);
420        m2.set_tool(4);
421
422        // Intersection should have only 2 and 3
423        let result = m1.intersect(&m2);
424        assert!(!result.has_tool(1));
425        assert!(result.has_tool(2));
426        assert!(result.has_tool(3));
427        assert!(!result.has_tool(4));
428
429        // Version should be minimum
430        assert_eq!(result.version, 1);
431    }
432
433    #[test]
434    fn test_round_trip() {
435        let mut manifest = CapabilityManifest::new(1);
436        manifest.set_tool(42);
437        manifest.set_tool(100);
438        manifest.set_resource(5);
439        manifest.set_prompt(10);
440        manifest.set_extension(3);
441
442        let bytes = manifest.as_bytes();
443        let parsed = CapabilityManifest::from_bytes(bytes).unwrap();
444
445        assert_eq!(parsed.version, 1);
446        assert!(parsed.has_tool(42));
447        assert!(parsed.has_tool(100));
448        assert!(parsed.has_resource(5));
449        assert!(parsed.has_prompt(10));
450        assert!(parsed.has_extension(3));
451    }
452
453    #[test]
454    fn test_counts() {
455        let mut manifest = CapabilityManifest::new(1);
456        manifest.set_tool(1);
457        manifest.set_tool(2);
458        manifest.set_tool(3);
459        manifest.set_resource(1);
460        manifest.set_resource(2);
461        manifest.set_prompt(1);
462        manifest.set_extension(0);
463        manifest.set_extension(1);
464
465        assert_eq!(manifest.tool_count(), 3);
466        assert_eq!(manifest.resource_count(), 2);
467        assert_eq!(manifest.prompt_count(), 1);
468        assert_eq!(manifest.extension_count(), 2);
469    }
470
471    #[test]
472    fn test_boundary_ids() {
473        let mut manifest = CapabilityManifest::new(1);
474
475        // Test boundary tool IDs
476        manifest.set_tool(0);
477        manifest.set_tool(63);
478        manifest.set_tool(64);
479        manifest.set_tool(8191);
480
481        assert!(manifest.has_tool(0));
482        assert!(manifest.has_tool(63));
483        assert!(manifest.has_tool(64));
484        assert!(manifest.has_tool(8191));
485
486        // Out of range should not panic
487        assert!(!manifest.has_tool(8192));
488    }
489
490    #[test]
491    fn test_iterators() {
492        let mut manifest = CapabilityManifest::new(1);
493        manifest.set_tool(5);
494        manifest.set_tool(100);
495        manifest.set_tool(1000);
496
497        let tool_ids: Vec<_> = manifest.tool_ids().collect();
498        assert_eq!(tool_ids, vec![5, 100, 1000]);
499    }
500}