1use std::{collections::BTreeMap, fmt::Write as _, sync::Arc};
19
20use serde::{Deserialize, Deserializer, Serialize};
21use sha2::{Digest, Sha256};
22use thiserror::Error;
23
24pub const HARD_MAX_PAYLOAD_STRATEGY_ID_BYTES: usize = 128;
26pub const HARD_MAX_PAYLOAD_ARTIFACT_BYTES: u32 = 64 * 1024;
28pub const DEFAULT_MAX_PAYLOAD_ARTIFACT_BYTES: u32 = 4 * 1024;
30
31#[derive(Debug, Clone, PartialEq, Eq, Error)]
33#[non_exhaustive]
34pub enum PayloadStrategyError {
35 #[error("payload strategy id must not be empty")]
37 EmptyStrategyId,
38
39 #[error("payload strategy id length {actual} exceeds maximum {maximum}")]
41 StrategyIdTooLong { actual: usize, maximum: usize },
42
43 #[error("payload strategy id must contain only ASCII letters, digits, '.', '-', '_', or ':'")]
45 InvalidStrategyId,
46
47 #[error("payload strategy revision must be greater than zero")]
49 ZeroRevision,
50
51 #[error("payload artifact limit {actual} exceeds maximum {maximum}")]
53 ArtifactLimitTooLarge { actual: u32, maximum: u32 },
54
55 #[error("payload seed length {actual} exceeds limit {maximum}")]
57 SeedTooLarge { actual: usize, maximum: u32 },
58
59 #[error("derived payload length {actual} exceeds limit {maximum}")]
61 ArtifactTooLarge { actual: usize, maximum: u32 },
62
63 #[error("payload strategy {strategy} is already registered")]
65 StrategyIdentityConflict { strategy: PayloadStrategyRef },
66
67 #[error("payload strategy {strategy} is not registered")]
69 UnknownStrategy { strategy: PayloadStrategyRef },
70
71 #[error("payload strategy implementation returned mismatched provenance")]
73 StrategyProvenanceMismatch,
74
75 #[error("payload strategy implementation returned mismatched variant role")]
77 VariantRoleMismatch,
78
79 #[error("payload strategy derivation failed")]
84 DerivationFailed,
85}
86
87#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
89pub struct PayloadStrategyRef {
90 id: String,
91 revision: u32,
92}
93
94impl PayloadStrategyRef {
95 pub fn new(id: impl Into<String>, revision: u32) -> Result<Self, PayloadStrategyError> {
97 let id = id.into();
98 if id.trim().is_empty() {
99 return Err(PayloadStrategyError::EmptyStrategyId);
100 }
101 if id.len() > HARD_MAX_PAYLOAD_STRATEGY_ID_BYTES {
102 return Err(PayloadStrategyError::StrategyIdTooLong {
103 actual: id.len(),
104 maximum: HARD_MAX_PAYLOAD_STRATEGY_ID_BYTES,
105 });
106 }
107 if !id
108 .bytes()
109 .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'-' | b'_' | b':'))
110 {
111 return Err(PayloadStrategyError::InvalidStrategyId);
112 }
113 if revision == 0 {
114 return Err(PayloadStrategyError::ZeroRevision);
115 }
116 Ok(Self { id, revision })
117 }
118
119 pub fn id(&self) -> &str {
121 &self.id
122 }
123
124 pub const fn revision(&self) -> u32 {
126 self.revision
127 }
128}
129
130impl std::fmt::Debug for PayloadStrategyRef {
131 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
132 formatter
133 .debug_struct("PayloadStrategyRef")
134 .field("id", &self.id)
135 .field("revision", &self.revision)
136 .finish()
137 }
138}
139
140impl std::fmt::Display for PayloadStrategyRef {
141 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
142 write!(formatter, "{}@{}", self.id, self.revision)
143 }
144}
145
146impl<'de> Deserialize<'de> for PayloadStrategyRef {
147 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
148 where
149 D: Deserializer<'de>,
150 {
151 #[derive(Deserialize)]
152 #[serde(deny_unknown_fields)]
153 struct WireRef {
154 id: String,
155 revision: u32,
156 }
157
158 let wire = WireRef::deserialize(deserializer)?;
159 Self::new(wire.id, wire.revision).map_err(serde::de::Error::custom)
160 }
161}
162
163#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
165#[serde(rename_all = "snake_case")]
166#[non_exhaustive]
167pub enum PayloadVariantRole {
168 Control,
170 Candidate,
172}
173
174#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
176pub struct PayloadStrategyLimits {
177 max_seed_bytes: u32,
178 max_artifact_bytes: u32,
179}
180
181impl PayloadStrategyLimits {
182 pub fn new(max_seed_bytes: u32, max_artifact_bytes: u32) -> Result<Self, PayloadStrategyError> {
184 for actual in [max_seed_bytes, max_artifact_bytes] {
185 if actual > HARD_MAX_PAYLOAD_ARTIFACT_BYTES {
186 return Err(PayloadStrategyError::ArtifactLimitTooLarge {
187 actual,
188 maximum: HARD_MAX_PAYLOAD_ARTIFACT_BYTES,
189 });
190 }
191 }
192 Ok(Self {
193 max_seed_bytes,
194 max_artifact_bytes,
195 })
196 }
197
198 pub const fn max_seed_bytes(self) -> u32 {
200 self.max_seed_bytes
201 }
202
203 pub const fn max_artifact_bytes(self) -> u32 {
205 self.max_artifact_bytes
206 }
207}
208
209impl Default for PayloadStrategyLimits {
210 fn default() -> Self {
211 Self {
212 max_seed_bytes: DEFAULT_MAX_PAYLOAD_ARTIFACT_BYTES,
213 max_artifact_bytes: DEFAULT_MAX_PAYLOAD_ARTIFACT_BYTES,
214 }
215 }
216}
217
218impl<'de> Deserialize<'de> for PayloadStrategyLimits {
219 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
220 where
221 D: Deserializer<'de>,
222 {
223 #[derive(Deserialize)]
224 #[serde(default, deny_unknown_fields)]
225 struct WireLimits {
226 max_seed_bytes: u32,
227 max_artifact_bytes: u32,
228 }
229
230 impl Default for WireLimits {
231 fn default() -> Self {
232 let limits = PayloadStrategyLimits::default();
233 Self {
234 max_seed_bytes: limits.max_seed_bytes,
235 max_artifact_bytes: limits.max_artifact_bytes,
236 }
237 }
238 }
239
240 let wire = WireLimits::deserialize(deserializer)?;
241 Self::new(wire.max_seed_bytes, wire.max_artifact_bytes).map_err(serde::de::Error::custom)
242 }
243}
244
245#[derive(Clone, PartialEq, Eq)]
247pub struct PayloadSeed(Vec<u8>);
248
249impl PayloadSeed {
250 pub fn new(
252 value: impl Into<Vec<u8>>,
253 limits: PayloadStrategyLimits,
254 ) -> Result<Self, PayloadStrategyError> {
255 let value = value.into();
256 if value.len() > usize::try_from(limits.max_seed_bytes).unwrap_or(usize::MAX) {
257 return Err(PayloadStrategyError::SeedTooLarge {
258 actual: value.len(),
259 maximum: limits.max_seed_bytes,
260 });
261 }
262 Ok(Self(value))
263 }
264
265 pub fn as_bytes(&self) -> &[u8] {
267 &self.0
268 }
269
270 pub fn len(&self) -> usize {
272 self.0.len()
273 }
274
275 pub fn is_empty(&self) -> bool {
277 self.0.is_empty()
278 }
279}
280
281impl std::fmt::Debug for PayloadSeed {
282 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
283 formatter
284 .debug_struct("PayloadSeed")
285 .field("bytes", &self.0.len())
286 .field("value", &"<redacted>")
287 .finish()
288 }
289}
290
291#[derive(Clone, PartialEq, Eq)]
293pub struct PayloadArtifact {
294 strategy: PayloadStrategyRef,
295 role: PayloadVariantRole,
296 bytes: Vec<u8>,
297 digest: [u8; 32],
298}
299
300impl PayloadArtifact {
301 pub fn new(
303 strategy: PayloadStrategyRef,
304 role: PayloadVariantRole,
305 bytes: impl Into<Vec<u8>>,
306 limits: PayloadStrategyLimits,
307 ) -> Result<Self, PayloadStrategyError> {
308 let bytes = bytes.into();
309 if bytes.len() > usize::try_from(limits.max_artifact_bytes).unwrap_or(usize::MAX) {
310 return Err(PayloadStrategyError::ArtifactTooLarge {
311 actual: bytes.len(),
312 maximum: limits.max_artifact_bytes,
313 });
314 }
315 let digest = Sha256::digest(&bytes).into();
316 Ok(Self {
317 strategy,
318 role,
319 bytes,
320 digest,
321 })
322 }
323
324 pub const fn strategy(&self) -> &PayloadStrategyRef {
326 &self.strategy
327 }
328
329 pub const fn role(&self) -> PayloadVariantRole {
331 self.role
332 }
333
334 pub fn as_bytes(&self) -> &[u8] {
336 &self.bytes
337 }
338
339 pub fn receipt(&self) -> PayloadArtifactReceipt {
341 PayloadArtifactReceipt {
342 strategy: self.strategy.clone(),
343 role: self.role,
344 byte_length: u32::try_from(self.bytes.len()).unwrap_or(u32::MAX),
345 sha256: encode_digest(self.digest),
346 }
347 }
348}
349
350impl std::fmt::Debug for PayloadArtifact {
351 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
352 formatter
353 .debug_struct("PayloadArtifact")
354 .field("strategy", &self.strategy)
355 .field("role", &self.role)
356 .field("byte_length", &self.bytes.len())
357 .field("sha256", &encode_digest(self.digest))
358 .field("value", &"<redacted>")
359 .finish()
360 }
361}
362
363#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
365pub struct PayloadArtifactReceipt {
366 strategy: PayloadStrategyRef,
367 role: PayloadVariantRole,
368 byte_length: u32,
369 sha256: String,
370}
371
372impl PayloadArtifactReceipt {
373 pub const fn strategy(&self) -> &PayloadStrategyRef {
375 &self.strategy
376 }
377
378 pub const fn role(&self) -> PayloadVariantRole {
380 self.role
381 }
382
383 pub const fn byte_length(&self) -> u32 {
385 self.byte_length
386 }
387
388 pub fn sha256(&self) -> &str {
390 &self.sha256
391 }
392}
393
394pub trait PayloadStrategy: Send + Sync {
401 fn strategy_ref(&self) -> &PayloadStrategyRef;
403
404 fn derive_one(
406 &self,
407 role: PayloadVariantRole,
408 seed: &PayloadSeed,
409 limits: PayloadStrategyLimits,
410 ) -> Result<PayloadArtifact, PayloadStrategyError>;
411}
412
413#[derive(Clone, Default)]
415pub struct PayloadStrategyRegistry {
416 strategies: BTreeMap<PayloadStrategyRef, Arc<dyn PayloadStrategy>>,
417}
418
419impl PayloadStrategyRegistry {
420 pub fn new() -> Self {
422 Self::default()
423 }
424
425 pub fn register(
427 &mut self,
428 strategy: Arc<dyn PayloadStrategy>,
429 ) -> Result<(), PayloadStrategyError> {
430 let reference = strategy.strategy_ref().clone();
431 if self.strategies.contains_key(&reference) {
432 return Err(PayloadStrategyError::StrategyIdentityConflict {
433 strategy: reference,
434 });
435 }
436 self.strategies.insert(reference, strategy);
437 Ok(())
438 }
439
440 pub fn contains(&self, reference: &PayloadStrategyRef) -> bool {
442 self.strategies.contains_key(reference)
443 }
444
445 pub fn len(&self) -> usize {
447 self.strategies.len()
448 }
449
450 pub fn is_empty(&self) -> bool {
452 self.strategies.is_empty()
453 }
454
455 pub fn derive_one(
457 &self,
458 reference: &PayloadStrategyRef,
459 role: PayloadVariantRole,
460 seed: &PayloadSeed,
461 limits: PayloadStrategyLimits,
462 ) -> Result<PayloadArtifact, PayloadStrategyError> {
463 if seed.len() > usize::try_from(limits.max_seed_bytes()).unwrap_or(usize::MAX) {
464 return Err(PayloadStrategyError::SeedTooLarge {
465 actual: seed.len(),
466 maximum: limits.max_seed_bytes(),
467 });
468 }
469 let strategy = self.strategies.get(reference).ok_or_else(|| {
470 PayloadStrategyError::UnknownStrategy {
471 strategy: reference.clone(),
472 }
473 })?;
474 let artifact = strategy.derive_one(role, seed, limits)?;
475 if artifact.strategy() != reference {
476 return Err(PayloadStrategyError::StrategyProvenanceMismatch);
477 }
478 if artifact.role() != role {
479 return Err(PayloadStrategyError::VariantRoleMismatch);
480 }
481 if artifact.as_bytes().len()
482 > usize::try_from(limits.max_artifact_bytes()).unwrap_or(usize::MAX)
483 {
484 return Err(PayloadStrategyError::ArtifactTooLarge {
485 actual: artifact.as_bytes().len(),
486 maximum: limits.max_artifact_bytes(),
487 });
488 }
489 Ok(artifact)
490 }
491}
492
493impl std::fmt::Debug for PayloadStrategyRegistry {
494 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
495 formatter
496 .debug_struct("PayloadStrategyRegistry")
497 .field("strategies", &self.strategies.keys().collect::<Vec<_>>())
498 .finish()
499 }
500}
501
502fn encode_digest(digest: [u8; 32]) -> String {
503 let mut output = String::with_capacity(64);
504 for byte in digest {
505 let _ = write!(output, "{byte:02x}");
506 }
507 output
508}
509
510#[cfg(test)]
511mod tests {
512 use super::*;
513
514 struct PrefixStrategy {
515 reference: PayloadStrategyRef,
516 }
517
518 impl PayloadStrategy for PrefixStrategy {
519 fn strategy_ref(&self) -> &PayloadStrategyRef {
520 &self.reference
521 }
522
523 fn derive_one(
524 &self,
525 role: PayloadVariantRole,
526 seed: &PayloadSeed,
527 limits: PayloadStrategyLimits,
528 ) -> Result<PayloadArtifact, PayloadStrategyError> {
529 let prefix = match role {
530 PayloadVariantRole::Control => b"control:".as_slice(),
531 PayloadVariantRole::Candidate => b"candidate:".as_slice(),
532 };
533 let mut bytes = Vec::with_capacity(prefix.len().saturating_add(seed.len()));
534 bytes.extend_from_slice(prefix);
535 bytes.extend_from_slice(seed.as_bytes());
536 PayloadArtifact::new(self.reference.clone(), role, bytes, limits)
537 }
538 }
539
540 fn reference() -> PayloadStrategyRef {
541 PayloadStrategyRef::new("differential.fixture", 1).unwrap()
542 }
543
544 #[test]
545 fn strategy_reference_and_limits_validate_wire_input() {
546 assert!(matches!(
547 PayloadStrategyRef::new("", 1),
548 Err(PayloadStrategyError::EmptyStrategyId)
549 ));
550 assert!(matches!(
551 PayloadStrategyRef::new("valid", 0),
552 Err(PayloadStrategyError::ZeroRevision)
553 ));
554 for invalid in [" leading", "trailing ", "line\nbreak", "non-ascii-ş"] {
555 assert!(matches!(
556 PayloadStrategyRef::new(invalid, 1),
557 Err(PayloadStrategyError::InvalidStrategyId)
558 ));
559 }
560 assert!(
561 serde_json::from_value::<PayloadStrategyRef>(serde_json::json!({
562 "id": "valid",
563 "revision": 0
564 }))
565 .is_err()
566 );
567 assert!(PayloadStrategyLimits::new(0, 0).is_ok());
568 assert!(PayloadStrategyLimits::new(HARD_MAX_PAYLOAD_ARTIFACT_BYTES + 1, 1).is_err());
569 }
570
571 #[test]
572 fn same_input_produces_same_artifact_and_receipt() {
573 let reference = reference();
574 let strategy = Arc::new(PrefixStrategy {
575 reference: reference.clone(),
576 });
577 let mut registry = PayloadStrategyRegistry::new();
578 registry.register(strategy).unwrap();
579 let limits = PayloadStrategyLimits::default();
580 let seed = PayloadSeed::new(b"fixture".to_vec(), limits).unwrap();
581
582 let first = registry
583 .derive_one(&reference, PayloadVariantRole::Candidate, &seed, limits)
584 .unwrap();
585 let second = registry
586 .derive_one(&reference, PayloadVariantRole::Candidate, &seed, limits)
587 .unwrap();
588
589 assert_eq!(first, second);
590 assert_eq!(first.receipt(), second.receipt());
591 assert_eq!(first.as_bytes(), b"candidate:fixture");
592 }
593
594 #[test]
595 fn raw_seed_and_artifact_never_enter_debug_or_receipt_json() {
596 let secret = b"not-for-logs".to_vec();
597 let limits = PayloadStrategyLimits::default();
598 let seed = PayloadSeed::new(secret.clone(), limits).unwrap();
599 let artifact =
600 PayloadArtifact::new(reference(), PayloadVariantRole::Control, secret, limits).unwrap();
601 let output = format!(
602 "{seed:?} {artifact:?} {}",
603 serde_json::to_string(&artifact.receipt()).unwrap()
604 );
605
606 assert!(!output.contains("not-for-logs"));
607 assert!(output.contains("<redacted>"));
608 assert_eq!(artifact.receipt().sha256().len(), 64);
609 }
610
611 #[test]
612 fn seed_and_derived_output_limits_fail_closed() {
613 let limits = PayloadStrategyLimits::new(3, 8).unwrap();
614 assert!(matches!(
615 PayloadSeed::new(b"four".to_vec(), limits),
616 Err(PayloadStrategyError::SeedTooLarge { .. })
617 ));
618
619 let reference = reference();
620 let mut registry = PayloadStrategyRegistry::new();
621 registry
622 .register(Arc::new(PrefixStrategy {
623 reference: reference.clone(),
624 }))
625 .unwrap();
626 let seed = PayloadSeed::new(b"abc".to_vec(), limits).unwrap();
627 assert!(matches!(
628 registry.derive_one(&reference, PayloadVariantRole::Candidate, &seed, limits),
629 Err(PayloadStrategyError::ArtifactTooLarge { .. })
630 ));
631
632 let wider = PayloadStrategyLimits::new(4, 16).unwrap();
633 let narrower = PayloadStrategyLimits::new(3, 16).unwrap();
634 let seed = PayloadSeed::new(b"four".to_vec(), wider).unwrap();
635 assert!(matches!(
636 registry.derive_one(&reference, PayloadVariantRole::Control, &seed, narrower),
637 Err(PayloadStrategyError::SeedTooLarge {
638 actual: 4,
639 maximum: 3
640 })
641 ));
642 }
643
644 #[test]
645 fn registry_rejects_unknown_and_duplicate_strategy_revisions() {
646 let reference = reference();
647 let strategy = Arc::new(PrefixStrategy {
648 reference: reference.clone(),
649 });
650 let mut registry = PayloadStrategyRegistry::new();
651 registry.register(strategy.clone()).unwrap();
652 assert!(matches!(
653 registry.register(strategy),
654 Err(PayloadStrategyError::StrategyIdentityConflict { .. })
655 ));
656
657 let missing = PayloadStrategyRef::new("missing", 1).unwrap();
658 let seed = PayloadSeed::new(Vec::new(), PayloadStrategyLimits::default()).unwrap();
659 assert!(matches!(
660 registry.derive_one(
661 &missing,
662 PayloadVariantRole::Control,
663 &seed,
664 PayloadStrategyLimits::default()
665 ),
666 Err(PayloadStrategyError::UnknownStrategy { .. })
667 ));
668 }
669}