vacp2p/status-rln-prover
1
0
Fork 0
mirror of https://github.com/vacp2p/status-rln-prover.git synced 2026-01-08 21:18:05 -05:00
Code Issues Packages Projects Releases Wiki Activity

Split unit tests in several new files (#21)

Browse Source
This commit is contained in:
Sydhds
2025-07-21 16:36:54 +02:00
committed by GitHub
parent 8b1540cf76
commit 1ec9f1f48d
12 changed files with 866 additions and 792 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
prover/src/epoch_service.rs
View File

@@ -17,7 +17,7 @@ use crate::metrics::{
/// Duration of an epoch (1 day)
const EPOCH_DURATION: Duration = Duration::from_secs(TimeDelta::days(1).num_seconds() as u64);
/// Minimum duration returned by EpochService::compute_wait_until()
const WAIT_UNTIL_MIN_DURATION: Duration = Duration::from_secs(2);
pub(crate) const WAIT_UNTIL_MIN_DURATION: Duration = Duration::from_secs(2);
/// EpochService::compute_wait_until() can return an error like TooLow (see WAIT_UNTIL_MIN_DURATION)
/// so the epoch service will retry X many times.
const WAIT_UNTIL_MAX_COMPUTE_ERROR: usize = 10;
@@ -287,10 +287,6 @@ impl Add<i64> for EpochSlice {
mod tests {
use super::*;
use chrono::{NaiveDate, NaiveDateTime, TimeDelta};
use claims::assert_ge;
use futures::TryFutureExt;
use std::sync::atomic::{AtomicU64, Ordering};
use tracing_test::traced_test;
#[test]
fn test_wait_until() {
@@ -476,53 +472,4 @@ mod tests {
3
);
}
#[derive(thiserror::Error, Debug)]
enum AppErrorExt {
#[error("AppError: {0}")]
AppError(#[from] AppError),
#[error("Future timeout")]
Elapsed,
}
#[tokio::test]
#[traced_test]
async fn test_notify() {
// Test epoch_service is really notifying when an epoch or epoch slice has just changed
let epoch_slice_duration = Duration::from_secs(10);
let epoch_service = EpochService::try_from((epoch_slice_duration, Utc::now())).unwrap();
let notifier = epoch_service.epoch_changes.clone();
let counter_0 = Arc::new(AtomicU64::new(0));
let counter = counter_0.clone();
let start = std::time::Instant::now();
let res = tokio::try_join!(
epoch_service
.listen_for_new_epoch()
.map_err(|e| AppErrorExt::AppError(e)),
// Wait for 3 epoch slices
// + WAIT_UNTIL_MIN_DURATION * 2 (expect a maximum of 2 retry)
// + 500 ms (to wait to receive notif + counter incr)
// Note: this might fail if there is more retry (see list_for_new_epoch code)
tokio::time::timeout(
epoch_slice_duration * 3 + WAIT_UNTIL_MIN_DURATION * 2 + Duration::from_millis(500),
async move {
loop {
notifier.notified().await;
// debug!("[Notified] Epoch update...");
let _v = counter.fetch_add(1, Ordering::SeqCst);
}
// Ok::<(), AppErrorExt>(())
}
)
.map_err(|_e| AppErrorExt::Elapsed)
);
debug!("Elapsed time: {}", start.elapsed().as_millis());
// debug!("res: {:?}", res);
assert!(matches!(res, Err(AppErrorExt::Elapsed)));
// Because the timeout is quite large - it is expected that sometimes the counter == 4
assert_ge!(counter_0.load(Ordering::SeqCst), 3);
}
}

66
prover/src/epoch_service_tests.rs Normal file
View File

@@ -0,0 +1,66 @@
#[cfg(test)]
mod epoch_service_tests {
// std
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
// third-party
use chrono::Utc;
use claims::assert_ge;
use futures::TryFutureExt;
use tracing::debug;
use tracing_test::traced_test;
// internal
use crate::epoch_service::{EpochService, WAIT_UNTIL_MIN_DURATION};
use crate::error::AppError;
#[derive(thiserror::Error, Debug)]
enum AppErrorExt {
#[error("AppError: {0}")]
AppError(#[from] AppError),
#[error("Future timeout")]
Elapsed,
}
#[tokio::test]
#[traced_test]
async fn test_notify() {
// Test epoch_service is really notifying when an epoch or epoch slice has just changed
let epoch_slice_duration = Duration::from_secs(10);
let epoch_service = EpochService::try_from((epoch_slice_duration, Utc::now())).unwrap();
let notifier = epoch_service.epoch_changes.clone();
let counter_0 = Arc::new(AtomicU64::new(0));
let counter = counter_0.clone();
let start = std::time::Instant::now();
let res = tokio::try_join!(
epoch_service
.listen_for_new_epoch()
.map_err(|e| AppErrorExt::AppError(e)),
// Wait for 3 epoch slices
// + WAIT_UNTIL_MIN_DURATION * 2 (expect a maximum of 2 retry)
// + 500 ms (to wait to receive notif + counter incr)
// Note: this might fail if there is more retry (see list_for_new_epoch code)
tokio::time::timeout(
epoch_slice_duration * 3 + WAIT_UNTIL_MIN_DURATION * 2 + Duration::from_millis(500),
async move {
loop {
notifier.notified().await;
// debug!("[Notified] Epoch update...");
let _v = counter.fetch_add(1, Ordering::SeqCst);
}
// Ok::<(), AppErrorExt>(())
}
)
.map_err(|_e| AppErrorExt::Elapsed)
);
debug!("Elapsed time: {}", start.elapsed().as_millis());
// debug!("res: {:?}", res);
assert!(matches!(res, Err(AppErrorExt::Elapsed)));
// Because the timeout is quite large - it is expected that sometimes the counter == 4
assert_ge!(counter_0.load(Ordering::SeqCst), 3);
}
}

67
prover/src/grpc_service.rs
View File

@@ -286,48 +286,6 @@ where
})),
}
}
/*
async fn set_tier_limits(
&self,
request: Request<SetTierLimitsRequest>,
) -> Result<Response<SetTierLimitsReply>, Status> {
debug!("request: {:?}", request);
let request = request.into_inner();
let tier_limits: Option<BTreeMap<KarmaAmount, (TierLimit, TierName)>> = request
.karma_amounts
.iter()
.zip(request.tiers)
.map(|(k, tier)| {
let karma_amount = U256::try_from_le_slice(k.value.as_slice())?;
let karma_amount = KarmaAmount::from(karma_amount);
let tier_info = (
TierLimit::from(tier.quota),
TierName::from(tier.name.clone()),
);
Some((karma_amount, tier_info))
})
.collect();
if tier_limits.is_none() {
return Err(Status::invalid_argument("Invalid tier limits"));
}
// unwrap safe - just tested if None
let reply = match self.user_db.on_new_tier_limits(tier_limits.unwrap()) {
Ok(_) => SetTierLimitsReply {
status: true,
error: "".to_string(),
},
Err(e) => SetTierLimitsReply {
status: false,
error: e.to_string(),
},
};
Ok(Response::new(reply))
}
*/
}
pub(crate) struct GrpcProverService {
@@ -513,28 +471,3 @@ where
}
}
}
#[cfg(test)]
mod tests {
use crate::grpc_service::prover_proto::Address;
use prost::Message;
const MAX_ADDRESS_SIZE_BYTES: usize = 20;
#[test]
#[ignore]
#[should_panic]
fn test_address_size_limit() {
// Check if an invalid address can be encoded (as Address grpc type)
let invalid_address = vec![0; MAX_ADDRESS_SIZE_BYTES + 1];
let addr = Address {
value: invalid_address,
};
let mut addr_encoded = vec![];
addr.encode(&mut addr_encoded).unwrap();
let _addr_decoded = Address::decode(&*addr_encoded).unwrap();
}
}

