Skip to main content

lichen_core/
account.rs

1// Lichen Core - Account Model
2// Based on Solana's account model with versioned PQ addresses
3
4use 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/// Versioned 32-byte address digest.
47#[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    /// Convert to Base58 string (native Lichen format)
64    pub fn to_base58(&self) -> String {
65        bs58::encode(self.0).into_string()
66    }
67
68    /// Convert to EVM-compatible hex address (0x...)
69    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]; // Last 20 bytes
73        format!("0x{}", hex::encode(evm_bytes))
74    }
75
76    /// Parse from Base58 string
77    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
254/// ML-DSA-65 keypair for signing native Lichen transactions.
255pub struct Keypair {
256    keypair: MlDsaSigningKey<MlDsa65>,
257    seed: [u8; 32],
258}
259
260impl Keypair {
261    /// Generate new random keypair
262    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    /// Alias for new() - generates random keypair
269    pub fn generate() -> Self {
270        Self::new()
271    }
272
273    /// Get secret key bytes (for serialization)
274    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    /// Create from seed bytes
283    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    /// Get account address.
296    pub fn pubkey(&self) -> Pubkey {
297        self.public_key().address()
298    }
299
300    /// Get the full PQ public key used for verification.
301    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    /// Get seed bytes (for saving to file)
308    pub fn to_seed(&self) -> [u8; 32] {
309        self.seed
310    }
311
312    /// Sign message with ML-DSA-65 and embed the verifying key.
313    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    /// Verify a native PQ signature against an address.
329    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/// Account structure with balance separation
374#[derive(Debug, Clone, Serialize, Deserialize)]
375pub struct Account {
376    /// Total balance in spores (1 LICN = 1_000_000_000 spores)
377    /// Total = spendable + staked + locked
378    pub spores: u64,
379
380    /// Spendable balance (available for transfers)
381    #[serde(default)]
382    pub spendable: u64,
383
384    /// Staked balance (locked in validator staking)
385    #[serde(default)]
386    pub staked: u64,
387
388    /// Locked balance (locked in contracts, escrow, multisig)
389    #[serde(default)]
390    pub locked: u64,
391
392    /// Arbitrary data storage
393    pub data: Vec<u8>,
394
395    /// Optional cache of the first PQ public key observed for this account.
396    #[serde(default)]
397    pub public_key: Option<PqPublicKey>,
398
399    /// Program that owns this account
400    pub owner: Pubkey,
401
402    /// Is this account an executable program?
403    pub executable: bool,
404
405    /// Last epoch when rent was assessed
406    pub rent_epoch: u64,
407
408    /// Whether this account is dormant (excluded from active state root)
409    #[serde(default)]
410    pub dormant: bool,
411
412    /// Consecutive epochs where rent could not be fully paid
413    #[serde(default)]
414    pub missed_rent_epochs: u64,
415}
416
417impl Account {
418    /// M11 fix: repair legacy accounts where spendable/staked/locked are all 0 but spores > 0.
419    /// This happens when deserializing accounts created before the balance separation fields existed.
420    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    /// Convert LICN to spores
427    pub const fn licn_to_spores(licn: u64) -> u64 {
428        licn.saturating_mul(1_000_000_000)
429    }
430
431    /// Convert spores to LICN (integer division — truncates fractional LICN).
432    /// AUDIT-FIX 3.2: Callers needing rounding should use
433    /// `(spores + 999_999_999) / 1_000_000_000` for round-up.
434    pub const fn spores_to_licn(spores: u64) -> u64 {
435        spores / 1_000_000_000
436    }
437
438    /// Create a new account with LICN balance (all spendable)
439    pub fn new(licn: u64, owner: Pubkey) -> Self {
440        let spores = Self::licn_to_spores(licn);
441        Account {
442            spores,
443            spendable: spores, // All balance is spendable initially
444            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    /// Stake some balance (moves from spendable to staked)
457    /// T3.3 fix: spores total is unchanged (just a reclassification)
458    /// AUDIT-FIX 1.1a: checked arithmetic, compute-before-mutate
459    pub fn stake(&mut self, amount: u64) -> Result<(), String> {
460        // AUDIT-FIX 3.1: Skip no-op zero-amount operations
461        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    /// Unstake balance (moves from staked to spendable)
485    /// AUDIT-FIX 1.1b: checked arithmetic, compute-before-mutate
486    pub fn unstake(&mut self, amount: u64) -> Result<(), String> {
487        // AUDIT-FIX 3.1: Skip no-op zero-amount operations
488        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    /// Lock balance (moves from spendable to locked)
510    /// AUDIT-FIX 1.1c: checked arithmetic, compute-before-mutate
511    pub fn lock(&mut self, amount: u64) -> Result<(), String> {
512        // AUDIT-FIX 3.1: Skip no-op zero-amount operations
513        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    /// Unlock balance (moves from locked to spendable)
537    /// AUDIT-FIX 1.1d: checked arithmetic, compute-before-mutate
538    pub fn unlock(&mut self, amount: u64) -> Result<(), String> {
539        // AUDIT-FIX 3.1: Skip no-op zero-amount operations
540        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    /// Add to spendable balance (for rewards, transfers)
562    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    /// Deduct from spendable balance (for transfers, fees)
581    /// AUDIT-FIX 1.1e: checked arithmetic, compute-before-mutate
582    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    /// Deduct from locked balance (burns/removes locked collateral)
601    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    /// Get balance in LICN
621    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        // Base58 format
643        let base58 = pubkey.to_base58();
644        assert!(!base58.is_empty());
645        println!("Base58: {}", base58);
646
647        // EVM format
648        let evm = pubkey.to_evm();
649        assert!(evm.starts_with("0x"));
650        assert_eq!(evm.len(), 42); // 0x + 40 hex chars
651        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}