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 create_merkle_candidate_quote(
254 &self,
255 data_size: usize,
256 data_type: u32,
257 merkle_payment_timestamp: u64,
258 ) -> Result<MerklePaymentCandidateNode> {
259 let sign_fn = self
260 .sign_fn
261 .as_ref()
262 .ok_or_else(|| Error::Payment("Quote signing not configured".to_string()))?;
263
264 let price = calculate_price(self.pricing_records_stored());
265
266 let msg = MerklePaymentCandidateNode::bytes_to_sign(
268 &price,
269 &self.rewards_address,
270 merkle_payment_timestamp,
271 );
272
273 let signature = sign_fn(&msg);
275 if signature.is_empty() {
276 return Err(Error::Payment(
277 "ML-DSA-65 signing produced empty signature for merkle candidate".to_string(),
278 ));
279 }
280
281 let candidate = MerklePaymentCandidateNode {
282 pub_key: self.pub_key.clone(),
283 price,
284 reward_address: self.rewards_address,
285 merkle_payment_timestamp,
286 signature,
287 };
288
289 if crate::logging::enabled!(crate::logging::Level::DEBUG) {
290 debug!(
291 "Generated ML-DSA-65 merkle candidate quote (size: {data_size}, type: {data_type}, ts: {merkle_payment_timestamp})"
292 );
293 }
294
295 Ok(candidate)
296 }
297}
298
299pub fn wire_ml_dsa_signer(
319 generator: &mut QuoteGenerator,
320 identity: &saorsa_core::identity::NodeIdentity,
321) -> Result<()> {
322 let pub_key_bytes = identity.public_key().as_bytes().to_vec();
323 let sk_bytes = identity.secret_key_bytes().to_vec();
324 let sk = MlDsaSecretKey::from_bytes(&sk_bytes)
325 .map_err(|e| Error::Crypto(format!("Failed to deserialize ML-DSA-65 secret key: {e}")))?;
326 let ml_dsa = MlDsa65::new();
327 generator.set_signer(pub_key_bytes, move |msg| match ml_dsa.sign(&sk, msg) {
328 Ok(sig) => sig.as_bytes().to_vec(),
329 Err(e) => {
330 crate::logging::error!("ML-DSA-65 signing failed: {e}");
331 vec![]
332 }
333 });
334 generator.probe_signer()?;
335 Ok(())
336}
337
338#[cfg(test)]
339#[allow(clippy::expect_used)]
340mod tests {
341 use super::*;
342 use crate::payment::metrics::QuotingMetricsTracker;
343 use ant_protocol::payment::verify::{
348 verify_merkle_candidate_signature, verify_quote_content, verify_quote_signature,
349 };
350 use evmlib::common::Amount;
351 use saorsa_pqc::pqc::types::MlDsaSecretKey;
352
353 fn create_test_generator() -> QuoteGenerator {
354 let rewards_address = RewardsAddress::new([1u8; 20]);
355 let metrics_tracker = QuotingMetricsTracker::new(100);
356
357 let mut generator = QuoteGenerator::new(rewards_address, metrics_tracker);
358
359 generator.set_signer(vec![0u8; 64], |bytes| {
361 let mut sig = vec![0u8; 64];
363 for (i, b) in bytes.iter().take(64).enumerate() {
364 sig[i] = *b;
365 }
366 sig
367 });
368
369 generator
370 }
371
372 #[tokio::test]
383 async fn test_pricing_tracks_attached_storage_not_side_counter() {
384 use crate::payment::pricing::derive_records_stored_from_price;
385 use crate::storage::{LmdbStorage, LmdbStorageConfig};
386 use tempfile::TempDir;
387
388 let temp_dir = TempDir::new().expect("temp dir");
389 let storage = Arc::new(
390 LmdbStorage::new(LmdbStorageConfig {
391 root_dir: temp_dir.path().to_path_buf(),
392 ..LmdbStorageConfig::test_default()
393 })
394 .await
395 .expect("create storage"),
396 );
397
398 let metrics_tracker = QuotingMetricsTracker::new(3);
402 let mut generator = QuoteGenerator::new(RewardsAddress::new([1u8; 20]), metrics_tracker);
403 generator.set_signer(vec![0u8; 64], |bytes| {
404 let mut sig = vec![0u8; 64];
405 for (i, b) in bytes.iter().take(64).enumerate() {
406 sig[i] = *b;
407 }
408 sig
409 });
410 generator.attach_storage(Arc::clone(&storage));
411
412 for i in 0..25u32 {
415 let content = format!("replicated-record-{i}");
416 let address = LmdbStorage::compute_address(content.as_bytes());
417 storage
418 .put(&address, content.as_bytes())
419 .await
420 .expect("put");
421 }
422 assert_eq!(
423 generator.records_stored(),
424 3,
425 "side counter must be untouched"
426 );
427 assert_eq!(storage.current_chunks().expect("count"), 25);
428
429 let quote = generator
430 .create_quote([42u8; 32], 1024, 0)
431 .expect("create quote");
432
433 assert_eq!(
435 quote.price,
436 calculate_price(25),
437 "price must be derived from current_chunks(), not metrics_tracker"
438 );
439 assert_eq!(
440 derive_records_stored_from_price(quote.price),
441 25,
442 "verifier's price-inverse must recover the store count, keeping the \
443 local price comparison aligned for a freshly issued quote"
444 );
445 }
446
447 #[test]
448 fn test_create_quote() {
449 let generator = create_test_generator();
450 let content = [42u8; 32];
451
452 let quote = generator.create_quote(content, 1024, 0);
453 assert!(quote.is_ok());
454
455 let quote = quote.expect("valid quote");
456 assert_eq!(quote.content.0, content);
457 }
458
459 #[test]
460 fn test_verify_quote_content() {
461 let generator = create_test_generator();
462 let content = [42u8; 32];
463
464 let quote = generator
465 .create_quote(content, 1024, 0)
466 .expect("valid quote");
467 assert!(verify_quote_content("e, &content));
468
469 let wrong_content = [99u8; 32];
471 assert!(!verify_quote_content("e, &wrong_content));
472 }
473
474 #[test]
475 fn test_generator_without_signer() {
476 let rewards_address = RewardsAddress::new([1u8; 20]);
477 let metrics_tracker = QuotingMetricsTracker::new(100);
478 let generator = QuoteGenerator::new(rewards_address, metrics_tracker);
479
480 assert!(!generator.can_sign());
481
482 let content = [42u8; 32];
483 let result = generator.create_quote(content, 1024, 0);
484 assert!(result.is_err());
485 }
486
487 #[test]
488 fn test_quote_signature_round_trip_real_keys() {
489 let ml_dsa = MlDsa65::new();
490 let (public_key, secret_key) = ml_dsa.generate_keypair().expect("keypair generation");
491
492 let rewards_address = RewardsAddress::new([2u8; 20]);
493 let metrics_tracker = QuotingMetricsTracker::new(100);
494 let mut generator = QuoteGenerator::new(rewards_address, metrics_tracker);
495
496 let pub_key_bytes = public_key.as_bytes().to_vec();
497 let sk_bytes = secret_key.as_bytes().to_vec();
498 generator.set_signer(pub_key_bytes, move |msg| {
499 let sk = MlDsaSecretKey::from_bytes(&sk_bytes).expect("secret key parse");
500 let ml_dsa = MlDsa65::new();
501 ml_dsa.sign(&sk, msg).expect("signing").as_bytes().to_vec()
502 });
503
504 let content = [7u8; 32];
505 let quote = generator
506 .create_quote(content, 2048, 0)
507 .expect("create quote");
508
509 assert!(verify_quote_signature("e));
511
512 let mut tampered_quote = quote;
514 if let Some(byte) = tampered_quote.signature.first_mut() {
515 *byte ^= 0xFF;
516 }
517 assert!(!verify_quote_signature(&tampered_quote));
518 }
519
520 #[test]
521 fn test_empty_signature_fails_verification() {
522 let generator = create_test_generator();
523 let content = [42u8; 32];
524
525 let quote = generator
526 .create_quote(content, 1024, 0)
527 .expect("create quote");
528
529 assert!(!verify_quote_signature("e));
532 }
533
534 #[test]
535 fn test_rewards_address_getter() {
536 let addr = RewardsAddress::new([42u8; 20]);
537 let metrics_tracker = QuotingMetricsTracker::new(0);
538 let generator = QuoteGenerator::new(addr, metrics_tracker);
539
540 assert_eq!(*generator.rewards_address(), addr);
541 }
542
543 #[test]
544 fn test_records_stored() {
545 let rewards_address = RewardsAddress::new([1u8; 20]);
546 let metrics_tracker = QuotingMetricsTracker::new(50);
547 let generator = QuoteGenerator::new(rewards_address, metrics_tracker);
548
549 assert_eq!(generator.records_stored(), 50);
550 }
551
552 #[test]
553 fn test_record_store_delegation() {
554 let rewards_address = RewardsAddress::new([1u8; 20]);
555 let metrics_tracker = QuotingMetricsTracker::new(0);
556 let generator = QuoteGenerator::new(rewards_address, metrics_tracker);
557
558 generator.record_store();
559 generator.record_store();
560 generator.record_store();
561
562 assert_eq!(generator.records_stored(), 3);
563 }
564
565 #[test]
566 fn test_create_quote_different_data_types() {
567 let generator = create_test_generator();
568 let content = [10u8; 32];
569
570 let q0 = generator.create_quote(content, 1024, 0).expect("type 0");
572 let q1 = generator.create_quote(content, 512, 1).expect("type 1");
573 let q2 = generator.create_quote(content, 256, 2).expect("type 2");
574
575 assert!(q0.price >= Amount::from(1u64));
577 assert!(q1.price >= Amount::from(1u64));
578 assert!(q2.price >= Amount::from(1u64));
579 }
580
581 #[test]
582 fn test_create_quote_zero_size() {
583 let generator = create_test_generator();
584 let content = [11u8; 32];
585
586 let quote = generator.create_quote(content, 0, 0).expect("zero size");
588 assert!(quote.price >= Amount::from(1u64));
589 }
590
591 #[test]
592 fn test_create_quote_large_size() {
593 let generator = create_test_generator();
594 let content = [12u8; 32];
595
596 let quote = generator
598 .create_quote(content, 10_000_000, 0)
599 .expect("large size");
600 assert!(quote.price >= Amount::from(1u64));
601 }
602
603 #[test]
604 fn test_verify_quote_signature_empty_pub_key() {
605 let quote = PaymentQuote {
606 content: xor_name::XorName([0u8; 32]),
607 timestamp: SystemTime::now(),
608 price: Amount::from(1u64),
609 rewards_address: RewardsAddress::new([0u8; 20]),
610 pub_key: vec![],
611 signature: vec![],
612 };
613
614 assert!(!verify_quote_signature("e));
616 }
617
618 #[test]
619 fn test_can_sign_after_set_signer() {
620 let rewards_address = RewardsAddress::new([1u8; 20]);
621 let metrics_tracker = QuotingMetricsTracker::new(0);
622 let mut generator = QuoteGenerator::new(rewards_address, metrics_tracker);
623
624 assert!(!generator.can_sign());
625
626 generator.set_signer(vec![0u8; 32], |_| vec![0u8; 32]);
627
628 assert!(generator.can_sign());
629 }
630
631 #[test]
632 fn test_wire_ml_dsa_signer_returns_ok_with_valid_identity() {
633 let identity = saorsa_core::identity::NodeIdentity::generate().expect("keypair generation");
634 let rewards_address = RewardsAddress::new([3u8; 20]);
635 let metrics_tracker = QuotingMetricsTracker::new(0);
636 let mut generator = QuoteGenerator::new(rewards_address, metrics_tracker);
637
638 let result = wire_ml_dsa_signer(&mut generator, &identity);
639 assert!(
640 result.is_ok(),
641 "wire_ml_dsa_signer should succeed: {result:?}"
642 );
643 assert!(generator.can_sign());
644 }
645
646 #[test]
647 fn test_probe_signer_fails_without_signer() {
648 let rewards_address = RewardsAddress::new([1u8; 20]);
649 let metrics_tracker = QuotingMetricsTracker::new(0);
650 let generator = QuoteGenerator::new(rewards_address, metrics_tracker);
651
652 let result = generator.probe_signer();
653 assert!(result.is_err());
654 }
655
656 #[test]
657 fn test_probe_signer_fails_with_empty_signature() {
658 let rewards_address = RewardsAddress::new([1u8; 20]);
659 let metrics_tracker = QuotingMetricsTracker::new(0);
660 let mut generator = QuoteGenerator::new(rewards_address, metrics_tracker);
661
662 generator.set_signer(vec![0u8; 32], |_| vec![]);
663
664 let result = generator.probe_signer();
665 assert!(result.is_err());
666 }
667
668 #[test]
669 fn test_create_merkle_candidate_quote_with_ml_dsa() {
670 let ml_dsa = MlDsa65::new();
671 let (public_key, secret_key) = ml_dsa.generate_keypair().expect("keypair generation");
672
673 let rewards_address = RewardsAddress::new([0x42u8; 20]);
674 let metrics_tracker = QuotingMetricsTracker::new(50);
675 let mut generator = QuoteGenerator::new(rewards_address, metrics_tracker);
676
677 let pub_key_bytes = public_key.as_bytes().to_vec();
679 let sk_bytes = secret_key.as_bytes().to_vec();
680 generator.set_signer(pub_key_bytes.clone(), move |msg| {
681 let sk = MlDsaSecretKey::from_bytes(&sk_bytes).expect("sk parse");
682 let ml_dsa = MlDsa65::new();
683 ml_dsa.sign(&sk, msg).expect("sign").as_bytes().to_vec()
684 });
685
686 let timestamp = std::time::SystemTime::now()
687 .duration_since(std::time::UNIX_EPOCH)
688 .expect("system time")
689 .as_secs();
690
691 let result = generator.create_merkle_candidate_quote(2048, 0, timestamp);
692
693 assert!(
694 result.is_ok(),
695 "create_merkle_candidate_quote should succeed: {result:?}"
696 );
697
698 let candidate = result.expect("valid candidate");
699
700 assert_eq!(candidate.reward_address, rewards_address);
702
703 assert_eq!(candidate.merkle_payment_timestamp, timestamp);
705
706 assert_eq!(candidate.price, calculate_price(50));
708
709 assert_eq!(
711 candidate.pub_key, pub_key_bytes,
712 "Public key should be raw ML-DSA-65 bytes"
713 );
714
715 assert!(
717 verify_merkle_candidate_signature(&candidate),
718 "ML-DSA-65 merkle candidate signature must be valid"
719 );
720
721 let mut tampered = candidate;
723 tampered.merkle_payment_timestamp = timestamp + 1;
724 assert!(
725 !verify_merkle_candidate_signature(&tampered),
726 "Tampered timestamp should invalidate the ML-DSA-65 signature"
727 );
728 }
729
730 fn make_valid_merkle_candidate() -> MerklePaymentCandidateNode {
736 let ml_dsa = MlDsa65::new();
737 let (public_key, secret_key) = ml_dsa.generate_keypair().expect("keygen");
738
739 let rewards_address = RewardsAddress::new([0xABu8; 20]);
740 let timestamp = std::time::SystemTime::now()
741 .duration_since(std::time::UNIX_EPOCH)
742 .expect("system time")
743 .as_secs();
744 let price = Amount::from(42u64);
745
746 let msg = MerklePaymentCandidateNode::bytes_to_sign(&price, &rewards_address, timestamp);
747 let sk = MlDsaSecretKey::from_bytes(secret_key.as_bytes()).expect("sk");
748 let signature = ml_dsa.sign(&sk, &msg).expect("sign").as_bytes().to_vec();
749
750 MerklePaymentCandidateNode {
751 pub_key: public_key.as_bytes().to_vec(),
752 price,
753 reward_address: rewards_address,
754 merkle_payment_timestamp: timestamp,
755 signature,
756 }
757 }
758
759 #[test]
760 fn test_verify_merkle_candidate_valid_signature() {
761 let candidate = make_valid_merkle_candidate();
762 assert!(
763 verify_merkle_candidate_signature(&candidate),
764 "Freshly signed merkle candidate must verify"
765 );
766 }
767
768 #[test]
769 fn test_verify_merkle_candidate_tampered_pub_key() {
770 let mut candidate = make_valid_merkle_candidate();
771 if let Some(byte) = candidate.pub_key.first_mut() {
773 *byte ^= 0xFF;
774 }
775 assert!(
776 !verify_merkle_candidate_signature(&candidate),
777 "Tampered pub_key must invalidate the signature"
778 );
779 }
780
781 #[test]
782 fn test_verify_merkle_candidate_tampered_reward_address() {
783 let mut candidate = make_valid_merkle_candidate();
784 candidate.reward_address = RewardsAddress::new([0xFFu8; 20]);
785 assert!(
786 !verify_merkle_candidate_signature(&candidate),
787 "Tampered reward_address must invalidate the signature"
788 );
789 }
790
791 #[test]
792 fn test_verify_merkle_candidate_tampered_price() {
793 let mut candidate = make_valid_merkle_candidate();
794 candidate.price = Amount::from(999_999u64);
795 assert!(
796 !verify_merkle_candidate_signature(&candidate),
797 "Tampered price must invalidate the signature"
798 );
799 }
800
801 #[test]
802 fn test_verify_merkle_candidate_tampered_signature_byte() {
803 let mut candidate = make_valid_merkle_candidate();
804 if let Some(byte) = candidate.signature.first_mut() {
805 *byte ^= 0xFF;
806 }
807 assert!(
808 !verify_merkle_candidate_signature(&candidate),
809 "Tampered signature byte must fail verification"
810 );
811 }
812
813 #[test]
814 fn test_verify_merkle_candidate_empty_pub_key() {
815 let mut candidate = make_valid_merkle_candidate();
816 candidate.pub_key = vec![];
817 assert!(
818 !verify_merkle_candidate_signature(&candidate),
819 "Empty pub_key must fail verification"
820 );
821 }
822
823 #[test]
824 fn test_verify_merkle_candidate_empty_signature() {
825 let mut candidate = make_valid_merkle_candidate();
826 candidate.signature = vec![];
827 assert!(
828 !verify_merkle_candidate_signature(&candidate),
829 "Empty signature must fail verification"
830 );
831 }
832
833 #[test]
834 fn test_verify_merkle_candidate_wrong_length_signature() {
835 let mut candidate = make_valid_merkle_candidate();
836 candidate.signature = vec![0xAA; 100];
838 assert!(
839 !verify_merkle_candidate_signature(&candidate),
840 "Wrong-length signature must fail verification"
841 );
842 }
843
844 #[test]
845 fn test_verify_merkle_candidate_wrong_length_pub_key() {
846 let mut candidate = make_valid_merkle_candidate();
847 candidate.pub_key = vec![0xBB; 100];
849 assert!(
850 !verify_merkle_candidate_signature(&candidate),
851 "Wrong-length pub_key must fail verification"
852 );
853 }
854
855 #[test]
856 fn test_verify_merkle_candidate_cross_key_rejection() {
857 let candidate = make_valid_merkle_candidate();
859 let ml_dsa = MlDsa65::new();
860 let (other_pk, _) = ml_dsa.generate_keypair().expect("keygen");
861
862 let mut swapped = candidate;
863 swapped.pub_key = other_pk.as_bytes().to_vec();
864 assert!(
865 !verify_merkle_candidate_signature(&swapped),
866 "Signature from key A must not verify under key B"
867 );
868 }
869}