268
prover/src/lib.rs
View File

@@ -17,6 +17,11 @@ mod user_db_serialization;
mod user_db_service;
mod user_db_types;
// tests
mod epoch_service_tests;
mod proof_service_tests;
mod user_db_tests;
// std
use std::net::SocketAddr;
use std::time::Duration;
@@ -86,29 +91,29 @@ pub async fn run_prover(
tier_limits,
)?;
if app_args.mock_sc.is_some() {
if let Some(user_filepath) = app_args.mock_user.as_ref() {
let mock_users = read_mock_user(user_filepath)?;
debug!("Mock - will register {} users", mock_users.len());
for mock_user in mock_users {
debug!(
"Registering user address: {} - tx count: {}",
mock_user.address, mock_user.tx_count
);
if app_args.mock_sc.is_some()
&& let Some(user_filepath) = app_args.mock_user.as_ref()
{
let mock_users = read_mock_user(user_filepath)?;
debug!("Mock - will register {} users", mock_users.len());
for mock_user in mock_users {
debug!(
"Registering user address: {} - tx count: {}",
mock_user.address, mock_user.tx_count
);
let user_db = user_db_service.get_user_db();
if let Err(e) = user_db.on_new_user(&mock_user.address) {
match e {
RegisterError::AlreadyRegistered(_) => {
debug!("User {} already registered", mock_user.address);
}
_ => {
return Err(Box::new(e));
}
let user_db = user_db_service.get_user_db();
if let Err(e) = user_db.on_new_user(&mock_user.address) {
match e {
RegisterError::AlreadyRegistered(_) => {
debug!("User {} already registered", mock_user.address);
}
_ => {
return Err(Box::new(e));
}
}
user_db.on_new_tx(&mock_user.address, Some(mock_user.tx_count))?;
}
user_db.on_new_tx(&mock_user.address, Some(mock_user.tx_count))?;
}
}
@@ -202,228 +207,3 @@ pub async fn run_prover(
let _ = set.join_all().await;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
// std
use std::net::{IpAddr, Ipv4Addr};
use std::str::FromStr;
use std::sync::Arc;
// third-party
use alloy::primitives::{Address, U256};
use futures::FutureExt;
use parking_lot::RwLock;
use tokio::task;
use tonic::Response;
use tracing::info;
use tracing_test::traced_test;
// internal
use crate::grpc_service::prover_proto::rln_prover_client::RlnProverClient;
use crate::grpc_service::prover_proto::{
Address as GrpcAddress, GetUserTierInfoReply, GetUserTierInfoRequest, RegisterUserReply,
RegisterUserRequest, RegistrationStatus, RlnProofFilter, RlnProofReply,
SendTransactionReply, SendTransactionRequest, U256 as GrpcU256, Wei as GrpcWei,
};
async fn proof_sender(port: u16, addresses: Vec<Address>, proof_count: usize) {
let chain_id = GrpcU256 {
// FIXME: LE or BE?
value: U256::from(1).to_le_bytes::<32>().to_vec(),
};
let url = format!("http://127.0.0.1:{}", port);
let mut client = RlnProverClient::connect(url).await.unwrap();
let addr = GrpcAddress {
value: addresses[0].to_vec(),
};
let wei = GrpcWei {
// FIXME: LE or BE?
value: U256::from(1000).to_le_bytes::<32>().to_vec(),
};
let tx_hash = U256::from(42).to_le_bytes::<32>().to_vec();
let request_0 = SendTransactionRequest {
gas_price: Some(wei),
sender: Some(addr),
chain_id: Some(chain_id),
transaction_hash: tx_hash,
};
let request = tonic::Request::new(request_0);
let response: Response<SendTransactionReply> =
client.send_transaction(request).await.unwrap();
assert_eq!(response.into_inner().result, true);
}
async fn proof_collector(port: u16) -> Vec<RlnProofReply> {
let result = Arc::new(RwLock::new(vec![]));
let url = format!("http://127.0.0.1:{}", port);
let mut client = RlnProverClient::connect(url).await.unwrap();
let request_0 = RlnProofFilter { address: None };
let request = tonic::Request::new(request_0);
let stream_ = client.get_proofs(request).await.unwrap();
let mut stream = stream_.into_inner();
let result_2 = result.clone();
let receiver = async move {
while let Some(response) = stream.message().await.unwrap() {
result_2.write().push(response);
}
};
let _res = tokio::time::timeout(Duration::from_secs(10), receiver).await;
std::mem::take(&mut *result.write())
}
async fn register_users(port: u16, addresses: Vec<Address>) {
let url = format!("http://127.0.0.1:{}", port);
let mut client = RlnProverClient::connect(url).await.unwrap();
for address in addresses {
let addr = GrpcAddress {
value: address.to_vec(),
};
let request_0 = RegisterUserRequest { user: Some(addr) };
let request = tonic::Request::new(request_0);
let response: Response<RegisterUserReply> =
client.register_user(request).await.unwrap();
assert_eq!(
RegistrationStatus::try_from(response.into_inner().status).unwrap(),
RegistrationStatus::Success
);
}
}
async fn query_user_info(port: u16, addresses: Vec<Address>) -> Vec<GetUserTierInfoReply> {
let url = format!("http://127.0.0.1:{}", port);
let mut client = RlnProverClient::connect(url).await.unwrap();
let mut result = vec![];
for address in addresses {
let addr = GrpcAddress {
value: address.to_vec(),
};
let request_0 = GetUserTierInfoRequest { user: Some(addr) };
let request = tonic::Request::new(request_0);
let resp: Response<GetUserTierInfoReply> =
client.get_user_tier_info(request).await.unwrap();
result.push(resp.into_inner());
}
result
}
#[tokio::test]
#[traced_test]
async fn test_grpc_register_users() {
let addresses = vec![
Address::from_str("0xd8da6bf26964af9d7eed9e03e53415d37aa96045").unwrap(),
Address::from_str("0xb20a608c624Ca5003905aA834De7156C68b2E1d0").unwrap(),
];
let temp_folder = tempfile::tempdir().unwrap();
let temp_folder_tree = tempfile::tempdir().unwrap();
let port = 50051;
let app_args = AppArgs {
ip: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
port,
ws_rpc_url: None,
db_path: temp_folder.path().to_path_buf(),
merkle_tree_path: temp_folder_tree.path().to_path_buf(),
ksc_address: None,
rlnsc_address: None,
tsc_address: None,
mock_sc: Some(true),
mock_user: None,
config_path: Default::default(),
no_config: Some(true),
metrics_ip: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
metrics_port: 30031,
broadcast_channel_size: 100,
proof_service_count: 8,
transaction_channel_size: 100,
proof_sender_channel_size: 100,
};
info!("Starting prover...");
let prover_handle = task::spawn(run_prover(app_args));
// Wait for the prover to be ready
// Note: if unit test is failing - maybe add an optional notification when service is ready
tokio::time::sleep(Duration::from_secs(5)).await;
info!("Registering some users...");
register_users(port, addresses.clone()).await;
info!("Query info for these new users...");
let res = query_user_info(port, addresses.clone()).await;
assert_eq!(res.len(), addresses.len());
info!("Aborting prover...");
prover_handle.abort();
tokio::time::sleep(Duration::from_secs(1)).await;
}
#[tokio::test]
#[traced_test]
async fn test_grpc_gen_proof() {
let addresses = vec![
Address::from_str("0xd8da6bf26964af9d7eed9e03e53415d37aa96045").unwrap(),
Address::from_str("0xb20a608c624Ca5003905aA834De7156C68b2E1d0").unwrap(),
];
let temp_folder = tempfile::tempdir().unwrap();
let temp_folder_tree = tempfile::tempdir().unwrap();
let port = 50052;
let app_args = AppArgs {
ip: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
port,
ws_rpc_url: None,
db_path: temp_folder.path().to_path_buf(),
merkle_tree_path: temp_folder_tree.path().to_path_buf(),
ksc_address: None,
rlnsc_address: None,
tsc_address: None,
mock_sc: Some(true),
mock_user: None,
config_path: Default::default(),
no_config: Some(true),
metrics_ip: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
metrics_port: 30031,
broadcast_channel_size: 100,
proof_service_count: 8,
transaction_channel_size: 100,
proof_sender_channel_size: 100,
};
info!("Starting prover...");
let prover_handle = task::spawn(run_prover(app_args));
// Wait for the prover to be ready
// Note: if unit test is failing - maybe add an optional notification when service is ready
tokio::time::sleep(Duration::from_secs(5)).await;
info!("Registering some users...");
register_users(port, addresses.clone()).await;
info!("Sending tx and collecting proofs...");
let proof_count = 1;
let mut set = JoinSet::new();
set.spawn(
proof_sender(port, addresses.clone(), proof_count).map(|_| vec![]), // JoinSet require having the same return type
);
set.spawn(proof_collector(port));
let res = set.join_all().await;
assert_eq!(res[1].len(), proof_count);
info!("Aborting prover...");
prover_handle.abort();
tokio::time::sleep(Duration::from_secs(1)).await;
}
}

303
prover/src/proof_service.rs
View File

@@ -168,7 +168,7 @@ mod tests {
use std::path::PathBuf;
// third-party
use alloy::primitives::{Address, address};
use ark_groth16::{Proof as ArkProof, Proof, VerifyingKey};
use ark_groth16::{Proof as ArkProof, VerifyingKey};
use ark_serialize::CanonicalDeserialize;
use claims::assert_matches;
use futures::TryFutureExt;
@@ -177,8 +177,7 @@ mod tests {
// third-party: zerokit
use rln::{
circuit::{Curve, zkey_from_folder},
error::ComputeIdSecretError,
protocol::{compute_id_secret, deserialize_proof_values, verify_proof},
protocol::{deserialize_proof_values, verify_proof},
};
// internal
use crate::user_db_service::UserDbService;
@@ -188,7 +187,6 @@ mod tests {
const ADDR_2: Address = address!("0xb20a608c624Ca5003905aA834De7156C68b2E1d0");
const TX_HASH_1: [u8; 32] = [0x011; 32];
const TX_HASH_1_2: [u8; 32] = [0x12; 32];
#[derive(thiserror::Error, Debug)]
enum AppErrorExt {
@@ -202,10 +200,6 @@ mod tests {
ProofVerification,
#[error("Exiting...")]
Exit,
#[error(transparent)]
RecoverSecretFailed(ComputeIdSecretError),
#[error("Recovered secret")]
RecoveredSecret(Fr),
}
async fn proof_sender(
@@ -326,297 +320,4 @@ mod tests {
// Everything ok if proof_verifier return AppErrorExt::Exit else there is a real error
assert_matches!(res, Err(AppErrorExt::Exit));
}
#[tokio::test]
#[tracing_test::traced_test]
async fn test_user_not_registered() {
// Ask for a proof for an unregistered user
// Queues
let (broadcast_sender, _broadcast_receiver) = broadcast::channel(2);
let mut broadcast_receiver = broadcast_sender.subscribe();
let (mut proof_tx, proof_rx) = async_channel::unbounded();
// Epoch
let epoch = Epoch::from(11);
let epoch_slice = EpochSlice::from(42);
let epoch_store = Arc::new(RwLock::new((epoch, epoch_slice)));
// User db
let temp_folder = tempfile::tempdir().unwrap();
let temp_folder_tree = tempfile::tempdir().unwrap();
let user_db_service = UserDbService::new(
PathBuf::from(temp_folder.path()),
PathBuf::from(temp_folder_tree.path()),
Default::default(),
epoch_store.clone(),
10.into(),
Default::default(),
)
.unwrap();
let user_db = user_db_service.get_user_db();
user_db.on_new_user(&ADDR_1).unwrap();
// user_db.on_new_user(ADDR_2).unwrap();
let rln_identifier = Arc::new(RlnIdentifier::new(b"foo bar baz"));
// Proof service
let proof_service = ProofService::new(
proof_rx,
broadcast_sender,
epoch_store,
user_db.clone(),
RateLimit::from(10),
0,
);
// Verification
let proving_key = zkey_from_folder();
let verification_key = &proving_key.0.vk;
info!("Starting...");
let res = tokio::try_join!(
proof_service.serve().map_err(AppErrorExt::AppError),
proof_verifier(&mut broadcast_receiver, verification_key),
proof_sender(ADDR_2, &mut proof_tx, rln_identifier.clone(), &user_db),
);
// Expect this error (any other error is a real error)
assert_matches!(
res,
Err(AppErrorExt::ProofGeneration(
ProofGenerationStringError::MerkleProofError(_)
))
);
}
async fn proof_reveal_secret(
broadcast_receiver: &mut broadcast::Receiver<
Result<ProofSendingData, ProofGenerationStringError>,
>,
// verifying_key: &VerifyingKey<Curve>,
) -> Result<(), AppErrorExt> {
// used by test_user_spamming unit test
debug!("Starting broadcast receiver...");
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
let mut proof_values_store = vec![];
loop {
let res =
tokio::time::timeout(std::time::Duration::from_secs(5), broadcast_receiver.recv())
.await
.map_err(|_e| AppErrorExt::Elapsed)?;
let res = res.unwrap();
let res = res?;
let mut proof_cursor = Cursor::new(&res.proof);
let _proof: Proof<Curve> = ArkProof::deserialize_compressed(&mut proof_cursor).unwrap();
let position = proof_cursor.position() as usize;
let proof_cursor_2 = &proof_cursor.get_ref().as_slice()[position..];
let (proof_values, _) = deserialize_proof_values(proof_cursor_2);
proof_values_store.push(proof_values);
if proof_values_store.len() >= 2 {
break;
}
}
debug!("Now recovering secret hash...");
let proof_values_0 = proof_values_store.get(0).unwrap();
let proof_values_1 = proof_values_store.get(1).unwrap();
println!("proof_values_0: {:?}", proof_values_0);
println!("proof_values_1: {:?}", proof_values_1);
let share1 = (proof_values_0.x, proof_values_0.y);
let share2 = (proof_values_1.x, proof_values_1.y);
// Note: if not in test, should check for external nullifier
let recovered_identity_secret_hash =
compute_id_secret(share1, share2).map_err(|e| AppErrorExt::RecoverSecretFailed(e))?;
debug!(
"recovered_identity_secret_hash: {:?}",
recovered_identity_secret_hash
);
// Exit after receiving one proof
Err::<(), AppErrorExt>(AppErrorExt::RecoveredSecret(recovered_identity_secret_hash))
}
async fn proof_sender_2(
proof_tx: &mut async_channel::Sender<ProofGenerationData>,
rln_identifier: Arc<RlnIdentifier>,
user_db: &UserDb,
sender: Address,
tx_hashes: ([u8; 32], [u8; 32]),
) -> Result<(), AppErrorExt> {
// used by test_proof_generation unit test
debug!("Starting proof sender 2...");
debug!("Waiting a bit before sending proof...");
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
debug!("Sending proof...");
proof_tx
.send(ProofGenerationData {
user_identity: user_db.get_user(&sender).unwrap(),
rln_identifier: rln_identifier.clone(),
tx_counter: 0,
tx_sender: sender.clone(),
tx_hash: tx_hashes.0.to_vec(),
})
.await
.unwrap();
debug!("Sending proof done");
debug!("Waiting a bit before sending 2nd proof...");
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
debug!("Sending 2nd proof...");
proof_tx
.send(ProofGenerationData {
user_identity: user_db.get_user(&sender).unwrap(),
rln_identifier,
tx_counter: 1,
tx_sender: sender,
tx_hash: tx_hashes.1.to_vec(),
})
.await
.unwrap();
debug!("Sending 2nd proof done");
Ok::<(), AppErrorExt>(())
}
#[tokio::test]
#[tracing_test::traced_test]
async fn test_user_spamming() {
// Recover secret from a user spamming the system
// Queues
let (broadcast_sender, _broadcast_receiver) = broadcast::channel(2);
let mut broadcast_receiver = broadcast_sender.subscribe();
let (mut proof_tx, proof_rx) = async_channel::unbounded();
// Epoch
let epoch = Epoch::from(11);
let epoch_slice = EpochSlice::from(42);
let epoch_store = Arc::new(RwLock::new((epoch, epoch_slice)));
// Limits
let rate_limit = RateLimit::from(1);
// User db
let temp_folder = tempfile::tempdir().unwrap();
let temp_folder_tree = tempfile::tempdir().unwrap();
let user_db_service = UserDbService::new(
PathBuf::from(temp_folder.path()),
PathBuf::from(temp_folder_tree.path()),
Default::default(),
epoch_store.clone(),
rate_limit,
Default::default(),
)
.unwrap();
let user_db = user_db_service.get_user_db();
user_db.on_new_user(&ADDR_1).unwrap();
let user_addr_1 = user_db.get_user(&ADDR_1).unwrap();
user_db.on_new_user(&ADDR_2).unwrap();
let rln_identifier = Arc::new(RlnIdentifier::new(b"foo bar baz"));
// Proof service
let proof_service = ProofService::new(
proof_rx,
broadcast_sender,
epoch_store,
user_db.clone(),
rate_limit,
0,
);
info!("Starting...");
let res = tokio::try_join!(
proof_service.serve().map_err(AppErrorExt::AppError),
proof_reveal_secret(&mut broadcast_receiver),
proof_sender_2(
&mut proof_tx,
rln_identifier.clone(),
&user_db,
ADDR_1,
(TX_HASH_1, TX_HASH_1_2)
),
);
match res {
Err(AppErrorExt::RecoveredSecret(secret_hash)) => {
assert_eq!(secret_hash, user_addr_1.secret_hash);
}
_ => {
panic!("Unexpected result");
}
}
}
#[tokio::test]
#[tracing_test::traced_test]
async fn test_user_spamming_same_signal() {
// Recover secret from a user spamming the system
// Queues
let (broadcast_sender, _broadcast_receiver) = broadcast::channel(2);
let mut broadcast_receiver = broadcast_sender.subscribe();
let (mut proof_tx, proof_rx) = async_channel::unbounded();
// Epoch
let epoch = Epoch::from(11);
let epoch_slice = EpochSlice::from(42);
let epoch_store = Arc::new(RwLock::new((epoch, epoch_slice)));
// Limits
let rate_limit = RateLimit::from(1);
// User db - limit is 1 message per epoch
let temp_folder = tempfile::tempdir().unwrap();
let temp_folder_tree = tempfile::tempdir().unwrap();
let user_db_service = UserDbService::new(
PathBuf::from(temp_folder.path()),
PathBuf::from(temp_folder_tree.path()),
Default::default(),
epoch_store.clone(),
rate_limit,
Default::default(),
)
.unwrap();
let user_db = user_db_service.get_user_db();
user_db.on_new_user(&ADDR_1).unwrap();
let user_addr_1 = user_db.get_user(&ADDR_1).unwrap();
debug!("user_addr_1: {:?}", user_addr_1);
user_db.on_new_user(&ADDR_2).unwrap();
let rln_identifier = Arc::new(RlnIdentifier::new(b"foo bar baz"));
// Proof service
let proof_service = ProofService::new(
proof_rx,
broadcast_sender,
epoch_store,
user_db.clone(),
rate_limit,
0,
);
info!("Starting...");
let res = tokio::try_join!(
proof_service.serve().map_err(AppErrorExt::AppError),
proof_reveal_secret(&mut broadcast_receiver),
proof_sender_2(
&mut proof_tx,
rln_identifier.clone(),
&user_db,
ADDR_1,
(TX_HASH_1, TX_HASH_1)
),
);
assert_matches!(res, Err(AppErrorExt::RecoverSecretFailed(_)));
}
}

408
prover/src/proof_service_tests.rs Normal file
View File

@@ -0,0 +1,408 @@
#[cfg(test)]
mod proof_service_tests {
use std::io::Cursor;
use std::path::PathBuf;
use std::sync::Arc;
// third-party
use alloy::primitives::{Address, address};
use ark_bn254::Fr;
use ark_groth16::{Proof as ArkProof, Proof, VerifyingKey};
use ark_serialize::CanonicalDeserialize;
use claims::assert_matches;
use futures::TryFutureExt;
use parking_lot::RwLock;
use rln::circuit::{Curve, zkey_from_folder};
use rln::error::ComputeIdSecretError;
use rln::protocol::{compute_id_secret, deserialize_proof_values, verify_proof};
use tokio::sync::broadcast;
use tracing::{debug, info};
// internal
use crate::epoch_service::{Epoch, EpochSlice};
use crate::error::{AppError, ProofGenerationStringError};
use crate::proof_generation::{ProofGenerationData, ProofSendingData};
use crate::proof_service::ProofService;
use crate::user_db::UserDb;
use crate::user_db_service::UserDbService;
use crate::user_db_types::RateLimit;
use rln_proof::RlnIdentifier;
const ADDR_1: Address = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045");
const ADDR_2: Address = address!("0xb20a608c624Ca5003905aA834De7156C68b2E1d0");
const TX_HASH_1: [u8; 32] = [0x011; 32];
const TX_HASH_1_2: [u8; 32] = [0x12; 32];
#[derive(thiserror::Error, Debug)]
enum AppErrorExt {
#[error("AppError: {0}")]
AppError(#[from] AppError),
#[error("Future timeout")]
Elapsed,
#[error("Proof generation failed: {0}")]
ProofGeneration(#[from] ProofGenerationStringError),
#[error("Proof verification failed")]
ProofVerification,
#[error("Exiting...")]
Exit,
#[error(transparent)]
RecoverSecretFailed(ComputeIdSecretError),
#[error("Recovered secret")]
RecoveredSecret(Fr),
}
async fn proof_sender(
sender: Address,
proof_tx: &mut async_channel::Sender<ProofGenerationData>,
rln_identifier: Arc<RlnIdentifier>,
user_db: &UserDb,
) -> Result<(), AppErrorExt> {
// used by test_proof_generation unit test
debug!("Starting proof sender...");
debug!("Waiting a bit before sending proof...");
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
debug!("Sending proof...");
proof_tx
.send(ProofGenerationData {
user_identity: user_db.get_user(&ADDR_1).unwrap(),
rln_identifier,
tx_counter: 0,
tx_sender: sender,
tx_hash: TX_HASH_1.to_vec(),
})
.await
.unwrap();
debug!("Sending proof done");
// tokio::time::sleep(std::time::Duration::from_secs(10)).await;
Ok::<(), AppErrorExt>(())
}
async fn proof_verifier(
broadcast_receiver: &mut broadcast::Receiver<
Result<ProofSendingData, ProofGenerationStringError>,
>,
verifying_key: &VerifyingKey<Curve>,
) -> Result<(), AppErrorExt> {
// used by test_proof_generation unit test
debug!("Starting broadcast receiver...");
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
let res =
tokio::time::timeout(std::time::Duration::from_secs(5), broadcast_receiver.recv())
.await
.map_err(|_e| AppErrorExt::Elapsed)?;
debug!("res: {:?}", res);
let res = res.unwrap();
let res = res?;
let mut proof_cursor = Cursor::new(&res.proof);
debug!("proof cursor: {:?}", proof_cursor);
let proof = ArkProof::deserialize_compressed(&mut proof_cursor).unwrap();
let position = proof_cursor.position() as usize;
let proof_cursor_2 = &proof_cursor.get_ref().as_slice()[position..];
let (proof_values, _) = deserialize_proof_values(proof_cursor_2);
debug!("[proof verifier] proof: {:?}", proof);
debug!("[proof verifier] proof_values: {:?}", proof_values);
let verified = verify_proof(verifying_key, &proof, &proof_values)
.map_err(|_e| AppErrorExt::ProofVerification)?;
debug!("verified: {:?}", verified);
// Exit after receiving one proof
Err::<(), AppErrorExt>(AppErrorExt::Exit)
}
#[tokio::test]
#[tracing_test::traced_test]
async fn test_user_not_registered() {
// Ask for a proof for an unregistered user
// Queues
let (broadcast_sender, _broadcast_receiver) = broadcast::channel(2);
let mut broadcast_receiver = broadcast_sender.subscribe();
let (mut proof_tx, proof_rx) = async_channel::unbounded();
// Epoch
let epoch = Epoch::from(11);
let epoch_slice = EpochSlice::from(42);
let epoch_store = Arc::new(RwLock::new((epoch, epoch_slice)));
// User db
let temp_folder = tempfile::tempdir().unwrap();
let temp_folder_tree = tempfile::tempdir().unwrap();
let user_db_service = UserDbService::new(
PathBuf::from(temp_folder.path()),
PathBuf::from(temp_folder_tree.path()),
Default::default(),
epoch_store.clone(),
10.into(),
Default::default(),
)
.unwrap();
let user_db = user_db_service.get_user_db();
user_db.on_new_user(&ADDR_1).unwrap();
// user_db.on_new_user(ADDR_2).unwrap();
let rln_identifier = Arc::new(RlnIdentifier::new(b"foo bar baz"));
// Proof service
let proof_service = ProofService::new(
proof_rx,
broadcast_sender,
epoch_store,
user_db.clone(),
RateLimit::from(10),
0,
);
// Verification
let proving_key = zkey_from_folder();
let verification_key = &proving_key.0.vk;
info!("Starting...");
let res = tokio::try_join!(
proof_service.serve().map_err(AppErrorExt::AppError),
proof_verifier(&mut broadcast_receiver, verification_key),
proof_sender(ADDR_2, &mut proof_tx, rln_identifier.clone(), &user_db),
);
// Expect this error (any other error is a real error)
assert_matches!(
res,
Err(AppErrorExt::ProofGeneration(
ProofGenerationStringError::MerkleProofError(_)
))
);
}
async fn proof_reveal_secret(
broadcast_receiver: &mut broadcast::Receiver<
Result<ProofSendingData, ProofGenerationStringError>,
>,
// verifying_key: &VerifyingKey<Curve>,
) -> Result<(), AppErrorExt> {
// used by test_user_spamming unit test
debug!("Starting broadcast receiver...");
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
let mut proof_values_store = vec![];
loop {
let res =
tokio::time::timeout(std::time::Duration::from_secs(5), broadcast_receiver.recv())
.await
.map_err(|_e| AppErrorExt::Elapsed)?;
let res = res.unwrap();
let res = res?;
let mut proof_cursor = Cursor::new(&res.proof);
let _proof: Proof<Curve> = ArkProof::deserialize_compressed(&mut proof_cursor).unwrap();
let position = proof_cursor.position() as usize;
let proof_cursor_2 = &proof_cursor.get_ref().as_slice()[position..];
let (proof_values, _) = deserialize_proof_values(proof_cursor_2);
proof_values_store.push(proof_values);
if proof_values_store.len() >= 2 {
break;
}
}
debug!("Now recovering secret hash...");
let proof_values_0 = proof_values_store.get(0).unwrap();
let proof_values_1 = proof_values_store.get(1).unwrap();
println!("proof_values_0: {:?}", proof_values_0);
println!("proof_values_1: {:?}", proof_values_1);
let share1 = (proof_values_0.x, proof_values_0.y);
let share2 = (proof_values_1.x, proof_values_1.y);
// Note: if not in test, should check for external nullifier
let recovered_identity_secret_hash =
compute_id_secret(share1, share2).map_err(|e| AppErrorExt::RecoverSecretFailed(e))?;
debug!(
"recovered_identity_secret_hash: {:?}",
recovered_identity_secret_hash
);
// Exit after receiving one proof
Err::<(), AppErrorExt>(AppErrorExt::RecoveredSecret(recovered_identity_secret_hash))
}
async fn proof_sender_2(
proof_tx: &mut async_channel::Sender<ProofGenerationData>,
rln_identifier: Arc<RlnIdentifier>,
user_db: &UserDb,
sender: Address,
tx_hashes: ([u8; 32], [u8; 32]),
) -> Result<(), AppErrorExt> {
// used by test_proof_generation unit test
debug!("Starting proof sender 2...");
debug!("Waiting a bit before sending proof...");
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
debug!("Sending proof...");
proof_tx
.send(ProofGenerationData {
user_identity: user_db.get_user(&sender).unwrap(),
rln_identifier: rln_identifier.clone(),
tx_counter: 0,
tx_sender: sender.clone(),
tx_hash: tx_hashes.0.to_vec(),
})
.await
.unwrap();
debug!("Sending proof done");
debug!("Waiting a bit before sending 2nd proof...");
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
debug!("Sending 2nd proof...");
proof_tx
.send(ProofGenerationData {
user_identity: user_db.get_user(&sender).unwrap(),
rln_identifier,
tx_counter: 1,
tx_sender: sender,
tx_hash: tx_hashes.1.to_vec(),
})
.await
.unwrap();
debug!("Sending 2nd proof done");
Ok::<(), AppErrorExt>(())
}
#[tokio::test]
#[tracing_test::traced_test]
async fn test_user_spamming() {
// Recover secret from a user spamming the system
// Queues
let (broadcast_sender, _broadcast_receiver) = broadcast::channel(2);
let mut broadcast_receiver = broadcast_sender.subscribe();
let (mut proof_tx, proof_rx) = async_channel::unbounded();
// Epoch
let epoch = Epoch::from(11);
let epoch_slice = EpochSlice::from(42);
let epoch_store = Arc::new(RwLock::new((epoch, epoch_slice)));
// Limits
let rate_limit = RateLimit::from(1);
// User db
let temp_folder = tempfile::tempdir().unwrap();
let temp_folder_tree = tempfile::tempdir().unwrap();
let user_db_service = UserDbService::new(
PathBuf::from(temp_folder.path()),
PathBuf::from(temp_folder_tree.path()),
Default::default(),
epoch_store.clone(),
rate_limit,
Default::default(),
)
.unwrap();
let user_db = user_db_service.get_user_db();
user_db.on_new_user(&ADDR_1).unwrap();
let user_addr_1 = user_db.get_user(&ADDR_1).unwrap();
user_db.on_new_user(&ADDR_2).unwrap();
let rln_identifier = Arc::new(RlnIdentifier::new(b"foo bar baz"));
// Proof service
let proof_service = ProofService::new(
proof_rx,
broadcast_sender,
epoch_store,
user_db.clone(),
rate_limit,
0,
);
info!("Starting...");
let res = tokio::try_join!(
proof_service.serve().map_err(AppErrorExt::AppError),
proof_reveal_secret(&mut broadcast_receiver),
proof_sender_2(
&mut proof_tx,
rln_identifier.clone(),
&user_db,
ADDR_1,
(TX_HASH_1, TX_HASH_1_2)
),
);
match res {
Err(AppErrorExt::RecoveredSecret(secret_hash)) => {
assert_eq!(secret_hash, user_addr_1.secret_hash);
}
_ => {
panic!("Unexpected result");
}
}
}
#[tokio::test]
#[tracing_test::traced_test]
async fn test_user_spamming_same_signal() {
// Recover secret from a user spamming the system
// Queues
let (broadcast_sender, _broadcast_receiver) = broadcast::channel(2);
let mut broadcast_receiver = broadcast_sender.subscribe();
let (mut proof_tx, proof_rx) = async_channel::unbounded();
// Epoch
let epoch = Epoch::from(11);
let epoch_slice = EpochSlice::from(42);
let epoch_store = Arc::new(RwLock::new((epoch, epoch_slice)));
// Limits
let rate_limit = RateLimit::from(1);
// User db - limit is 1 message per epoch
let temp_folder = tempfile::tempdir().unwrap();
let temp_folder_tree = tempfile::tempdir().unwrap();
let user_db_service = UserDbService::new(
PathBuf::from(temp_folder.path()),
PathBuf::from(temp_folder_tree.path()),
Default::default(),
epoch_store.clone(),
rate_limit,
Default::default(),
)
.unwrap();
let user_db = user_db_service.get_user_db();
user_db.on_new_user(&ADDR_1).unwrap();
let user_addr_1 = user_db.get_user(&ADDR_1).unwrap();
debug!("user_addr_1: {:?}", user_addr_1);
user_db.on_new_user(&ADDR_2).unwrap();
let rln_identifier = Arc::new(RlnIdentifier::new(b"foo bar baz"));
// Proof service
let proof_service = ProofService::new(
proof_rx,
broadcast_sender,
epoch_store,
user_db.clone(),
rate_limit,
0,
);
info!("Starting...");
let res = tokio::try_join!(
proof_service.serve().map_err(AppErrorExt::AppError),
proof_reveal_secret(&mut broadcast_receiver),
proof_sender_2(
&mut proof_tx,
rln_identifier.clone(),
&user_db,
ADDR_1,
(TX_HASH_1, TX_HASH_1)
),
);
assert_matches!(res, Err(AppErrorExt::RecoverSecretFailed(_)));
}
}

1
prover/src/rocksdb_operands.rs
View File

@@ -46,6 +46,7 @@ impl EpochCounterSerializer {
size_of::<EpochCounters>()
}
#[allow(dead_code)]
pub(crate) fn size_hint(&self) -> usize {
Self::size_hint_()
}

145
prover/src/user_db.rs
View File

@@ -194,7 +194,7 @@ impl UserDb {
self.db.cf_handle(TIER_LIMITS_CF).unwrap()
}
fn register(&self, address: Address) -> Result<Fr, RegisterError> {
pub(crate) fn register(&self, address: Address) -> Result<Fr, RegisterError> {
let rln_identity_serializer = RlnUserIdentitySerializer {};
let merkle_index_serializer = MerkleTreeIndexSerializer {};
let merkle_index_deserializer = MerkleTreeIndexDeserializer {};
@@ -286,7 +286,7 @@ impl UserDb {
Ok(id_commitment)
}
fn has_user(&self, address: &Address) -> Result<bool, rocksdb::Error> {
pub(crate) fn has_user(&self, address: &Address) -> Result<bool, rocksdb::Error> {
let cf_user = self.get_user_cf();
self.db
.get_pinned_cf(cf_user, address.as_slice())
@@ -403,7 +403,7 @@ impl UserDb {
Ok(epoch_slice_counter)
}
fn get_tx_counter(
pub(crate) fn get_tx_counter(
&self,
address: &Address,
) -> Result<(EpochCounter, EpochSliceCounter), TxCounterError> {
@@ -472,7 +472,7 @@ impl UserDb {
}
#[cfg(test)]
fn get_merkle_tree_index(&self) -> Result<MerkleTreeIndex, MerkleTreeIndexError> {
pub(crate) fn get_merkle_tree_index(&self) -> Result<MerkleTreeIndex, MerkleTreeIndexError> {
let cf_mtree = self.get_mtree_cf();
Self::get_merkle_tree_index_(self.db.clone(), cf_mtree)
}
@@ -784,97 +784,33 @@ mod tests {
assert_eq!(tier_info.epoch_slice_tx_count, 1);
}
#[tokio::test]
async fn test_persistent_storage() {
#[test]
fn test_user_remove() {
let temp_folder = tempfile::tempdir().unwrap();
let temp_folder_tree = tempfile::tempdir().unwrap();
let epoch_store = Arc::new(RwLock::new(Default::default()));
let addr = Address::new([0; 20]);
{
let user_db = UserDb::new(
PathBuf::from(temp_folder.path()),
PathBuf::from(temp_folder_tree.path()),
epoch_store.clone(),
Default::default(),
Default::default(),
)
.unwrap();
let user_db = UserDb::new(
PathBuf::from(temp_folder.path()),
PathBuf::from(temp_folder_tree.path()),
epoch_store.clone(),
Default::default(),
Default::default(),
)
.unwrap();
assert_eq!(
user_db.get_merkle_tree_index().unwrap(),
MerkleTreeIndex::from(0)
);
// Register user
user_db.register(ADDR_1).unwrap();
assert_eq!(
user_db.get_merkle_tree_index().unwrap(),
MerkleTreeIndex::from(1)
);
// + 1 user
user_db.register(ADDR_2).unwrap();
assert_eq!(
user_db.get_merkle_tree_index().unwrap(),
MerkleTreeIndex::from(2)
);
assert_eq!(
user_db.get_user_merkle_tree_index(&ADDR_1).unwrap(),
MerkleTreeIndex::from(0)
);
assert_eq!(
user_db.get_user_merkle_tree_index(&ADDR_2).unwrap(),
MerkleTreeIndex::from(1)
);
assert_eq!(
user_db.on_new_tx(&ADDR_1, Some(2)),
Ok(EpochSliceCounter::from(2))
);
assert_eq!(
user_db.on_new_tx(&ADDR_2, Some(1000)),
Ok(EpochSliceCounter::from(1000))
);
// Should be dropped but let's make it explicit
drop(user_db);
}
{
// Reopen Db and check that is inside
let user_db = UserDb::new(
PathBuf::from(temp_folder.path()),
PathBuf::from(temp_folder_tree.path()),
epoch_store,
Default::default(),
Default::default(),
)
.unwrap();
assert_eq!(user_db.has_user(&addr).unwrap(), false);
assert_eq!(user_db.has_user(&ADDR_1).unwrap(), true);
assert_eq!(user_db.has_user(&ADDR_2).unwrap(), true);
assert_eq!(
user_db.get_tx_counter(&ADDR_1).unwrap(),
(2.into(), 2.into())
);
assert_eq!(
user_db.get_tx_counter(&ADDR_2).unwrap(),
(1000.into(), 1000.into())
);
assert_eq!(
user_db.get_merkle_tree_index().unwrap(),
MerkleTreeIndex::from(2)
);
assert_eq!(
user_db.get_user_merkle_tree_index(&ADDR_1).unwrap(),
MerkleTreeIndex::from(0)
);
assert_eq!(
user_db.get_user_merkle_tree_index(&ADDR_2).unwrap(),
MerkleTreeIndex::from(1)
);
}
user_db.register(ADDR_1).unwrap();
let mtree_index_add_addr_1 = user_db.merkle_tree.read().leaves_set();
user_db.register(ADDR_2).unwrap();
let mtree_index_add_addr_2 = user_db.merkle_tree.read().leaves_set();
assert_ne!(mtree_index_add_addr_1, mtree_index_add_addr_2);
user_db.remove_user(&ADDR_2, false);
let mtree_index_after_rm_addr_2 = user_db.merkle_tree.read().leaves_set();
assert_eq!(user_db.has_user(&ADDR_1), Ok(true));
assert_eq!(user_db.has_user(&ADDR_2), Ok(false));
// No reuse of index in PmTree (as this is a generic impl and could lead to security issue:
// like replay attack...)
assert_eq!(mtree_index_after_rm_addr_2, mtree_index_add_addr_2);
}
#[test]
@@ -924,33 +860,4 @@ mod tests {
assert_eq!(user_db.has_user(&ADDR_2), Ok(false));
assert_eq!(tree.read().leaves_set(), 2);
}
#[test]
fn test_user_remove() {
let temp_folder = tempfile::tempdir().unwrap();
let temp_folder_tree = tempfile::tempdir().unwrap();
let epoch_store = Arc::new(RwLock::new(Default::default()));
let user_db = UserDb::new(
PathBuf::from(temp_folder.path()),
PathBuf::from(temp_folder_tree.path()),
epoch_store.clone(),
Default::default(),
Default::default(),
)
.unwrap();
user_db.register(ADDR_1).unwrap();
let mtree_index_add_addr_1 = user_db.merkle_tree.read().leaves_set();
user_db.register(ADDR_2).unwrap();
let mtree_index_add_addr_2 = user_db.merkle_tree.read().leaves_set();
assert_ne!(mtree_index_add_addr_1, mtree_index_add_addr_2);
user_db.remove_user(&ADDR_2, false);
let mtree_index_after_rm_addr_2 = user_db.merkle_tree.read().leaves_set();
assert_eq!(user_db.has_user(&ADDR_1), Ok(true));
assert_eq!(user_db.has_user(&ADDR_2), Ok(false));
// No reuse of index in PmTree (as this is a generic impl and could lead to security issue:
// like replay attack...)
assert_eq!(mtree_index_after_rm_addr_2, mtree_index_add_addr_2);
}
}

108
prover/src/user_db_tests.rs Normal file
View File

@@ -0,0 +1,108 @@
#[cfg(test)]
mod user_db_tests {
// std
use std::path::PathBuf;
use std::sync::Arc;
// third-party
use alloy::primitives::{Address, address};
use parking_lot::RwLock;
// internal
use crate::user_db::UserDb;
use crate::user_db_types::{EpochSliceCounter, MerkleTreeIndex};
const ADDR_1: Address = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045");
const ADDR_2: Address = address!("0xb20a608c624Ca5003905aA834De7156C68b2E1d0");
#[tokio::test]
async fn test_persistent_storage() {
let temp_folder = tempfile::tempdir().unwrap();
let temp_folder_tree = tempfile::tempdir().unwrap();
let epoch_store = Arc::new(RwLock::new(Default::default()));
let addr = Address::new([0; 20]);
{
let user_db = UserDb::new(
PathBuf::from(temp_folder.path()),
PathBuf::from(temp_folder_tree.path()),
epoch_store.clone(),
Default::default(),
Default::default(),
)
.unwrap();
assert_eq!(
user_db.get_merkle_tree_index().unwrap(),
MerkleTreeIndex::from(0)
);
// Register user
user_db.register(ADDR_1).unwrap();
assert_eq!(
user_db.get_merkle_tree_index().unwrap(),
MerkleTreeIndex::from(1)
);
// + 1 user
user_db.register(ADDR_2).unwrap();
assert_eq!(
user_db.get_merkle_tree_index().unwrap(),
MerkleTreeIndex::from(2)
);
assert_eq!(
user_db.get_user_merkle_tree_index(&ADDR_1).unwrap(),
MerkleTreeIndex::from(0)
);
assert_eq!(
user_db.get_user_merkle_tree_index(&ADDR_2).unwrap(),
MerkleTreeIndex::from(1)
);
assert_eq!(
user_db.on_new_tx(&ADDR_1, Some(2)),
Ok(EpochSliceCounter::from(2))
);
assert_eq!(
user_db.on_new_tx(&ADDR_2, Some(1000)),
Ok(EpochSliceCounter::from(1000))
);
// Should be dropped but let's make it explicit
drop(user_db);
}
{
// Reopen Db and check that is inside
let user_db = UserDb::new(
PathBuf::from(temp_folder.path()),
PathBuf::from(temp_folder_tree.path()),
epoch_store,
Default::default(),
Default::default(),
)
.unwrap();
assert_eq!(user_db.has_user(&addr).unwrap(), false);
assert_eq!(user_db.has_user(&ADDR_1).unwrap(), true);
assert_eq!(user_db.has_user(&ADDR_2).unwrap(), true);
assert_eq!(
user_db.get_tx_counter(&ADDR_1).unwrap(),
(2.into(), 2.into())
);
assert_eq!(
user_db.get_tx_counter(&ADDR_2).unwrap(),
(1000.into(), 1000.into())
);
assert_eq!(
user_db.get_merkle_tree_index().unwrap(),
MerkleTreeIndex::from(2)
);
assert_eq!(
user_db.get_user_merkle_tree_index(&ADDR_1).unwrap(),
MerkleTreeIndex::from(0)
);
assert_eq!(
user_db.get_user_merkle_tree_index(&ADDR_2).unwrap(),
MerkleTreeIndex::from(1)
);
}
}
}

222
prover/tests/grpc_e2e.rs Normal file
View File

@@ -0,0 +1,222 @@
use alloy::primitives::{Address, U256};
use futures::FutureExt;
use parking_lot::RwLock;
use prover::{AppArgs, run_prover};
use std::net::{IpAddr, Ipv4Addr};
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;
use tokio::task;
use tokio::task::JoinSet;
use tonic::Response;
use tracing::info;
use tracing_test::traced_test;
pub mod prover_proto {
// Include generated code (see build.rs)
tonic::include_proto!("prover");
}
use crate::prover_proto::{
Address as GrpcAddress, GetUserTierInfoReply, GetUserTierInfoRequest, RegisterUserReply,
RegisterUserRequest, RegistrationStatus, RlnProofFilter, RlnProofReply, SendTransactionReply,
SendTransactionRequest, U256 as GrpcU256, Wei as GrpcWei, rln_prover_client::RlnProverClient,
};
async fn register_users(port: u16, addresses: Vec<Address>) {
let url = format!("http://127.0.0.1:{}", port);
let mut client = RlnProverClient::connect(url).await.unwrap();
for address in addresses {
let addr = GrpcAddress {
value: address.to_vec(),
};
let request_0 = RegisterUserRequest { user: Some(addr) };
let request = tonic::Request::new(request_0);
let response: Response<RegisterUserReply> = client.register_user(request).await.unwrap();
assert_eq!(
RegistrationStatus::try_from(response.into_inner().status).unwrap(),
RegistrationStatus::Success
);
}
}
async fn query_user_info(port: u16, addresses: Vec<Address>) -> Vec<GetUserTierInfoReply> {
let url = format!("http://127.0.0.1:{}", port);
let mut client = RlnProverClient::connect(url).await.unwrap();
let mut result = vec![];
for address in addresses {
let addr = GrpcAddress {
value: address.to_vec(),
};
let request_0 = GetUserTierInfoRequest { user: Some(addr) };
let request = tonic::Request::new(request_0);
let resp: Response<GetUserTierInfoReply> =
client.get_user_tier_info(request).await.unwrap();
result.push(resp.into_inner());
}
result
}
#[tokio::test]
#[traced_test]
async fn test_grpc_register_users() {
let addresses = vec![
Address::from_str("0xd8da6bf26964af9d7eed9e03e53415d37aa96045").unwrap(),
Address::from_str("0xb20a608c624Ca5003905aA834De7156C68b2E1d0").unwrap(),
];
let temp_folder = tempfile::tempdir().unwrap();
let temp_folder_tree = tempfile::tempdir().unwrap();
let port = 50051;
let app_args = AppArgs {
ip: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
port,
ws_rpc_url: None,
db_path: temp_folder.path().to_path_buf(),
merkle_tree_path: temp_folder_tree.path().to_path_buf(),
ksc_address: None,
rlnsc_address: None,
tsc_address: None,
mock_sc: Some(true),
mock_user: None,
config_path: Default::default(),
no_config: Some(true),
metrics_ip: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
metrics_port: 30031,
broadcast_channel_size: 100,
proof_service_count: 8,
transaction_channel_size: 100,
proof_sender_channel_size: 100,
};
info!("Starting prover...");
let prover_handle = task::spawn(run_prover(app_args));
// Wait for the prover to be ready
// Note: if unit test is failing - maybe add an optional notification when service is ready
tokio::time::sleep(Duration::from_secs(5)).await;
info!("Registering some users...");
register_users(port, addresses.clone()).await;
info!("Query info for these new users...");
let res = query_user_info(port, addresses.clone()).await;
assert_eq!(res.len(), addresses.len());
info!("Aborting prover...");
prover_handle.abort();
tokio::time::sleep(Duration::from_secs(1)).await;
}
async fn proof_sender(port: u16, addresses: Vec<Address>, _proof_count: usize) {
let chain_id = GrpcU256 {
// FIXME: LE or BE?
value: U256::from(1).to_le_bytes::<32>().to_vec(),
};
let url = format!("http://127.0.0.1:{}", port);
let mut client = RlnProverClient::connect(url).await.unwrap();
let addr = GrpcAddress {
value: addresses[0].to_vec(),
};
let wei = GrpcWei {
// FIXME: LE or BE?
value: U256::from(1000).to_le_bytes::<32>().to_vec(),
};
let tx_hash = U256::from(42).to_le_bytes::<32>().to_vec();
let request_0 = SendTransactionRequest {
gas_price: Some(wei),
sender: Some(addr),
chain_id: Some(chain_id),
transaction_hash: tx_hash,
};
let request = tonic::Request::new(request_0);
let response: Response<SendTransactionReply> = client.send_transaction(request).await.unwrap();
assert_eq!(response.into_inner().result, true);
}
async fn proof_collector(port: u16) -> Vec<RlnProofReply> {
let result = Arc::new(RwLock::new(vec![]));
let url = format!("http://127.0.0.1:{}", port);
let mut client = RlnProverClient::connect(url).await.unwrap();
let request_0 = RlnProofFilter { address: None };
let request = tonic::Request::new(request_0);
let stream_ = client.get_proofs(request).await.unwrap();
let mut stream = stream_.into_inner();
let result_2 = result.clone();
let receiver = async move {
while let Some(response) = stream.message().await.unwrap() {
result_2.write().push(response);
}
};
let _res = tokio::time::timeout(Duration::from_secs(10), receiver).await;
std::mem::take(&mut *result.write())
}
#[tokio::test]
#[traced_test]
async fn test_grpc_gen_proof() {
let addresses = vec![
Address::from_str("0xd8da6bf26964af9d7eed9e03e53415d37aa96045").unwrap(),
Address::from_str("0xb20a608c624Ca5003905aA834De7156C68b2E1d0").unwrap(),
];
let temp_folder = tempfile::tempdir().unwrap();
let temp_folder_tree = tempfile::tempdir().unwrap();
let port = 50052;
let app_args = AppArgs {
ip: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
port,
ws_rpc_url: None,
db_path: temp_folder.path().to_path_buf(),
merkle_tree_path: temp_folder_tree.path().to_path_buf(),
ksc_address: None,
rlnsc_address: None,
tsc_address: None,
mock_sc: Some(true),
mock_user: None,
config_path: Default::default(),
no_config: Some(true),
metrics_ip: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
metrics_port: 30031,
broadcast_channel_size: 100,
proof_service_count: 8,
transaction_channel_size: 100,
proof_sender_channel_size: 100,
};
info!("Starting prover...");
let prover_handle = task::spawn(run_prover(app_args));
// Wait for the prover to be ready
// Note: if unit test is failing - maybe add an optional notification when service is ready
tokio::time::sleep(Duration::from_secs(5)).await;
info!("Registering some users...");
register_users(port, addresses.clone()).await;
info!("Sending tx and collecting proofs...");
let proof_count = 1;
let mut set = JoinSet::new();
set.spawn(
proof_sender(port, addresses.clone(), proof_count).map(|_| vec![]), // JoinSet require having the same return type
);
set.spawn(proof_collector(port));
let res = set.join_all().await;
assert_eq!(res[1].len(), proof_count);
info!("Aborting prover...");
prover_handle.abort();
tokio::time::sleep(Duration::from_secs(1)).await;
}

1
prover/tests/mod.rs Normal file
View File

@@ -0,0 +1 @@