1use ml_dsa::{
5 EncodedVerifyingKey, KeyGen, MlDsa65, Signature as MlDsaSignature,
6 SigningKey as MlDsaSigningKey, VerifyingKey as MlDsaVerifyingKey, B32,
7};
8use serde::{Deserialize, Deserializer, Serialize, Serializer};
9use sha2::{Digest, Sha256};
10use slh_dsa::{
11 Shake128f as SlhDsa128f, Signature as SlhDsaSignature, VerifyingKey as SlhDsaVerifyingKey,
12};
13use std::fmt;
14
15pub const PQ_SCHEME_ML_DSA_65: u8 = 0x01;
16pub const PQ_SCHEME_SLH_DSA_128F: u8 = 0x02;
17
18pub const ML_DSA_65_PUBLIC_KEY_BYTES: usize = 1952;
19pub const ML_DSA_65_SIGNATURE_BYTES: usize = 3309;
20pub const SLH_DSA_128F_PUBLIC_KEY_BYTES: usize = 32;
21pub const SLH_DSA_128F_SIGNATURE_BYTES: usize = 17_088;
22
23fn serialize_pq_blob<S>(bytes: &[u8], serializer: S) -> Result<S::Ok, S::Error>
24where
25 S: Serializer,
26{
27 if serializer.is_human_readable() {
28 String::serialize(&hex::encode(bytes), serializer)
29 } else {
30 serializer.serialize_bytes(bytes)
31 }
32}
33
34fn deserialize_pq_blob<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
35where
36 D: Deserializer<'de>,
37{
38 if deserializer.is_human_readable() {
39 let encoded = String::deserialize(deserializer)?;
40 hex::decode(encoded).map_err(serde::de::Error::custom)
41 } else {
42 Vec::<u8>::deserialize(deserializer)
43 }
44}
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
48pub struct Pubkey(pub [u8; 32]);
49
50pub type Address = Pubkey;
51
52impl AsRef<[u8]> for Pubkey {
53 fn as_ref(&self) -> &[u8] {
54 &self.0
55 }
56}
57
58impl Pubkey {
59 pub const fn new(bytes: [u8; 32]) -> Self {
60 Pubkey(bytes)
61 }
62
63 pub fn to_base58(&self) -> String {
65 bs58::encode(self.0).into_string()
66 }
67
68 pub fn to_evm(&self) -> String {
70 use sha3::{Digest, Keccak256};
71 let hash = Keccak256::digest(self.0);
72 let evm_bytes = &hash[12..32]; format!("0x{}", hex::encode(evm_bytes))
74 }
75
76 pub fn from_base58(s: &str) -> Result<Self, String> {
78 let bytes = bs58::decode(s)
79 .into_vec()
80 .map_err(|e| format!("Invalid base58: {}", e))?;
81
82 if bytes.len() != 32 {
83 return Err(format!("Invalid length: {} (expected 32)", bytes.len()));
84 }
85
86 let mut pubkey = [0u8; 32];
87 pubkey.copy_from_slice(&bytes);
88 Ok(Pubkey(pubkey))
89 }
90}
91
92impl fmt::Display for Pubkey {
93 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
94 write!(f, "{}", self.to_base58())
95 }
96}
97
98#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
99pub struct PqPublicKey {
100 pub scheme_version: u8,
101 #[serde(
102 serialize_with = "serialize_pq_blob",
103 deserialize_with = "deserialize_pq_blob"
104 )]
105 pub bytes: Vec<u8>,
106}
107
108impl PqPublicKey {
109 pub fn new(scheme_version: u8, bytes: Vec<u8>) -> Result<Self, String> {
110 let key = Self {
111 scheme_version,
112 bytes,
113 };
114 key.validate()?;
115 Ok(key)
116 }
117
118 pub fn validate(&self) -> Result<(), String> {
119 let expected_len = match self.scheme_version {
120 PQ_SCHEME_ML_DSA_65 => ML_DSA_65_PUBLIC_KEY_BYTES,
121 PQ_SCHEME_SLH_DSA_128F => SLH_DSA_128F_PUBLIC_KEY_BYTES,
122 other => return Err(format!("Unsupported PQ public key scheme: 0x{other:02x}")),
123 };
124
125 if self.bytes.len() != expected_len {
126 return Err(format!(
127 "Invalid PQ public key length for scheme 0x{:02x}: {} (expected {})",
128 self.scheme_version,
129 self.bytes.len(),
130 expected_len
131 ));
132 }
133
134 Ok(())
135 }
136
137 pub fn address(&self) -> Pubkey {
138 let digest = Sha256::digest(&self.bytes);
139 let mut address = [0u8; 32];
140 address[0] = self.scheme_version;
141 address[1..].copy_from_slice(&digest[..31]);
142 Pubkey(address)
143 }
144
145 pub fn from_ml_dsa(verifying_key: &MlDsaVerifyingKey<MlDsa65>) -> Self {
146 Self {
147 scheme_version: PQ_SCHEME_ML_DSA_65,
148 bytes: verifying_key.encode().as_slice().to_vec(),
149 }
150 }
151
152 fn as_ml_dsa_verifying_key(&self) -> Option<MlDsaVerifyingKey<MlDsa65>> {
153 if self.scheme_version != PQ_SCHEME_ML_DSA_65 {
154 return None;
155 }
156
157 let encoded = EncodedVerifyingKey::<MlDsa65>::try_from(self.bytes.as_slice()).ok()?;
158 Some(MlDsaVerifyingKey::<MlDsa65>::decode(&encoded))
159 }
160
161 fn as_slh_dsa_verifying_key(&self) -> Option<SlhDsaVerifyingKey<SlhDsa128f>> {
162 if self.scheme_version != PQ_SCHEME_SLH_DSA_128F {
163 return None;
164 }
165
166 SlhDsaVerifyingKey::<SlhDsa128f>::try_from(self.bytes.as_slice()).ok()
167 }
168}
169
170#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
171pub struct PqSignature {
172 pub scheme_version: u8,
173 pub public_key: PqPublicKey,
174 #[serde(
175 serialize_with = "serialize_pq_blob",
176 deserialize_with = "deserialize_pq_blob"
177 )]
178 pub sig: Vec<u8>,
179}
180
181impl PqSignature {
182 pub fn new(scheme_version: u8, public_key: PqPublicKey, sig: Vec<u8>) -> Result<Self, String> {
183 let signature = Self {
184 scheme_version,
185 public_key,
186 sig,
187 };
188 signature.validate()?;
189 Ok(signature)
190 }
191
192 pub fn validate(&self) -> Result<(), String> {
193 self.public_key.validate()?;
194
195 if self.public_key.scheme_version != self.scheme_version {
196 return Err(format!(
197 "PQ signature/public-key scheme mismatch: 0x{:02x} vs 0x{:02x}",
198 self.scheme_version, self.public_key.scheme_version
199 ));
200 }
201
202 let expected_len = match self.scheme_version {
203 PQ_SCHEME_ML_DSA_65 => ML_DSA_65_SIGNATURE_BYTES,
204 PQ_SCHEME_SLH_DSA_128F => SLH_DSA_128F_SIGNATURE_BYTES,
205 other => return Err(format!("Unsupported PQ signature scheme: 0x{other:02x}")),
206 };
207
208 if self.sig.len() != expected_len {
209 return Err(format!(
210 "Invalid PQ signature length for scheme 0x{:02x}: {} (expected {})",
211 self.scheme_version,
212 self.sig.len(),
213 expected_len
214 ));
215 }
216
217 Ok(())
218 }
219
220 pub fn signer_address(&self) -> Pubkey {
221 self.public_key.address()
222 }
223
224 fn as_ml_dsa_signature(&self) -> Option<MlDsaSignature<MlDsa65>> {
225 if self.scheme_version != PQ_SCHEME_ML_DSA_65 {
226 return None;
227 }
228
229 MlDsaSignature::<MlDsa65>::try_from(self.sig.as_slice()).ok()
230 }
231
232 fn as_slh_dsa_signature(&self) -> Option<SlhDsaSignature<SlhDsa128f>> {
233 if self.scheme_version != PQ_SCHEME_SLH_DSA_128F {
234 return None;
235 }
236
237 SlhDsaSignature::<SlhDsa128f>::try_from(self.sig.as_slice()).ok()
238 }
239
240 #[cfg(test)]
241 pub fn test_fixture(fill: u8) -> Self {
242 let public_key = PqPublicKey {
243 scheme_version: PQ_SCHEME_ML_DSA_65,
244 bytes: vec![fill; ML_DSA_65_PUBLIC_KEY_BYTES],
245 };
246 Self {
247 scheme_version: PQ_SCHEME_ML_DSA_65,
248 public_key,
249 sig: vec![fill; ML_DSA_65_SIGNATURE_BYTES],
250 }
251 }
252}
253
254pub struct Keypair {
256 keypair: MlDsaSigningKey<MlDsa65>,
257 seed: [u8; 32],
258}
259
260impl Keypair {
261 pub fn new() -> Self {
263 let mut seed = [0u8; 32];
264 getrandom::fill(&mut seed).expect("Failed to generate random seed");
265 Self::from_seed(&seed)
266 }
267
268 pub fn generate() -> Self {
270 Self::new()
271 }
272
273 pub fn secret_key(&self) -> &[u8; 32] {
275 &self.seed
276 }
277
278 pub fn secret(&self) -> &[u8; 32] {
279 &self.seed
280 }
281
282 pub fn from_seed(seed: &[u8; 32]) -> Self {
284 let pq_seed = match B32::try_from(seed.as_slice()) {
285 Ok(seed) => seed,
286 Err(_) => unreachable!("ML-DSA seed length must be 32 bytes"),
287 };
288 let keypair = MlDsa65::from_seed(&pq_seed);
289 Keypair {
290 keypair,
291 seed: *seed,
292 }
293 }
294
295 pub fn pubkey(&self) -> Pubkey {
297 self.public_key().address()
298 }
299
300 pub fn public_key(&self) -> PqPublicKey {
302 let verifying_key =
303 <MlDsaSigningKey<MlDsa65> as ml_dsa::signature::Keypair>::verifying_key(&self.keypair);
304 PqPublicKey::from_ml_dsa(&verifying_key)
305 }
306
307 pub fn to_seed(&self) -> [u8; 32] {
309 self.seed
310 }
311
312 pub fn sign(&self, message: &[u8]) -> PqSignature {
314 let signature = self
315 .keypair
316 .signing_key()
317 .sign_deterministic(message, &[])
318 .expect("ML-DSA-65 deterministic signing failed");
319
320 PqSignature::new(
321 PQ_SCHEME_ML_DSA_65,
322 self.public_key(),
323 signature.encode().as_slice().to_vec(),
324 )
325 .expect("ML-DSA-65 signature encoding produced invalid output")
326 }
327
328 pub fn verify(address: &Pubkey, message: &[u8], signature: &PqSignature) -> bool {
330 if signature.validate().is_err() {
331 return false;
332 }
333
334 if signature.signer_address() != *address {
335 return false;
336 }
337
338 match signature.scheme_version {
339 PQ_SCHEME_ML_DSA_65 => {
340 let verifying_key = match signature.public_key.as_ml_dsa_verifying_key() {
341 Some(verifying_key) => verifying_key,
342 None => return false,
343 };
344 let ml_signature = match signature.as_ml_dsa_signature() {
345 Some(signature) => signature,
346 None => return false,
347 };
348 verifying_key.verify_with_context(message, &[], &ml_signature)
349 }
350 PQ_SCHEME_SLH_DSA_128F => {
351 let verifying_key = match signature.public_key.as_slh_dsa_verifying_key() {
352 Some(verifying_key) => verifying_key,
353 None => return false,
354 };
355 let slh_signature = match signature.as_slh_dsa_signature() {
356 Some(signature) => signature,
357 None => return false,
358 };
359 slh_dsa::signature::Verifier::verify(&verifying_key, message, &slh_signature)
360 .is_ok()
361 }
362 _ => false,
363 }
364 }
365}
366
367impl Default for Keypair {
368 fn default() -> Self {
369 Self::new()
370 }
371}
372
373#[derive(Debug, Clone, Serialize, Deserialize)]
375pub struct Account {
376 pub spores: u64,
379
380 #[serde(default)]
382 pub spendable: u64,
383
384 #[serde(default)]
386 pub staked: u64,
387
388 #[serde(default)]
390 pub locked: u64,
391
392 pub data: Vec<u8>,
394
395 #[serde(default)]
397 pub public_key: Option<PqPublicKey>,
398
399 pub owner: Pubkey,
401
402 pub executable: bool,
404
405 pub rent_epoch: u64,
407
408 #[serde(default)]
410 pub dormant: bool,
411
412 #[serde(default)]
414 pub missed_rent_epochs: u64,
415}
416
417impl Account {
418 pub fn fixup_legacy(&mut self) {
421 if self.spores > 0 && self.spendable == 0 && self.staked == 0 && self.locked == 0 {
422 self.spendable = self.spores;
423 }
424 }
425
426 pub const fn licn_to_spores(licn: u64) -> u64 {
428 licn.saturating_mul(1_000_000_000)
429 }
430
431 pub const fn spores_to_licn(spores: u64) -> u64 {
435 spores / 1_000_000_000
436 }
437
438 pub fn new(licn: u64, owner: Pubkey) -> Self {
440 let spores = Self::licn_to_spores(licn);
441 Account {
442 spores,
443 spendable: spores, staked: 0,
445 locked: 0,
446 data: Vec::new(),
447 public_key: None,
448 owner,
449 executable: false,
450 rent_epoch: 0,
451 dormant: false,
452 missed_rent_epochs: 0,
453 }
454 }
455
456 pub fn stake(&mut self, amount: u64) -> Result<(), String> {
460 if amount == 0 {
462 return Ok(());
463 }
464 let new_spendable = self.spendable.checked_sub(amount).ok_or_else(|| {
465 format!(
466 "Insufficient spendable balance: {} < {}",
467 self.spendable, amount
468 )
469 })?;
470 let new_staked = self.staked.checked_add(amount).ok_or_else(|| {
471 format!(
472 "Overflow adding {} to staked balance {}",
473 amount, self.staked
474 )
475 })?;
476 self.spendable = new_spendable;
477 self.staked = new_staked;
478 if self.spores != self.spendable + self.staked + self.locked {
479 return Err("Account invariant violated after stake".to_string());
480 }
481 Ok(())
482 }
483
484 pub fn unstake(&mut self, amount: u64) -> Result<(), String> {
487 if amount == 0 {
489 return Ok(());
490 }
491 let new_staked = self
492 .staked
493 .checked_sub(amount)
494 .ok_or_else(|| format!("Insufficient staked balance: {} < {}", self.staked, amount))?;
495 let new_spendable = self.spendable.checked_add(amount).ok_or_else(|| {
496 format!(
497 "Overflow adding {} to spendable balance {}",
498 amount, self.spendable
499 )
500 })?;
501 self.staked = new_staked;
502 self.spendable = new_spendable;
503 if self.spores != self.spendable + self.staked + self.locked {
504 return Err("Account invariant violated after unstake".to_string());
505 }
506 Ok(())
507 }
508
509 pub fn lock(&mut self, amount: u64) -> Result<(), String> {
512 if amount == 0 {
514 return Ok(());
515 }
516 let new_spendable = self.spendable.checked_sub(amount).ok_or_else(|| {
517 format!(
518 "Insufficient spendable balance: {} < {}",
519 self.spendable, amount
520 )
521 })?;
522 let new_locked = self.locked.checked_add(amount).ok_or_else(|| {
523 format!(
524 "Overflow adding {} to locked balance {}",
525 amount, self.locked
526 )
527 })?;
528 self.spendable = new_spendable;
529 self.locked = new_locked;
530 if self.spores != self.spendable + self.staked + self.locked {
531 return Err("Account invariant violated after lock".to_string());
532 }
533 Ok(())
534 }
535
536 pub fn unlock(&mut self, amount: u64) -> Result<(), String> {
539 if amount == 0 {
541 return Ok(());
542 }
543 let new_locked = self
544 .locked
545 .checked_sub(amount)
546 .ok_or_else(|| format!("Insufficient locked balance: {} < {}", self.locked, amount))?;
547 let new_spendable = self.spendable.checked_add(amount).ok_or_else(|| {
548 format!(
549 "Overflow adding {} to spendable balance {}",
550 amount, self.spendable
551 )
552 })?;
553 self.locked = new_locked;
554 self.spendable = new_spendable;
555 if self.spores != self.spendable + self.staked + self.locked {
556 return Err("Account invariant violated after unlock".to_string());
557 }
558 Ok(())
559 }
560
561 pub fn add_spendable(&mut self, amount: u64) -> Result<(), String> {
563 let new_spores = self.spores.checked_add(amount).ok_or_else(|| {
564 format!(
565 "Overflow adding {} to spores balance {}",
566 amount, self.spores
567 )
568 })?;
569 let new_spendable = self.spendable.checked_add(amount).ok_or_else(|| {
570 format!(
571 "Overflow adding {} to spendable balance {}",
572 amount, self.spendable
573 )
574 })?;
575 self.spores = new_spores;
576 self.spendable = new_spendable;
577 Ok(())
578 }
579
580 pub fn deduct_spendable(&mut self, amount: u64) -> Result<(), String> {
583 let new_spendable = self.spendable.checked_sub(amount).ok_or_else(|| {
584 format!(
585 "Insufficient spendable balance: {} < {}",
586 self.spendable, amount
587 )
588 })?;
589 let new_spores = self.spores.checked_sub(amount).ok_or_else(|| {
590 format!(
591 "Underflow subtracting {} from spores balance {}",
592 amount, self.spores
593 )
594 })?;
595 self.spendable = new_spendable;
596 self.spores = new_spores;
597 Ok(())
598 }
599
600 pub fn deduct_locked(&mut self, amount: u64) -> Result<(), String> {
602 let new_locked = self
603 .locked
604 .checked_sub(amount)
605 .ok_or_else(|| format!("Insufficient locked balance: {} < {}", self.locked, amount))?;
606 let new_spores = self.spores.checked_sub(amount).ok_or_else(|| {
607 format!(
608 "Underflow subtracting {} from spores balance {}",
609 amount, self.spores
610 )
611 })?;
612 self.locked = new_locked;
613 self.spores = new_spores;
614 if self.spores != self.spendable + self.staked + self.locked {
615 return Err("Account invariant violated after locked deduction".to_string());
616 }
617 Ok(())
618 }
619
620 pub fn balance_licn(&self) -> u64 {
622 Self::spores_to_licn(self.spores)
623 }
624}
625
626#[cfg(test)]
627mod tests {
628 use super::*;
629
630 #[test]
631 fn test_licn_spores_conversion() {
632 assert_eq!(Account::licn_to_spores(1), 1_000_000_000);
633 assert_eq!(Account::licn_to_spores(100), 100_000_000_000);
634 assert_eq!(Account::spores_to_licn(1_000_000_000), 1);
635 assert_eq!(Account::spores_to_licn(100_000_000_000), 100);
636 }
637
638 #[test]
639 fn test_dual_address_format() {
640 let pubkey = Pubkey([1u8; 32]);
641
642 let base58 = pubkey.to_base58();
644 assert!(!base58.is_empty());
645 println!("Base58: {}", base58);
646
647 let evm = pubkey.to_evm();
649 assert!(evm.starts_with("0x"));
650 assert_eq!(evm.len(), 42); println!("EVM: {}", evm);
652 }
653
654 #[test]
655 fn test_base58_roundtrip() {
656 let original = Pubkey([42u8; 32]);
657 let base58 = original.to_base58();
658 let parsed = Pubkey::from_base58(&base58).unwrap();
659 assert_eq!(original, parsed);
660 }
661
662 #[test]
663 fn test_pq_sign_and_verify_roundtrip() {
664 let keypair = Keypair::new();
665 let message = b"lichen-native-pq";
666 let signature = keypair.sign(message);
667
668 assert!(Keypair::verify(&keypair.pubkey(), message, &signature));
669 assert!(!Keypair::verify(&Pubkey([7u8; 32]), message, &signature));
670 assert!(!Keypair::verify(
671 &keypair.pubkey(),
672 b"different",
673 &signature
674 ));
675 }
676
677 #[test]
678 fn test_slh_verify_roundtrip() {
679 use slh_dsa::signature::Signer;
680
681 let signing_key = slh_dsa::SigningKey::<SlhDsa128f>::slh_keygen_internal(
682 &[1u8; 16], &[2u8; 16], &[3u8; 16],
683 );
684 let message = b"lichen-native-slh";
685 let slh_signature = signing_key.sign(message);
686
687 let public_key =
688 PqPublicKey::new(PQ_SCHEME_SLH_DSA_128F, signing_key.as_ref().to_vec()).unwrap();
689 let signature =
690 PqSignature::new(PQ_SCHEME_SLH_DSA_128F, public_key, slh_signature.to_vec()).unwrap();
691
692 assert!(Keypair::verify(
693 &signature.signer_address(),
694 message,
695 &signature
696 ));
697 assert!(!Keypair::verify(
698 &signature.signer_address(),
699 b"different",
700 &signature,
701 ));
702 }
703}