1use std::collections::BTreeMap;
45
46use serde::{Deserialize, Serialize};
47
48use crate::crypto::SigningKey;
49use crate::dsse::{self, DsseEnvelope};
50use crate::types::AuthorId;
51use crate::{AionError, Result};
52
53pub const AIBOM_PAYLOAD_TYPE: &str = "application/vnd.aion.aibom.v1+json";
55
56pub const AIBOM_SCHEMA_VERSION: &str = "aion.aibom.v1";
58
59#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
61pub struct ModelRef {
62 pub name: String,
64 pub version: String,
66 pub hash_algorithm: String,
68 #[serde(with = "hex_bytes32")]
70 pub hash: [u8; 32],
71 pub size: u64,
73 pub format: String,
75}
76
77#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
79pub struct FrameworkRef {
80 pub name: String,
82 pub version: String,
84 pub cpe: Option<String>,
86}
87
88#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
90pub struct DatasetRef {
91 pub name: String,
93 pub hash_algorithm: Option<String>,
95 #[serde(with = "hex_bytes32_opt")]
97 pub hash: Option<[u8; 32]>,
98 pub size: Option<u64>,
100 pub uri: Option<String>,
102 pub license_spdx_id: Option<String>,
104}
105
106#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
108pub enum LicenseScope {
109 Weights,
111 SourceCode,
113 TrainingData,
115 Documentation,
117 Combined,
119}
120
121#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
123pub struct License {
124 pub spdx_id: String,
126 pub scope: LicenseScope,
128 pub text_uri: Option<String>,
130}
131
132#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
134pub struct SafetyAttestation {
135 pub name: String,
137 pub result: String,
139 pub report_hash_algorithm: Option<String>,
141 #[serde(with = "hex_bytes32_opt")]
143 pub report_hash: Option<[u8; 32]>,
144 pub report_uri: Option<String>,
146}
147
148#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
150pub struct ExportControl {
151 pub regime: String,
153 pub classification: String,
156 pub notes: Option<String>,
158}
159
160#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
162pub struct ExternalReference {
163 pub kind: String,
165 pub uri: String,
167}
168
169#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
171pub struct AiBom {
172 pub schema_version: String,
175 pub model: ModelRef,
177 pub frameworks: Vec<FrameworkRef>,
179 pub datasets: Vec<DatasetRef>,
181 pub licenses: Vec<License>,
183 pub hyperparameters: BTreeMap<String, serde_json::Value>,
185 pub safety_attestations: Vec<SafetyAttestation>,
187 pub export_controls: Vec<ExportControl>,
189 pub references: Vec<ExternalReference>,
191 pub created_at_version: u64,
195}
196
197impl AiBom {
198 #[must_use]
201 pub const fn builder(model: ModelRef, created_at_version: u64) -> AiBomBuilder {
202 AiBomBuilder {
203 model,
204 frameworks: Vec::new(),
205 datasets: Vec::new(),
206 licenses: Vec::new(),
207 hyperparameters: BTreeMap::new(),
208 safety_attestations: Vec::new(),
209 export_controls: Vec::new(),
210 references: Vec::new(),
211 created_at_version,
212 }
213 }
214
215 pub fn to_json(&self) -> Result<String> {
222 serde_json::to_string(self).map_err(|e| AionError::InvalidFormat {
223 reason: format!("AIBOM JSON serialize failed: {e}"),
224 })
225 }
226
227 pub fn from_json(s: &str) -> Result<Self> {
233 serde_json::from_str(s).map_err(|e| AionError::InvalidFormat {
234 reason: format!("AIBOM JSON parse failed: {e}"),
235 })
236 }
237
238 pub fn canonical_bytes(&self) -> Result<Vec<u8>> {
246 serde_json::to_vec(self).map_err(|e| AionError::InvalidFormat {
247 reason: format!("AIBOM canonical bytes failed: {e}"),
248 })
249 }
250
251 pub fn to_jcs_bytes(&self) -> Result<Vec<u8>> {
260 crate::jcs::to_jcs_bytes(self)
261 }
262}
263
264#[derive(Debug)]
266pub struct AiBomBuilder {
267 model: ModelRef,
268 frameworks: Vec<FrameworkRef>,
269 datasets: Vec<DatasetRef>,
270 licenses: Vec<License>,
271 hyperparameters: BTreeMap<String, serde_json::Value>,
272 safety_attestations: Vec<SafetyAttestation>,
273 export_controls: Vec<ExportControl>,
274 references: Vec<ExternalReference>,
275 created_at_version: u64,
276}
277
278impl AiBomBuilder {
279 pub fn add_framework(&mut self, f: FrameworkRef) -> &mut Self {
281 self.frameworks.push(f);
282 self
283 }
284
285 pub fn add_dataset(&mut self, d: DatasetRef) -> &mut Self {
287 self.datasets.push(d);
288 self
289 }
290
291 pub fn add_license(&mut self, l: License) -> &mut Self {
293 self.licenses.push(l);
294 self
295 }
296
297 pub fn hyperparameter(&mut self, k: impl Into<String>, v: serde_json::Value) -> &mut Self {
300 self.hyperparameters.insert(k.into(), v);
301 self
302 }
303
304 pub fn add_safety_attestation(&mut self, s: SafetyAttestation) -> &mut Self {
306 self.safety_attestations.push(s);
307 self
308 }
309
310 pub fn add_export_control(&mut self, e: ExportControl) -> &mut Self {
312 self.export_controls.push(e);
313 self
314 }
315
316 pub fn add_reference(&mut self, r: ExternalReference) -> &mut Self {
318 self.references.push(r);
319 self
320 }
321
322 #[must_use]
324 pub fn build(self) -> AiBom {
325 AiBom {
326 schema_version: AIBOM_SCHEMA_VERSION.to_string(),
327 model: self.model,
328 frameworks: self.frameworks,
329 datasets: self.datasets,
330 licenses: self.licenses,
331 hyperparameters: self.hyperparameters,
332 safety_attestations: self.safety_attestations,
333 export_controls: self.export_controls,
334 references: self.references,
335 created_at_version: self.created_at_version,
336 }
337 }
338}
339
340pub fn wrap_aibom_dsse(aibom: &AiBom, signer: AuthorId, key: &SigningKey) -> Result<DsseEnvelope> {
346 let payload = aibom.canonical_bytes()?;
347 Ok(dsse::sign_envelope(
348 &payload,
349 AIBOM_PAYLOAD_TYPE,
350 signer,
351 key,
352 ))
353}
354
355pub fn unwrap_aibom_dsse(envelope: &DsseEnvelope) -> Result<AiBom> {
364 if envelope.payload_type != AIBOM_PAYLOAD_TYPE {
365 return Err(AionError::InvalidFormat {
366 reason: format!(
367 "envelope payloadType is '{}', expected '{}'",
368 envelope.payload_type, AIBOM_PAYLOAD_TYPE
369 ),
370 });
371 }
372 let payload_str =
373 std::str::from_utf8(&envelope.payload).map_err(|e| AionError::InvalidFormat {
374 reason: format!("AIBOM DSSE payload is not valid UTF-8: {e}"),
375 })?;
376 AiBom::from_json(payload_str)
377}
378
379mod hex_bytes32 {
381 use serde::{Deserialize, Deserializer, Serializer};
382
383 pub fn serialize<S: Serializer>(bytes: &[u8; 32], serializer: S) -> Result<S::Ok, S::Error> {
384 serializer.serialize_str(&hex::encode(bytes))
385 }
386
387 pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<[u8; 32], D::Error> {
388 let s = String::deserialize(deserializer)?;
389 let v = hex::decode(&s).map_err(serde::de::Error::custom)?;
390 if v.len() != 32 {
391 return Err(serde::de::Error::custom(format!(
392 "hash hex length is {} (expected 64 chars = 32 bytes)",
393 v.len()
394 )));
395 }
396 let mut out = [0u8; 32];
397 out.copy_from_slice(&v);
398 Ok(out)
399 }
400}
401
402mod hex_bytes32_opt {
404 use serde::{Deserialize, Deserializer, Serializer};
405
406 pub fn serialize<S: Serializer>(
407 bytes: &Option<[u8; 32]>,
408 serializer: S,
409 ) -> Result<S::Ok, S::Error> {
410 match bytes {
411 Some(b) => serializer.serialize_str(&hex::encode(b)),
412 None => serializer.serialize_none(),
413 }
414 }
415
416 pub fn deserialize<'de, D: Deserializer<'de>>(
417 deserializer: D,
418 ) -> Result<Option<[u8; 32]>, D::Error> {
419 let maybe: Option<String> = Option::deserialize(deserializer)?;
420 match maybe {
421 None => Ok(None),
422 Some(s) => {
423 let v = hex::decode(&s).map_err(serde::de::Error::custom)?;
424 if v.len() != 32 {
425 return Err(serde::de::Error::custom(format!(
426 "hash hex length is {} (expected 64 chars = 32 bytes)",
427 v.len()
428 )));
429 }
430 let mut out = [0u8; 32];
431 out.copy_from_slice(&v);
432 Ok(Some(out))
433 }
434 }
435 }
436}
437
438#[cfg(test)]
439#[allow(clippy::unwrap_used)]
440mod tests {
441 use super::*;
442 use crate::dsse::verify_envelope;
443 use crate::key_registry::KeyRegistry;
444 use serde_json::json;
445
446 fn reg_pinning(signer: AuthorId, key: &SigningKey) -> KeyRegistry {
448 let mut reg = KeyRegistry::new();
449 let master = SigningKey::generate();
450 reg.register_author(signer, master.verifying_key(), key.verifying_key(), 0)
451 .unwrap();
452 reg
453 }
454
455 fn sample_model() -> ModelRef {
456 ModelRef {
457 name: "acme-7b-chat".to_string(),
458 version: "0.3.1".to_string(),
459 hash_algorithm: "BLAKE3-256".to_string(),
460 hash: [0xABu8; 32],
461 size: 1_000,
462 format: "safetensors".to_string(),
463 }
464 }
465
466 fn sample_aibom() -> AiBom {
467 let mut b = AiBom::builder(sample_model(), 42);
468 b.add_framework(FrameworkRef {
469 name: "pytorch".to_string(),
470 version: "2.3.1".to_string(),
471 cpe: None,
472 });
473 b.add_dataset(DatasetRef {
474 name: "c4-en-v2".to_string(),
475 hash_algorithm: Some("BLAKE3-256".to_string()),
476 hash: Some([0xCDu8; 32]),
477 size: None,
478 uri: Some("s3://acme-datasets/c4-en-v2/".to_string()),
479 license_spdx_id: Some("ODC-By-1.0".to_string()),
480 });
481 b.add_license(License {
482 spdx_id: "Apache-2.0".to_string(),
483 scope: LicenseScope::Weights,
484 text_uri: None,
485 });
486 b.hyperparameter("context_length", json!(8192));
487 b.add_export_control(ExportControl {
488 regime: "US-ECCN".to_string(),
489 classification: "EAR99".to_string(),
490 notes: None,
491 });
492 b.build()
493 }
494
495 #[test]
496 fn builds_with_schema_version() {
497 let aibom = sample_aibom();
498 assert_eq!(aibom.schema_version, AIBOM_SCHEMA_VERSION);
499 }
500
501 #[test]
502 fn json_round_trip_preserves_fields() {
503 let aibom = sample_aibom();
504 let json = aibom.to_json().unwrap();
505 let parsed = AiBom::from_json(&json).unwrap();
506 assert_eq!(parsed, aibom);
507 }
508
509 #[test]
510 fn canonical_bytes_are_deterministic() {
511 let aibom = sample_aibom();
512 let a = aibom.canonical_bytes().unwrap();
513 let b = aibom.canonical_bytes().unwrap();
514 assert_eq!(a, b);
515 }
516
517 #[test]
518 fn dsse_wrap_and_verify_round_trip() {
519 let aibom = sample_aibom();
520 let signer = AuthorId::new(1001);
521 let key = SigningKey::generate();
522 let env = wrap_aibom_dsse(&aibom, signer, &key).unwrap();
523 assert_eq!(env.payload_type, AIBOM_PAYLOAD_TYPE);
524 let reg = reg_pinning(signer, &key);
525 let verified = verify_envelope(&env, ®, 1).unwrap();
526 assert_eq!(verified.len(), 1);
527 let back = unwrap_aibom_dsse(&env).unwrap();
528 assert_eq!(back, aibom);
529 }
530
531 #[test]
532 fn unwrap_rejects_wrong_payload_type() {
533 let key = SigningKey::generate();
534 let env = dsse::sign_envelope(b"not aibom", "text/plain", AuthorId::new(1), &key);
535 assert!(unwrap_aibom_dsse(&env).is_err());
536 }
537
538 #[test]
539 fn hash_field_survives_hex_round_trip() {
540 let aibom = sample_aibom();
541 let json = aibom.to_json().unwrap();
542 let parsed = AiBom::from_json(&json).unwrap();
543 assert_eq!(parsed.model.hash, [0xABu8; 32]);
544 }
545
546 mod properties {
547 use super::*;
548 use hegel::generators as gs;
549
550 fn draw_model(tc: &hegel::TestCase) -> ModelRef {
551 let hash_v = tc.draw(gs::binary().min_size(32).max_size(32));
552 let mut hash = [0u8; 32];
553 hash.copy_from_slice(&hash_v);
554 ModelRef {
555 name: tc.draw(gs::text().max_size(32)),
556 version: tc.draw(gs::text().max_size(16)),
557 hash_algorithm: "BLAKE3-256".to_string(),
558 hash,
559 size: tc.draw(gs::integers::<u64>()),
560 format: tc.draw(gs::text().max_size(16)),
561 }
562 }
563
564 fn draw_aibom(tc: &hegel::TestCase) -> AiBom {
565 let mut b = AiBom::builder(draw_model(tc), tc.draw(gs::integers::<u64>()));
566 let n_frameworks = tc.draw(gs::integers::<usize>().max_value(3));
567 for _ in 0..n_frameworks {
568 b.add_framework(FrameworkRef {
569 name: tc.draw(gs::text().max_size(16)),
570 version: tc.draw(gs::text().max_size(8)),
571 cpe: None,
572 });
573 }
574 let n_datasets = tc.draw(gs::integers::<usize>().max_value(3));
575 for _ in 0..n_datasets {
576 b.add_dataset(DatasetRef {
577 name: tc.draw(gs::text().max_size(16)),
578 hash_algorithm: None,
579 hash: None,
580 size: None,
581 uri: None,
582 license_spdx_id: None,
583 });
584 }
585 b.build()
586 }
587
588 #[hegel::test]
589 fn prop_aibom_json_roundtrip(tc: hegel::TestCase) {
590 let aibom = draw_aibom(&tc);
591 let json = aibom.to_json().unwrap_or_else(|_| std::process::abort());
592 let parsed = AiBom::from_json(&json).unwrap_or_else(|_| std::process::abort());
593 assert_eq!(parsed, aibom);
594 }
595
596 #[hegel::test]
597 fn prop_aibom_canonical_bytes_deterministic(tc: hegel::TestCase) {
598 let aibom = draw_aibom(&tc);
599 let a = aibom
600 .canonical_bytes()
601 .unwrap_or_else(|_| std::process::abort());
602 let b = aibom
603 .canonical_bytes()
604 .unwrap_or_else(|_| std::process::abort());
605 assert_eq!(a, b);
606 }
607
608 #[hegel::test]
609 fn prop_aibom_model_hash_survives_json(tc: hegel::TestCase) {
610 let aibom = draw_aibom(&tc);
611 let expected = aibom.model.hash;
612 let json = aibom.to_json().unwrap_or_else(|_| std::process::abort());
613 let parsed = AiBom::from_json(&json).unwrap_or_else(|_| std::process::abort());
614 assert_eq!(parsed.model.hash, expected);
615 }
616
617 #[hegel::test]
618 fn prop_aibom_dsse_roundtrip(tc: hegel::TestCase) {
619 let aibom = draw_aibom(&tc);
620 let signer = AuthorId::new(tc.draw(gs::integers::<u64>().min_value(1)));
621 let key = SigningKey::generate();
622 let env =
623 wrap_aibom_dsse(&aibom, signer, &key).unwrap_or_else(|_| std::process::abort());
624 let reg = reg_pinning(signer, &key);
625 let verified = verify_envelope(&env, ®, 1).unwrap_or_else(|_| std::process::abort());
626 assert_eq!(verified.len(), 1);
627 let back = unwrap_aibom_dsse(&env).unwrap_or_else(|_| std::process::abort());
628 assert_eq!(back, aibom);
629 }
630
631 #[hegel::test]
632 fn prop_aibom_tampered_json_rejects(tc: hegel::TestCase) {
633 let aibom = draw_aibom(&tc);
634 let signer = AuthorId::new(tc.draw(gs::integers::<u64>().min_value(1)));
635 let key = SigningKey::generate();
636 let mut env =
637 wrap_aibom_dsse(&aibom, signer, &key).unwrap_or_else(|_| std::process::abort());
638 let max_idx = env.payload.len().saturating_sub(1);
639 let idx = tc.draw(gs::integers::<usize>().max_value(max_idx));
640 if let Some(b) = env.payload.get_mut(idx) {
641 *b ^= 0x01;
642 }
643 let reg = reg_pinning(signer, &key);
644 assert!(verify_envelope(&env, ®, 1).is_err());
645 }
646
647 #[hegel::test]
648 fn prop_aibom_multi_signer_envelope(tc: hegel::TestCase) {
649 let aibom = draw_aibom(&tc);
650 let s1 = (
651 AuthorId::new(tc.draw(gs::integers::<u64>().min_value(1).max_value(1 << 20))),
652 SigningKey::generate(),
653 );
654 let s2 = (
655 AuthorId::new(s1.0.as_u64().saturating_add(1)),
656 SigningKey::generate(),
657 );
658 let mut env =
659 wrap_aibom_dsse(&aibom, s1.0, &s1.1).unwrap_or_else(|_| std::process::abort());
660 dsse::add_signature(&mut env, s2.0, &s2.1);
661 let mut reg = KeyRegistry::new();
662 for (signer, key) in [(s1.0, &s1.1), (s2.0, &s2.1)] {
663 let master = SigningKey::generate();
664 reg.register_author(signer, master.verifying_key(), key.verifying_key(), 0)
665 .unwrap_or_else(|_| std::process::abort());
666 }
667 let verified = verify_envelope(&env, ®, 1).unwrap_or_else(|_| std::process::abort());
668 assert_eq!(verified.len(), 2);
669 }
670
671 #[hegel::test]
672 fn prop_aibom_payload_type_is_aion_aibom(tc: hegel::TestCase) {
673 let aibom = draw_aibom(&tc);
674 let signer = AuthorId::new(tc.draw(gs::integers::<u64>().min_value(1)));
675 let key = SigningKey::generate();
676 let env =
677 wrap_aibom_dsse(&aibom, signer, &key).unwrap_or_else(|_| std::process::abort());
678 assert_eq!(env.payload_type, AIBOM_PAYLOAD_TYPE);
679 }
680
681 #[hegel::test]
682 fn prop_aibom_to_jcs_bytes_matches_helper(tc: hegel::TestCase) {
683 let aibom = draw_aibom(&tc);
684 let from_method = aibom
685 .to_jcs_bytes()
686 .unwrap_or_else(|_| std::process::abort());
687 let from_helper =
688 crate::jcs::to_jcs_bytes(&aibom).unwrap_or_else(|_| std::process::abort());
689 assert_eq!(from_method, from_helper);
690 }
691 }
692}