dcp/capability/
manifest.rs1use crate::{DCPError, SecurityError};
7
8#[repr(C)]
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct CapabilityManifest {
17 pub version: u16,
19 _reserved: u16,
21 _reserved2: u32,
23 pub tools: [u64; 128],
25 pub resources: [u64; 16],
27 pub prompts: [u64; 8],
29 pub extensions: u64,
31 pub signature: [u8; 64],
33}
34
35impl CapabilityManifest {
36 pub const SIZE: usize = 8 + 1024 + 128 + 64 + 8 + 64; pub const MAX_TOOLS: usize = 8192;
41 pub const MAX_RESOURCES: usize = 1024;
43 pub const MAX_PROMPTS: usize = 512;
45
46 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 #[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 Ok(unsafe { &*(bytes.as_ptr() as *const Self) })
68 }
69
70 #[inline(always)]
72 pub fn as_bytes(&self) -> &[u8] {
73 unsafe { std::slice::from_raw_parts(self as *const Self as *const u8, Self::SIZE) }
75 }
76
77 pub fn signed_bytes(&self) -> &[u8] {
79 &self.as_bytes()[..Self::SIZE - 64]
80 }
81
82 #[inline]
86 pub fn intersect(&self, other: &Self) -> Self {
87 let mut result = Self::new(self.version.min(other.version));
88
89 for i in 0..128 {
91 result.tools[i] = self.tools[i] & other.tools[i];
92 }
93
94 for i in 0..16 {
96 result.resources[i] = self.resources[i] & other.resources[i];
97 }
98
99 for i in 0..8 {
101 result.prompts[i] = self.prompts[i] & other.prompts[i];
102 }
103
104 result.extensions = self.extensions & other.extensions;
106
107 result
108 }
109
110 #[inline]
112 pub fn negotiate(client: &Self, server: &Self) -> Self {
113 client.intersect(server)
114 }
115
116 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 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 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 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[inline]
256 pub fn set_extension(&mut self, bit: u8) {
257 if bit < 64 {
258 self.extensions |= 1u64 << bit;
259 }
260 }
261
262 #[inline]
264 pub fn clear_extension(&mut self, bit: u8) {
265 if bit < 64 {
266 self.extensions &= !(1u64 << bit);
267 }
268 }
269
270 #[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 pub fn tool_count(&self) -> u32 {
281 self.tools.iter().map(|w| w.count_ones()).sum()
282 }
283
284 pub fn resource_count(&self) -> u32 {
286 self.resources.iter().map(|w| w.count_ones()).sum()
287 }
288
289 pub fn prompt_count(&self) -> u32 {
291 self.prompts.iter().map(|w| w.count_ones()).sum()
292 }
293
294 pub fn extension_count(&self) -> u32 {
296 self.extensions.count_ones()
297 }
298
299 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 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 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 m1.set_tool(1);
414 m1.set_tool(2);
415 m1.set_tool(3);
416
417 m2.set_tool(2);
419 m2.set_tool(3);
420 m2.set_tool(4);
421
422 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 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 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 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}