logos-blockchain/lssa
1
0
Fork 0
mirror of https://github.com/logos-blockchain/lssa.git synced 2026-01-02 13:23:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into Pravdyvy/sparse-merkle-tree-storage-preparation

Browse Source
This commit is contained in:
Oleksandr Pravdyvyi 2025-05-21 08:31:15 +03:00
commit 77062e2e4e
8 changed files with 50 additions and 375 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
accounts/src/account_core/mod.rs
View File

@ -5,10 +5,7 @@ use common::{merkle_tree_public::TreeHashType, nullifier::UTXONullifier, transac
use k256::AffinePoint;
use log::info;
use serde::Serialize;
use utxo::{
utxo_core::{UTXOPayload, UTXO},
utxo_tree::{UTXOSparseMerkleTree, UTXOTreeInput},
};
use utxo::utxo_core::{UTXOPayload, UTXO};
use crate::key_management::{
constants_types::{CipherText, Nonce},
@ -23,7 +20,7 @@ pub struct Account {
pub key_holder: AddressKeyHolder,
pub address: AccountAddress,
pub balance: u64,
pub utxo_tree: UTXOSparseMerkleTree,
pub utxos: HashMap<TreeHashType, UTXO>,
}
///A strucure, which represents all the visible(public) information
@ -42,26 +39,26 @@ impl Account {
let key_holder = AddressKeyHolder::new_os_random();
let address = key_holder.address;
let balance = 0;
let utxo_tree = UTXOSparseMerkleTree::new();
let utxos = HashMap::new();
Self {
key_holder,
address,
balance,
utxo_tree,
utxos,
}
}
pub fn new_with_balance(balance: u64) -> Self {
let key_holder = AddressKeyHolder::new_os_random();
let address = key_holder.address;
let utxo_tree = UTXOSparseMerkleTree::new();
let utxos = HashMap::new();
Self {
key_holder,
address,
balance,
utxo_tree,
utxos,
}
}
@ -92,7 +89,7 @@ impl Account {
utxo_nullifier_map: HashMap<TreeHashType, UTXONullifier>,
) -> Result<()> {
for (hash, nullifier) in utxo_nullifier_map {
if let Some(utxo_entry) = self.utxo_tree.store.get_mut(&hash) {
if let Some(utxo_entry) = self.utxos.get_mut(&hash) {
utxo_entry.consume_utxo(nullifier)?;
}
}
@ -100,8 +97,14 @@ impl Account {
Ok(())
}
pub fn add_new_utxo_outputs(&mut self, utxos: Vec<UTXOTreeInput>) -> Result<()> {
Ok(self.utxo_tree.insert_items(utxos)?)
pub fn add_new_utxo_outputs(&mut self, utxos: Vec<UTXO>) -> Result<()> {
for utxo in utxos {
if self.utxos.contains_key(&utxo.hash) {
return Err(anyhow::anyhow!("UTXO already exists"));
}
self.utxos.insert(utxo.hash, utxo);
}
return Ok(());
}
pub fn update_public_balance(&mut self, new_balance: u64) {
@ -126,14 +129,7 @@ impl Account {
let asset_utxo = UTXO::create_utxo_from_payload(payload_with_asset)?;
let input_utxo = UTXOTreeInput {
utxo_id,
tx_id,
block_id,
utxo: asset_utxo,
};
self.utxo_tree.insert_item(input_utxo)?;
self.utxos.insert(asset_utxo.hash, asset_utxo);
Ok(())
}
@ -213,7 +209,6 @@ mod tests {
let result = account.mark_spent_utxo(utxo_nullifier_map);
assert!(result.is_ok());
assert!(account.utxo_tree.store.get(&account.address).is_none());
}
#[test]
@ -225,7 +220,7 @@ mod tests {
let result = account.add_new_utxo_outputs(vec![utxo1.clone(), utxo2.clone()]);
assert!(result.is_ok());
assert_eq!(account.utxo_tree.store.len(), 2);
assert_eq!(account.utxos.len(), 2);
}
#[test]
@ -245,6 +240,6 @@ mod tests {
let result = account.add_asset(asset, amount, false, 1, 1, 1);
assert!(result.is_ok());
assert_eq!(account.utxo_tree.store.len(), 1);
assert_eq!(account.utxos.len(), 1);
}
}

13
accounts/src/key_management/ephemeral_key_holder.rs
View File

@ -1,4 +1,5 @@
use aes_gcm::{aead::Aead, AeadCore, Aes256Gcm, Key, KeyInit};
use aes_gcm::{aead::Aead, AeadCore, Aes256Gcm, KeyInit};
use elliptic_curve::point::AffineCoordinates;
use elliptic_curve::PrimeField;
use k256::{AffinePoint, FieldBytes, Scalar};
use log::info;
@ -39,14 +40,8 @@ impl EphemeralKeyHolder {
viewing_public_key_receiver: AffinePoint,
data: &[u8],
) -> (CipherText, Nonce) {
let key_point = self.calculate_shared_secret_sender(viewing_public_key_receiver);
let binding = serde_json::to_vec(&key_point).unwrap();
let key_raw = &binding.as_slice()[..32];
let key_raw_adjust: [u8; 32] = key_raw.try_into().unwrap();
let key: Key<Aes256Gcm> = key_raw_adjust.into();
let cipher = Aes256Gcm::new(&key);
let shared_secret = self.calculate_shared_secret_sender(viewing_public_key_receiver);
let cipher = Aes256Gcm::new(&shared_secret.x());
let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
(cipher.encrypt(&nonce, data).unwrap(), nonce)

55
accounts/src/key_management/mod.rs
View File

@ -1,6 +1,7 @@
use aes_gcm::{aead::Aead, Aes256Gcm, Key, KeyInit};
use aes_gcm::{aead::Aead, Aes256Gcm, KeyInit};
use common::merkle_tree_public::TreeHashType;
use constants_types::{CipherText, Nonce};
use elliptic_curve::point::AffineCoordinates;
use ephemeral_key_holder::EphemeralKeyHolder;
use k256::AffinePoint;
use log::info;
@ -63,14 +64,8 @@ impl AddressKeyHolder {
ciphertext: CipherText,
nonce: Nonce,
) -> Result<Vec<u8>, aes_gcm::Error> {
let key_point = self.calculate_shared_secret_receiver(ephemeral_public_key_sender);
let binding = serde_json::to_vec(&key_point).unwrap();
let key_raw = &binding.as_slice()[..32];
let key_raw_adjust: [u8; 32] = key_raw.try_into().unwrap();
let key: Key<Aes256Gcm> = key_raw_adjust.into();
let cipher = Aes256Gcm::new(&key);
let shared_secret = self.calculate_shared_secret_receiver(ephemeral_public_key_sender);
let cipher = Aes256Gcm::new(&shared_secret.x());
cipher.decrypt(&nonce, ciphertext.as_slice())
}
@ -115,6 +110,7 @@ mod tests {
use constants_types::{NULLIFIER_SECRET_CONST, VIEWING_SECRET_CONST};
use elliptic_curve::ff::Field;
use elliptic_curve::group::prime::PrimeCurveAffine;
use elliptic_curve::point::AffineCoordinates;
use k256::{AffinePoint, ProjectivePoint, Scalar};
use super::*;
@ -154,22 +150,14 @@ mod tests {
let address_key_holder = AddressKeyHolder::new_os_random();
// Generate an ephemeral key and shared secret
let scalar = Scalar::random(OsRng);
let ephemeral_public_key_sender = address_key_holder
.produce_ephemeral_key_holder()
.generate_ephemeral_public_key();
let shared_secret =
address_key_holder.calculate_shared_secret_receiver(ephemeral_public_key_sender);
// Prepare the encryption key from shared secret
let key_raw = serde_json::to_vec(&shared_secret).unwrap();
let key_raw_adjust_pre = &key_raw.as_slice()[..32];
let key_raw_adjust: [u8; 32] = key_raw_adjust_pre.try_into().unwrap();
let key: Key<Aes256Gcm> = key_raw_adjust.into();
let cipher = Aes256Gcm::new(&key);
// Encrypt sample data
let cipher = Aes256Gcm::new(&shared_secret.x());
let nonce = Nonce::from_slice(b"unique nonce");
let plaintext = b"Sensitive data";
let ciphertext = cipher
@ -225,19 +213,12 @@ mod tests {
// Generate ephemeral public key and shared secret
let scalar = Scalar::random(OsRng);
let ephemeral_public_key_sender = (ProjectivePoint::generator() * scalar).to_affine();
let ephemeral_public_key_sender = (ProjectivePoint::GENERATOR * scalar).to_affine();
let shared_secret =
address_key_holder.calculate_shared_secret_receiver(ephemeral_public_key_sender);
// Prepare the encryption key from shared secret
let key_raw = serde_json::to_vec(&shared_secret).unwrap();
let key_raw_adjust_pre = &key_raw.as_slice()[..32];
let key_raw_adjust: [u8; 32] = key_raw_adjust_pre.try_into().unwrap();
let key: Key<Aes256Gcm> = key_raw_adjust.into();
let cipher = Aes256Gcm::new(&key);
// Encrypt sample data with a specific nonce
let cipher = Aes256Gcm::new(&shared_secret.x());
let nonce = Nonce::from_slice(b"unique nonce");
let plaintext = b"Sensitive data";
let ciphertext = cipher
@ -265,19 +246,12 @@ mod tests {
// Generate ephemeral public key and shared secret
let scalar = Scalar::random(OsRng);
let ephemeral_public_key_sender = (ProjectivePoint::generator() * scalar).to_affine();
let ephemeral_public_key_sender = (ProjectivePoint::GENERATOR * scalar).to_affine();
let shared_secret =
address_key_holder.calculate_shared_secret_receiver(ephemeral_public_key_sender);
// Prepare the encryption key from shared secret
let key_raw = serde_json::to_vec(&shared_secret).unwrap();
let key_raw_adjust_pre = &key_raw.as_slice()[..32];
let key_raw_adjust: [u8; 32] = key_raw_adjust_pre.try_into().unwrap();
let key: Key<Aes256Gcm> = key_raw_adjust.into();
let cipher = Aes256Gcm::new(&key);
// Encrypt sample data
let cipher = Aes256Gcm::new(&shared_secret.x());
let nonce = Nonce::from_slice(b"unique nonce");
let plaintext = b"Sensitive data";
let ciphertext = cipher
@ -307,7 +281,7 @@ mod tests {
// Generate ephemeral key and shared secret
let scalar = Scalar::random(OsRng);
let ephemeral_public_key_sender = (ProjectivePoint::generator() * scalar).to_affine();
let ephemeral_public_key_sender = (ProjectivePoint::GENERATOR * scalar).to_affine();
// Encrypt sample data
let plaintext = b"Round-trip test data";
@ -315,12 +289,7 @@ mod tests {
let shared_secret =
address_key_holder.calculate_shared_secret_receiver(ephemeral_public_key_sender);
// Prepare the encryption key from shared secret
let key_raw = serde_json::to_vec(&shared_secret).unwrap();
let key_raw_adjust_pre = &key_raw.as_slice()[..32];
let key_raw_adjust: [u8; 32] = key_raw_adjust_pre.try_into().unwrap();
let key: Key<Aes256Gcm> = key_raw_adjust.into();
let cipher = Aes256Gcm::new(&key);
let cipher = Aes256Gcm::new(&shared_secret.x());
let ciphertext = cipher
.encrypt(nonce, plaintext.as_ref())

8
node_core/src/chain_storage/mod.rs
View File

@ -133,13 +133,7 @@ impl NodeChainStore {
serde_json::from_slice::<UTXO>(&decoded_data_curr_acc);
if let Ok(utxo) = decoded_utxo_try {
if &utxo.owner == acc_id {
let input_utxo = UTXOTreeInput {
utxo_id: utxo_id as u64,
tx_id: tx_id as u64,
block_id: block.block_id,
utxo,
};
acc.utxo_tree.insert_item(input_utxo)?;
acc.utxos.insert(utxo.hash, utxo);
}
}
}

24
node_core/src/lib.rs
View File

@ -1062,7 +1062,7 @@ impl NodeCore {
let acc = write_guard.acc_map.get_mut(&acc_addr).unwrap();
acc.utxo_tree.get_item(new_utxo_hash)?.unwrap().clone()
acc.utxos.get(&new_utxo_hash).unwrap().clone()
};
new_utxo.log();
@ -1101,12 +1101,7 @@ impl NodeCore {
.map(|new_utxo_hash| {
let acc = write_guard.acc_map.get_mut(&acc_addr).unwrap();
let new_utxo = acc
.utxo_tree
.get_item(new_utxo_hash)
.unwrap()
.unwrap()
.clone();
let new_utxo = acc.utxos.get(&new_utxo_hash).unwrap().clone();
new_utxo.log();
info!(
@ -1238,7 +1233,7 @@ impl NodeCore {
let acc = write_guard.acc_map.get_mut(&acc_addr_rec).unwrap();
acc.log();
acc.utxo_tree.get_item(new_utxo_hash)?.unwrap().clone()
acc.utxos.get(&new_utxo_hash).unwrap().clone()
};
new_utxo.log();
info!(
@ -1278,7 +1273,7 @@ impl NodeCore {
let acc = write_guard.acc_map.get_mut(&acc_addr_rec).unwrap();
acc.log();
acc.utxo_tree.get_item(new_utxo_hash)?.unwrap().clone()
acc.utxos.get(&new_utxo_hash).unwrap().clone()
};
new_utxo.log();
info!(
@ -1323,7 +1318,7 @@ impl NodeCore {
let acc = write_guard.acc_map.get_mut(&acc_addr_rec).unwrap();
acc.log();
let new_utxo = acc.utxo_tree.get_item(new_utxo_hash)?.unwrap().clone();
let new_utxo = acc.utxos.get(&new_utxo_hash).unwrap().clone();
new_utxo.log();
info!(
@ -1343,7 +1338,7 @@ impl NodeCore {
let acc = write_guard.acc_map.get_mut(&acc_addr).unwrap();
acc.log();
let new_utxo = acc.utxo_tree.get_item(new_utxo_hash)?.unwrap().clone();
let new_utxo = acc.utxos.get(&new_utxo_hash).unwrap().clone();
new_utxo.log();
info!(
@ -1556,12 +1551,7 @@ impl NodeCore {
.map(|(acc_addr_rec, new_utxo_hash)| {
let acc = write_guard.acc_map.get_mut(&acc_addr_rec).unwrap();
let new_utxo = acc
.utxo_tree
.get_item(new_utxo_hash)
.unwrap()
.unwrap()
.clone();
let new_utxo = acc.utxos.get(&new_utxo_hash).unwrap().clone();
new_utxo.log();
info!(

28
node_rpc/src/process.rs
View File

@ -268,11 +268,8 @@ impl JsonHandler {
.ok_or(RpcError::new_internal_error(None, ACCOUNT_NOT_FOUND))?;
let utxo = acc
.utxo_tree
.get_item(utxo_hash)
.map_err(|err| {
RpcError::new_internal_error(None, &format!("DB fetch failure {err:?}"))
})?
.utxos
.get(&utxo_hash)
.ok_or(RpcError::new_internal_error(
None,
"UTXO does not exist in the tree",
@ -512,11 +509,8 @@ impl JsonHandler {
.get_mut(&acc_addr_sender)
.ok_or(RpcError::new_internal_error(None, ACCOUNT_NOT_FOUND))?;
acc.utxo_tree
.get_item(utxo_hash)
.map_err(|err| {
RpcError::new_internal_error(None, &format!("DB fetch failure {err:?}"))
})?
acc.utxos
.get(&utxo_hash)
.ok_or(RpcError::new_internal_error(
None,
"UTXO does not exist in tree",
@ -647,11 +641,8 @@ impl JsonHandler {
.get_mut(&acc_addr_sender)
.ok_or(RpcError::new_internal_error(None, ACCOUNT_NOT_FOUND))?;
acc.utxo_tree
.get_item(utxo_hash)
.map_err(|err| {
RpcError::new_internal_error(None, &format!("DB fetch failure {err:?}"))
})?
acc.utxos
.get(&utxo_hash)
.ok_or(RpcError::new_internal_error(
None,
"UTXO does not exist in tree",
@ -735,11 +726,8 @@ impl JsonHandler {
.get_mut(&acc_addr_sender)
.ok_or(RpcError::new_internal_error(None, ACCOUNT_NOT_FOUND))?;
acc.utxo_tree
.get_item(utxo_hash)
.map_err(|err| {
RpcError::new_internal_error(None, &format!("DB fetch failure {err:?}"))
})?
acc.utxos
.get(&utxo_hash)
.ok_or(RpcError::new_internal_error(
None,
"UTXO does not exist in tree",

1
utxo/src/lib.rs
View File

@ -1,2 +1 @@
pub mod utxo_core;
pub mod utxo_tree;

255
utxo/src/utxo_tree.rs
View File

@ -1,255 +0,0 @@
use std::collections::{BTreeMap, HashMap};
use common::merkle_tree_public::TreeHashType;
use monotree::database::MemoryDB;
use monotree::hasher::Blake3;
use monotree::{Hasher, Monotree, Proof};
use crate::utxo_core::UTXO;
#[derive(Debug, Clone)]
pub struct UTXOTreeInput {
pub utxo_id: u64,
pub tx_id: u64,
pub block_id: u64,
pub utxo: UTXO,
}
#[derive(Debug, Clone)]
pub struct TreeTxWithUTXOId {
pub id: u64,
pub utxos: BTreeMap<u64, UTXO>,
}
#[derive(Debug, Clone)]
pub struct TreeBlockWithTxId {
pub id: u64,
pub txs: BTreeMap<u64, TreeTxWithUTXOId>,
}
pub struct UTXOSparseMerkleTree {
pub curr_root: Option<TreeHashType>,
pub tree: Monotree<MemoryDB, Blake3>,
pub hasher: Blake3,
pub store: HashMap<TreeHashType, UTXO>,
pub leafs: BTreeMap<u64, TreeBlockWithTxId>,
}
impl UTXOSparseMerkleTree {
pub fn new() -> Self {
UTXOSparseMerkleTree {
curr_root: None,
tree: Monotree::default(),
hasher: Blake3::new(),
store: HashMap::new(),
leafs: BTreeMap::new(),
}
}
pub fn modify_leavs_with_nullifier_input(&mut self, tree_utxo: UTXOTreeInput) {
self.leafs
.entry(tree_utxo.block_id)
.and_modify(|tree_block| {
tree_block
.txs
.entry(tree_utxo.tx_id)
.and_modify(|tree_tx| {
tree_tx.utxos.insert(tree_utxo.utxo_id, tree_utxo.utxo);
})
.or_insert(TreeTxWithUTXOId {
id: tree_utxo.tx_id,
utxos: BTreeMap::new(),
});
})
.or_insert(TreeBlockWithTxId {
id: tree_utxo.block_id,
txs: BTreeMap::new(),
});
}
pub fn insert_item(&mut self, tree_utxo: UTXOTreeInput) -> Result<(), monotree::Errors> {
let root = self.curr_root.as_ref();
let new_root = self
.tree
.insert(root, &tree_utxo.utxo.hash, &tree_utxo.utxo.hash)?;
self.curr_root = new_root;
self.store
.insert(tree_utxo.utxo.hash, tree_utxo.utxo.clone());
self.modify_leavs_with_nullifier_input(tree_utxo);
Ok(())
}
pub fn insert_items(&mut self, tree_utxos: Vec<UTXOTreeInput>) -> Result<(), monotree::Errors> {
let root = self.curr_root.as_ref();
let hashes: Vec<TreeHashType> = tree_utxos.iter().map(|item| item.utxo.hash).collect();
let new_root = self.tree.inserts(root, &hashes, &hashes)?;
for tree_utxo in tree_utxos {
self.store
.insert(tree_utxo.utxo.hash, tree_utxo.utxo.clone());
self.modify_leavs_with_nullifier_input(tree_utxo);
}
self.curr_root = new_root;
Ok(())
}
pub fn get_item(&mut self, hash: TreeHashType) -> Result<Option<&UTXO>, monotree::Errors> {
let hash = self.tree.get(self.curr_root.as_ref(), &hash)?;
Ok(hash.and_then(|hash| self.store.get(&hash)))
}
pub fn get_membership_proof(
&mut self,
nullifier_hash: TreeHashType,
) -> Result<Option<Proof>, monotree::Errors> {
self.tree
.get_merkle_proof(self.curr_root.as_ref(), &nullifier_hash)
}
}
impl Default for UTXOSparseMerkleTree {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use common::AccountId;
use super::*;
use crate::utxo_core::{UTXOPayload, UTXO};
fn sample_utxo_payload(amount: u128) -> UTXOPayload {
UTXOPayload {
owner: AccountId::default(),
asset: vec![1, 2, 3],
amount,
privacy_flag: false,
}
}
fn sample_utxo(amount: u128) -> anyhow::Result<UTXO> {
UTXO::create_utxo_from_payload(sample_utxo_payload(amount))
}
fn sample_utxo_input(
utxo_id: u64,
tx_id: u64,
block_id: u64,
amount: u128,
) -> anyhow::Result<UTXOTreeInput> {
sample_utxo(amount).map(|utxo| UTXOTreeInput {
utxo_id,
tx_id,
block_id,
utxo,
})
}
#[test]
fn test_utxo_sparse_merkle_tree_new() {
let smt = UTXOSparseMerkleTree::new();
assert!(smt.curr_root.is_none());
assert_eq!(smt.store.len(), 0);
}
#[test]
fn test_insert_item() {
let mut smt = UTXOSparseMerkleTree::new();
let utxo = sample_utxo_input(1, 1, 1, 10).unwrap();
let result = smt.insert_item(utxo.clone());
// Test insertion is successful
assert!(result.is_ok());
// Test UTXO is now stored in the tree
assert_eq!(smt.store.get(&utxo.utxo.hash).unwrap().hash, utxo.utxo.hash);
// Test curr_root is updated
assert!(smt.curr_root.is_some());
}
#[test]
fn test_insert_items() {
let mut smt = UTXOSparseMerkleTree::new();
let utxo1 = sample_utxo_input(1, 1, 1, 10).unwrap();
let utxo2 = sample_utxo_input(2, 1, 1, 11).unwrap();
let result = smt.insert_items(vec![utxo1.clone(), utxo2.clone()]);
// Test insertion of multiple items is successful
assert!(result.is_ok());
// Test UTXOs are now stored in the tree
assert!(smt.store.get(&utxo1.utxo.hash).is_some());
assert!(smt.store.get(&utxo2.utxo.hash).is_some());
// Test curr_root is updated
assert!(smt.curr_root.is_some());
}
#[test]
fn test_get_item_exists() {
let mut smt = UTXOSparseMerkleTree::new();
let utxo = sample_utxo_input(1, 1, 1, 10).unwrap();
smt.insert_item(utxo.clone()).unwrap();
// Test that the UTXO can be retrieved by hash
let retrieved_utxo = smt.get_item(utxo.utxo.hash).unwrap();
assert!(retrieved_utxo.is_some());
assert_eq!(retrieved_utxo.unwrap().hash, utxo.utxo.hash);
}
#[test]
fn test_get_item_not_exists() {
let mut smt = UTXOSparseMerkleTree::new();
let utxo = sample_utxo_input(1, 1, 1, 10).unwrap();
// Insert one UTXO and try to fetch a different hash
smt.insert_item(utxo).unwrap();
let non_existent_hash = TreeHashType::default();
let result = smt.get_item(non_existent_hash).unwrap();
// Test that retrieval for a non-existent UTXO returns None
assert!(result.is_none());
}
#[test]
fn test_get_membership_proof() {
let mut smt = UTXOSparseMerkleTree::new();
let utxo = sample_utxo_input(1, 1, 1, 10).unwrap();
smt.insert_item(utxo.clone()).unwrap();
// Fetch membership proof for the inserted UTXO
let proof = smt.get_membership_proof(utxo.utxo.hash).unwrap();
// Test proof is generated successfully
assert!(proof.is_some());
}
#[test]
fn test_get_membership_proof_not_exists() {
let mut smt = UTXOSparseMerkleTree::new();
// Try fetching proof for a non-existent UTXO hash
let non_existent_hash = TreeHashType::default();
let proof = smt.get_membership_proof(non_existent_hash).unwrap();
// Test no proof is generated for a non-existent UTXO
assert!(proof.is_none());
}
}