1use crate::error::{Error, Result};
11use crate::logging::debug;
12use crate::payment::metrics::QuotingMetricsTracker;
13use crate::payment::pricing::calculate_price;
14use crate::storage::lmdb::LmdbStorage;
15use evmlib::merkle_payments::MerklePaymentCandidateNode;
16use evmlib::PaymentQuote;
17use evmlib::RewardsAddress;
18use parking_lot::RwLock;
19use saorsa_core::MlDsa65;
20use saorsa_pqc::pqc::types::MlDsaSecretKey;
21use saorsa_pqc::pqc::MlDsaOperations;
22use std::sync::Arc;
23use std::time::SystemTime;
24
25pub type XorName = [u8; 32];
27
28pub type SignFn = Box<dyn Fn(&[u8]) -> Vec<u8> + Send + Sync>;
30
31pub struct QuoteGenerator {
36 rewards_address: RewardsAddress,
38 metrics_tracker: QuotingMetricsTracker,
44 storage: RwLock<Option<Arc<LmdbStorage>>>,
56 sign_fn: Option<SignFn>,
59 pub_key: Vec<u8>,
61}
62
63impl QuoteGenerator {
64 #[must_use]
73 pub fn new(rewards_address: RewardsAddress, metrics_tracker: QuotingMetricsTracker) -> Self {
74 Self {
75 rewards_address,
76 metrics_tracker,
77 storage: RwLock::new(None),
78 sign_fn: None,
79 pub_key: Vec::new(),
80 }
81 }
82
83 pub fn attach_storage(&self, storage: Arc<LmdbStorage>) {
92 *self.storage.write() = Some(storage);
93 debug!("QuoteGenerator: LmdbStorage attached for current-records pricing");
94 }
95
96 fn pricing_records_stored(&self) -> usize {
104 if let Some(storage) = self.storage.read().as_ref() {
105 match storage.current_chunks() {
106 Ok(n) => return usize::try_from(n).unwrap_or(usize::MAX),
107 Err(e) => {
108 debug!(
109 "QuoteGenerator: current_chunks() failed ({e}); \
110 falling back to metrics_tracker for pricing"
111 );
112 }
113 }
114 }
115 self.metrics_tracker.records_stored()
116 }
117
118 pub fn set_signer<F>(&mut self, pub_key: Vec<u8>, sign_fn: F)
125 where
126 F: Fn(&[u8]) -> Vec<u8> + Send + Sync + 'static,
127 {
128 self.pub_key = pub_key;
129 self.sign_fn = Some(Box::new(sign_fn));
130 }
131
132 #[must_use]
134 pub fn can_sign(&self) -> bool {
135 self.sign_fn.is_some()
136 }
137
138 pub fn probe_signer(&self) -> Result<()> {
144 let sign_fn = self
145 .sign_fn
146 .as_ref()
147 .ok_or_else(|| Error::Payment("Signer not set".to_string()))?;
148 let test_msg = b"ant-signing-probe";
149 let test_sig = sign_fn(test_msg);
150 if test_sig.is_empty() {
151 return Err(Error::Payment(
152 "ML-DSA-65 signing probe failed: empty signature produced".to_string(),
153 ));
154 }
155 Ok(())
156 }
157
158 pub fn create_quote(
174 &self,
175 content: XorName,
176 data_size: usize,
177 data_type: u32,
178 ) -> Result<PaymentQuote> {
179 let sign_fn = self
180 .sign_fn
181 .as_ref()
182 .ok_or_else(|| Error::Payment("Quote signing not configured".to_string()))?;
183
184 let timestamp = SystemTime::now();
185
186 let price = calculate_price(self.pricing_records_stored());
190
191 let xor_name = xor_name::XorName(content);
193
194 let bytes =
196 PaymentQuote::bytes_for_signing(xor_name, timestamp, &price, &self.rewards_address);
197
198 let signature = sign_fn(&bytes);
200 if signature.is_empty() {
201 return Err(Error::Payment(
202 "Signing produced empty signature".to_string(),
203 ));
204 }
205
206 let quote = PaymentQuote {
207 content: xor_name,
208 timestamp,
209 price,
210 pub_key: self.pub_key.clone(),
211 rewards_address: self.rewards_address,
212 signature,
213 };
214
215 if crate::logging::enabled!(crate::logging::Level::DEBUG) {
216 let content_hex = hex::encode(content);
217 debug!("Generated quote for {content_hex} (size: {data_size}, type: {data_type})");
218 }
219
220 Ok(quote)
221 }
222
223 #[must_use]
225 pub fn rewards_address(&self) -> &RewardsAddress {
226 &self.rewards_address
227 }
228
229 #[must_use]
231 pub fn records_stored(&self) -> usize {
232 self.metrics_tracker.records_stored()
233 }
234
235 pub fn record_store(&self) {
237 self.metrics_tracker.record_store();
238 }
239
240 pub fn create_merkle_candidate_quote(
255 &self,
256 data_size: usize,
257 data_type: u32,
258 merkle_payment_timestamp: u64,
259 ) -> Result<MerklePaymentCandidateNode> {
260 let sign_fn = self
261 .sign_fn
262 .as_ref()
263 .ok_or_else(|| Error::Payment("Quote signing not configured".to_string()))?;
264
265 let price = calculate_price(self.pricing_records_stored());
266
267 let msg = MerklePaymentCandidateNode::bytes_to_sign(
269 &price,
270 &self.rewards_address,
271 merkle_payment_timestamp,
272 );
273
274 let signature = sign_fn(&msg);
276 if signature.is_empty() {
277 return Err(Error::Payment(
278 "ML-DSA-65 signing produced empty signature for merkle candidate".to_string(),
279 ));
280 }
281
282 let candidate = MerklePaymentCandidateNode {
283 pub_key: self.pub_key.clone(),
284 price,
285 reward_address: self.rewards_address,
286 merkle_payment_timestamp,
287 signature,
288 };
289
290 if crate::logging::enabled!(crate::logging::Level::DEBUG) {
291 debug!(
292 "Generated ML-DSA-65 merkle candidate quote (size: {data_size}, type: {data_type}, ts: {merkle_payment_timestamp})"
293 );
294 }
295
296 Ok(candidate)
297 }
298}
299
300pub fn wire_ml_dsa_signer(
320 generator: &mut QuoteGenerator,
321 identity: &saorsa_core::identity::NodeIdentity,
322) -> Result<()> {
323 let pub_key_bytes = identity.public_key().as_bytes().to_vec();
324 let sk_bytes = identity.secret_key_bytes().to_vec();
325 let sk = MlDsaSecretKey::from_bytes(&sk_bytes)
326 .map_err(|e| Error::Crypto(format!("Failed to deserialize ML-DSA-65 secret key: {e}")))?;
327 let ml_dsa = MlDsa65::new();
328 generator.set_signer(pub_key_bytes, move |msg| match ml_dsa.sign(&sk, msg) {
329 Ok(sig) => sig.as_bytes().to_vec(),
330 Err(e) => {
331 crate::logging::error!("ML-DSA-65 signing failed: {e}");
332 vec![]
333 }
334 });
335 generator.probe_signer()?;
336 Ok(())
337}
338
339#[cfg(test)]
340#[allow(clippy::expect_used)]
341mod tests {
342 use super::*;
343 use crate::payment::metrics::QuotingMetricsTracker;
344 use ant_protocol::payment::verify::{
349 verify_merkle_candidate_signature, verify_quote_content, verify_quote_signature,
350 };
351 use evmlib::common::Amount;
352 use saorsa_pqc::pqc::types::MlDsaSecretKey;
353
354 fn create_test_generator() -> QuoteGenerator {
355 let rewards_address = RewardsAddress::new([1u8; 20]);
356 let metrics_tracker = QuotingMetricsTracker::new(100);
357
358 let mut generator = QuoteGenerator::new(rewards_address, metrics_tracker);
359
360 generator.set_signer(vec![0u8; 64], |bytes| {
362 let mut sig = vec![0u8; 64];
364 for (i, b) in bytes.iter().take(64).enumerate() {
365 sig[i] = *b;
366 }
367 sig
368 });
369
370 generator
371 }
372
373 #[tokio::test]
384 async fn test_pricing_tracks_attached_storage_not_side_counter() {
385 use crate::payment::pricing::derive_records_stored_from_price;
386 use crate::storage::{LmdbStorage, LmdbStorageConfig};
387 use tempfile::TempDir;
388
389 let temp_dir = TempDir::new().expect("temp dir");
390 let storage = Arc::new(
391 LmdbStorage::new(LmdbStorageConfig {
392 root_dir: temp_dir.path().to_path_buf(),
393 ..LmdbStorageConfig::test_default()
394 })
395 .await
396 .expect("create storage"),
397 );
398
399 let metrics_tracker = QuotingMetricsTracker::new(3);
403 let mut generator = QuoteGenerator::new(RewardsAddress::new([1u8; 20]), metrics_tracker);
404 generator.set_signer(vec![0u8; 64], |bytes| {
405 let mut sig = vec![0u8; 64];
406 for (i, b) in bytes.iter().take(64).enumerate() {
407 sig[i] = *b;
408 }
409 sig
410 });
411 generator.attach_storage(Arc::clone(&storage));
412
413 for i in 0..25u32 {
416 let content = format!("replicated-record-{i}");
417 let address = LmdbStorage::compute_address(content.as_bytes());
418 storage
419 .put(&address, content.as_bytes())
420 .await
421 .expect("put");
422 }
423 assert_eq!(
424 generator.records_stored(),
425 3,
426 "side counter must be untouched"
427 );
428 assert_eq!(storage.current_chunks().expect("count"), 25);
429
430 let quote = generator
431 .create_quote([42u8; 32], 1024, 0)
432 .expect("create quote");
433
434 assert_eq!(
436 quote.price,
437 calculate_price(25),
438 "price must be derived from current_chunks(), not metrics_tracker"
439 );
440 assert_eq!(
441 derive_records_stored_from_price(quote.price),
442 25,
443 "verifier's price-inverse must recover the store count, keeping the \
444 freshness delta at ~0 for a freshly issued quote"
445 );
446 }
447
448 #[test]
449 fn test_create_quote() {
450 let generator = create_test_generator();
451 let content = [42u8; 32];
452
453 let quote = generator.create_quote(content, 1024, 0);
454 assert!(quote.is_ok());
455
456 let quote = quote.expect("valid quote");
457 assert_eq!(quote.content.0, content);
458 }
459
460 #[test]
461 fn test_verify_quote_content() {
462 let generator = create_test_generator();
463 let content = [42u8; 32];
464
465 let quote = generator
466 .create_quote(content, 1024, 0)
467 .expect("valid quote");
468 assert!(verify_quote_content("e, &content));
469
470 let wrong_content = [99u8; 32];
472 assert!(!verify_quote_content("e, &wrong_content));
473 }
474
475 #[test]
476 fn test_generator_without_signer() {
477 let rewards_address = RewardsAddress::new([1u8; 20]);
478 let metrics_tracker = QuotingMetricsTracker::new(100);
479 let generator = QuoteGenerator::new(rewards_address, metrics_tracker);
480
481 assert!(!generator.can_sign());
482
483 let content = [42u8; 32];
484 let result = generator.create_quote(content, 1024, 0);
485 assert!(result.is_err());
486 }
487
488 #[test]
489 fn test_quote_signature_round_trip_real_keys() {
490 let ml_dsa = MlDsa65::new();
491 let (public_key, secret_key) = ml_dsa.generate_keypair().expect("keypair generation");
492
493 let rewards_address = RewardsAddress::new([2u8; 20]);
494 let metrics_tracker = QuotingMetricsTracker::new(100);
495 let mut generator = QuoteGenerator::new(rewards_address, metrics_tracker);
496
497 let pub_key_bytes = public_key.as_bytes().to_vec();
498 let sk_bytes = secret_key.as_bytes().to_vec();
499 generator.set_signer(pub_key_bytes, move |msg| {
500 let sk = MlDsaSecretKey::from_bytes(&sk_bytes).expect("secret key parse");
501 let ml_dsa = MlDsa65::new();
502 ml_dsa.sign(&sk, msg).expect("signing").as_bytes().to_vec()
503 });
504
505 let content = [7u8; 32];
506 let quote = generator
507 .create_quote(content, 2048, 0)
508 .expect("create quote");
509
510 assert!(verify_quote_signature("e));
512
513 let mut tampered_quote = quote;
515 if let Some(byte) = tampered_quote.signature.first_mut() {
516 *byte ^= 0xFF;
517 }
518 assert!(!verify_quote_signature(&tampered_quote));
519 }
520
521 #[test]
522 fn test_empty_signature_fails_verification() {
523 let generator = create_test_generator();
524 let content = [42u8; 32];
525
526 let quote = generator
527 .create_quote(content, 1024, 0)
528 .expect("create quote");
529
530 assert!(!verify_quote_signature("e));
533 }
534
535 #[test]
536 fn test_rewards_address_getter() {
537 let addr = RewardsAddress::new([42u8; 20]);
538 let metrics_tracker = QuotingMetricsTracker::new(0);
539 let generator = QuoteGenerator::new(addr, metrics_tracker);
540
541 assert_eq!(*generator.rewards_address(), addr);
542 }
543
544 #[test]
545 fn test_records_stored() {
546 let rewards_address = RewardsAddress::new([1u8; 20]);
547 let metrics_tracker = QuotingMetricsTracker::new(50);
548 let generator = QuoteGenerator::new(rewards_address, metrics_tracker);
549
550 assert_eq!(generator.records_stored(), 50);
551 }
552
553 #[test]
554 fn test_record_store_delegation() {
555 let rewards_address = RewardsAddress::new([1u8; 20]);
556 let metrics_tracker = QuotingMetricsTracker::new(0);
557 let generator = QuoteGenerator::new(rewards_address, metrics_tracker);
558
559 generator.record_store();
560 generator.record_store();
561 generator.record_store();
562
563 assert_eq!(generator.records_stored(), 3);
564 }
565
566 #[test]
567 fn test_create_quote_different_data_types() {
568 let generator = create_test_generator();
569 let content = [10u8; 32];
570
571 let q0 = generator.create_quote(content, 1024, 0).expect("type 0");
573 let q1 = generator.create_quote(content, 512, 1).expect("type 1");
574 let q2 = generator.create_quote(content, 256, 2).expect("type 2");
575
576 assert!(q0.price >= Amount::from(1u64));
578 assert!(q1.price >= Amount::from(1u64));
579 assert!(q2.price >= Amount::from(1u64));
580 }
581
582 #[test]
583 fn test_create_quote_zero_size() {
584 let generator = create_test_generator();
585 let content = [11u8; 32];
586
587 let quote = generator.create_quote(content, 0, 0).expect("zero size");
589 assert!(quote.price >= Amount::from(1u64));
590 }
591
592 #[test]
593 fn test_create_quote_large_size() {
594 let generator = create_test_generator();
595 let content = [12u8; 32];
596
597 let quote = generator
599 .create_quote(content, 10_000_000, 0)
600 .expect("large size");
601 assert!(quote.price >= Amount::from(1u64));
602 }
603
604 #[test]
605 fn test_verify_quote_signature_empty_pub_key() {
606 let quote = PaymentQuote {
607 content: xor_name::XorName([0u8; 32]),
608 timestamp: SystemTime::now(),
609 price: Amount::from(1u64),
610 rewards_address: RewardsAddress::new([0u8; 20]),
611 pub_key: vec![],
612 signature: vec![],
613 };
614
615 assert!(!verify_quote_signature("e));
617 }
618
619 #[test]
620 fn test_can_sign_after_set_signer() {
621 let rewards_address = RewardsAddress::new([1u8; 20]);
622 let metrics_tracker = QuotingMetricsTracker::new(0);
623 let mut generator = QuoteGenerator::new(rewards_address, metrics_tracker);
624
625 assert!(!generator.can_sign());
626
627 generator.set_signer(vec![0u8; 32], |_| vec![0u8; 32]);
628
629 assert!(generator.can_sign());
630 }
631
632 #[test]
633 fn test_wire_ml_dsa_signer_returns_ok_with_valid_identity() {
634 let identity = saorsa_core::identity::NodeIdentity::generate().expect("keypair generation");
635 let rewards_address = RewardsAddress::new([3u8; 20]);
636 let metrics_tracker = QuotingMetricsTracker::new(0);
637 let mut generator = QuoteGenerator::new(rewards_address, metrics_tracker);
638
639 let result = wire_ml_dsa_signer(&mut generator, &identity);
640 assert!(
641 result.is_ok(),
642 "wire_ml_dsa_signer should succeed: {result:?}"
643 );
644 assert!(generator.can_sign());
645 }
646
647 #[test]
648 fn test_probe_signer_fails_without_signer() {
649 let rewards_address = RewardsAddress::new([1u8; 20]);
650 let metrics_tracker = QuotingMetricsTracker::new(0);
651 let generator = QuoteGenerator::new(rewards_address, metrics_tracker);
652
653 let result = generator.probe_signer();
654 assert!(result.is_err());
655 }
656
657 #[test]
658 fn test_probe_signer_fails_with_empty_signature() {
659 let rewards_address = RewardsAddress::new([1u8; 20]);
660 let metrics_tracker = QuotingMetricsTracker::new(0);
661 let mut generator = QuoteGenerator::new(rewards_address, metrics_tracker);
662
663 generator.set_signer(vec![0u8; 32], |_| vec![]);
664
665 let result = generator.probe_signer();
666 assert!(result.is_err());
667 }
668
669 #[test]
670 fn test_create_merkle_candidate_quote_with_ml_dsa() {
671 let ml_dsa = MlDsa65::new();
672 let (public_key, secret_key) = ml_dsa.generate_keypair().expect("keypair generation");
673
674 let rewards_address = RewardsAddress::new([0x42u8; 20]);
675 let metrics_tracker = QuotingMetricsTracker::new(50);
676 let mut generator = QuoteGenerator::new(rewards_address, metrics_tracker);
677
678 let pub_key_bytes = public_key.as_bytes().to_vec();
680 let sk_bytes = secret_key.as_bytes().to_vec();
681 generator.set_signer(pub_key_bytes.clone(), move |msg| {
682 let sk = MlDsaSecretKey::from_bytes(&sk_bytes).expect("sk parse");
683 let ml_dsa = MlDsa65::new();
684 ml_dsa.sign(&sk, msg).expect("sign").as_bytes().to_vec()
685 });
686
687 let timestamp = std::time::SystemTime::now()
688 .duration_since(std::time::UNIX_EPOCH)
689 .expect("system time")
690 .as_secs();
691
692 let result = generator.create_merkle_candidate_quote(2048, 0, timestamp);
693
694 assert!(
695 result.is_ok(),
696 "create_merkle_candidate_quote should succeed: {result:?}"
697 );
698
699 let candidate = result.expect("valid candidate");
700
701 assert_eq!(candidate.reward_address, rewards_address);
703
704 assert_eq!(candidate.merkle_payment_timestamp, timestamp);
706
707 assert_eq!(candidate.price, calculate_price(50));
709
710 assert_eq!(
712 candidate.pub_key, pub_key_bytes,
713 "Public key should be raw ML-DSA-65 bytes"
714 );
715
716 assert!(
718 verify_merkle_candidate_signature(&candidate),
719 "ML-DSA-65 merkle candidate signature must be valid"
720 );
721
722 let mut tampered = candidate;
724 tampered.merkle_payment_timestamp = timestamp + 1;
725 assert!(
726 !verify_merkle_candidate_signature(&tampered),
727 "Tampered timestamp should invalidate the ML-DSA-65 signature"
728 );
729 }
730
731 fn make_valid_merkle_candidate() -> MerklePaymentCandidateNode {
737 let ml_dsa = MlDsa65::new();
738 let (public_key, secret_key) = ml_dsa.generate_keypair().expect("keygen");
739
740 let rewards_address = RewardsAddress::new([0xABu8; 20]);
741 let timestamp = std::time::SystemTime::now()
742 .duration_since(std::time::UNIX_EPOCH)
743 .expect("system time")
744 .as_secs();
745 let price = Amount::from(42u64);
746
747 let msg = MerklePaymentCandidateNode::bytes_to_sign(&price, &rewards_address, timestamp);
748 let sk = MlDsaSecretKey::from_bytes(secret_key.as_bytes()).expect("sk");
749 let signature = ml_dsa.sign(&sk, &msg).expect("sign").as_bytes().to_vec();
750
751 MerklePaymentCandidateNode {
752 pub_key: public_key.as_bytes().to_vec(),
753 price,
754 reward_address: rewards_address,
755 merkle_payment_timestamp: timestamp,
756 signature,
757 }
758 }
759
760 #[test]
761 fn test_verify_merkle_candidate_valid_signature() {
762 let candidate = make_valid_merkle_candidate();
763 assert!(
764 verify_merkle_candidate_signature(&candidate),
765 "Freshly signed merkle candidate must verify"
766 );
767 }
768
769 #[test]
770 fn test_verify_merkle_candidate_tampered_pub_key() {
771 let mut candidate = make_valid_merkle_candidate();
772 if let Some(byte) = candidate.pub_key.first_mut() {
774 *byte ^= 0xFF;
775 }
776 assert!(
777 !verify_merkle_candidate_signature(&candidate),
778 "Tampered pub_key must invalidate the signature"
779 );
780 }
781
782 #[test]
783 fn test_verify_merkle_candidate_tampered_reward_address() {
784 let mut candidate = make_valid_merkle_candidate();
785 candidate.reward_address = RewardsAddress::new([0xFFu8; 20]);
786 assert!(
787 !verify_merkle_candidate_signature(&candidate),
788 "Tampered reward_address must invalidate the signature"
789 );
790 }
791
792 #[test]
793 fn test_verify_merkle_candidate_tampered_price() {
794 let mut candidate = make_valid_merkle_candidate();
795 candidate.price = Amount::from(999_999u64);
796 assert!(
797 !verify_merkle_candidate_signature(&candidate),
798 "Tampered price must invalidate the signature"
799 );
800 }
801
802 #[test]
803 fn test_verify_merkle_candidate_tampered_signature_byte() {
804 let mut candidate = make_valid_merkle_candidate();
805 if let Some(byte) = candidate.signature.first_mut() {
806 *byte ^= 0xFF;
807 }
808 assert!(
809 !verify_merkle_candidate_signature(&candidate),
810 "Tampered signature byte must fail verification"
811 );
812 }
813
814 #[test]
815 fn test_verify_merkle_candidate_empty_pub_key() {
816 let mut candidate = make_valid_merkle_candidate();
817 candidate.pub_key = vec![];
818 assert!(
819 !verify_merkle_candidate_signature(&candidate),
820 "Empty pub_key must fail verification"
821 );
822 }
823
824 #[test]
825 fn test_verify_merkle_candidate_empty_signature() {
826 let mut candidate = make_valid_merkle_candidate();
827 candidate.signature = vec![];
828 assert!(
829 !verify_merkle_candidate_signature(&candidate),
830 "Empty signature must fail verification"
831 );
832 }
833
834 #[test]
835 fn test_verify_merkle_candidate_wrong_length_signature() {
836 let mut candidate = make_valid_merkle_candidate();
837 candidate.signature = vec![0xAA; 100];
839 assert!(
840 !verify_merkle_candidate_signature(&candidate),
841 "Wrong-length signature must fail verification"
842 );
843 }
844
845 #[test]
846 fn test_verify_merkle_candidate_wrong_length_pub_key() {
847 let mut candidate = make_valid_merkle_candidate();
848 candidate.pub_key = vec![0xBB; 100];
850 assert!(
851 !verify_merkle_candidate_signature(&candidate),
852 "Wrong-length pub_key must fail verification"
853 );
854 }
855
856 #[test]
857 fn test_verify_merkle_candidate_cross_key_rejection() {
858 let candidate = make_valid_merkle_candidate();
860 let ml_dsa = MlDsa65::new();
861 let (other_pk, _) = ml_dsa.generate_keypair().expect("keygen");
862
863 let mut swapped = candidate;
864 swapped.pub_key = other_pk.as_bytes().to_vec();
865 assert!(
866 !verify_merkle_candidate_signature(&swapped),
867 "Signature from key A must not verify under key B"
868 );
869 }
870}