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