use confidential_transfer::{
ConfidentialTransfer as OtherConfidentialTransfer, ConfidentialTransferCircuit,
ConfidentialTransferTransaction,
};
use frame_support::traits::StorageMapShim;
use frame_support::{assert_ok, construct_runtime, parameter_types};
use pallet_encrypted_balance::{Account, AccountData};
use pallet_plonk::FullcodecRng;
use rand::Rng;
use rand_core::{OsRng, SeedableRng};
use sp_core::H256;
use sp_runtime::{
testing::Header,
traits::{BlakeTwo256, IdentityLookup},
};
use zero_crypto::behave::Group;
use zero_elgamal::EncryptedNumber;
use zero_jubjub::{Fp, JubJubAffine, GENERATOR_EXTENDED};
use zero_plonk::prelude::Compiler;
type UncheckedExtrinsic = frame_system::mocking::MockUncheckedExtrinsic<TestRuntime>;
type Block = frame_system::mocking::MockBlock<TestRuntime>;
construct_runtime!(
pub enum TestRuntime where
Block = Block,
NodeBlock = Block,
UncheckedExtrinsic = UncheckedExtrinsic,
{
System: frame_system::{Module, Call, Config, Storage, Event<T>},
Plonk: pallet_plonk::{Module, Call, Storage, Event<T>},
EncryptedCurrency: pallet_encrypted_balance::{Module, Call, Storage, Config<T>, Event<T>},
ConfidentialTransfer: confidential_transfer::{Module, Call, Storage, Event<T>},
}
);
parameter_types! {
pub const BlockHashCount: u64 = 250;
pub BlockWeights: frame_system::limits::BlockWeights =
frame_system::limits::BlockWeights::simple_max(1024);
}
impl frame_system::Config for TestRuntime {
type BaseCallFilter = ();
type BlockWeights = ();
type BlockLength = ();
type Origin = Origin;
type Index = u64;
type Call = Call;
type BlockNumber = u64;
type Hash = H256;
type Hashing = BlakeTwo256;
type AccountId = u64;
type Lookup = IdentityLookup<Self::AccountId>;
type Header = Header;
type Event = Event;
type BlockHashCount = BlockHashCount;
type DbWeight = ();
type Version = ();
type PalletInfo = PalletInfo;
type AccountData = ();
type OnNewAccount = ();
type OnKilledAccount = ();
type SystemWeightInfo = ();
type SS58Prefix = ();
}
impl pallet_plonk::Config for TestRuntime {
type CustomCircuit = ConfidentialTransferCircuit;
type Event = Event;
}
impl pallet_encrypted_balance::Config for TestRuntime {
type EncryptedBalance = EncryptedNumber;
type Event = Event;
type AccountStore = StorageMapShim<
Account<TestRuntime>,
frame_system::Provider<TestRuntime>,
u64,
AccountData<Self::EncryptedBalance>,
>;
type WeightInfo = ();
}
impl confidential_transfer::Config for TestRuntime {
type Plonk = Plonk;
type EncryptedCurrency = EncryptedCurrency;
type Event = Event;
}
fn new_test_ext(
alice_address: u64,
alice_private_key: Fp,
alice_balance: u16,
alice_radomness: Fp,
bob_private_key: Fp,
bob_address: u64,
bob_balance: u16,
bob_radomness: Fp,
) -> sp_io::TestExternalities {
let alice_balance =
EncryptedNumber::encrypt(alice_private_key, alice_balance.into(), alice_radomness);
let bob_balance = EncryptedNumber::encrypt(bob_private_key, bob_balance.into(), bob_radomness);
let mut t = frame_system::GenesisConfig::default()
.build_storage::<TestRuntime>()
.unwrap();
pallet_encrypted_balance::GenesisConfig::<TestRuntime> {
balances: vec![(alice_address, alice_balance), (bob_address, bob_balance)],
}
.assimilate_storage(&mut t)
.unwrap();
let mut ext = sp_io::TestExternalities::new(t);
ext.execute_with(|| System::set_block_number(1));
ext
}
fn generate_default_test_data() -> (u64, Fp, u16, Fp, Fp, u64, u16, Fp, u16, u16, u16, Fp) {
let mut rng = rand::thread_rng();
let alice_address = rng.gen::<u64>();
let alice_private_key = Fp::random(OsRng);
let alice_balance = rng.gen::<u16>();
let alice_radomness = Fp::random(OsRng);
let bob_private_key = Fp::random(OsRng);
let bob_address = rng.gen::<u64>();
let bob_balance = rng.gen::<u16>();
let bob_radomness = Fp::random(OsRng);
let transfer_amount = rng.gen_range(0..alice_balance);
let alice_after_balance = alice_balance - transfer_amount;
let bob_after_balance = bob_balance + transfer_amount;
let transfer_randomness = Fp::random(OsRng);
(
alice_address,
alice_private_key,
alice_balance,
alice_radomness,
bob_private_key,
bob_address,
bob_balance,
bob_radomness,
transfer_amount,
alice_after_balance,
bob_after_balance,
transfer_randomness,
)
}
fn get_rng() -> FullcodecRng {
FullcodecRng::from_seed([
0x59, 0x62, 0xbe, 0x5d, 0x76, 0x3d, 0x31, 0x8d, 0x17, 0xdb, 0x37, 0x32, 0x54, 0x06, 0xbc,
0xe5,
])
}
fn main() {
let k = 14;
let label = b"verify";
let mut rng = get_rng();
let (
alice_address,
alice_private_key,
alice_balance,
alice_radomness,
bob_private_key,
bob_address,
bob_balance,
bob_radomness,
transfer_amount,
alice_after_balance,
bob_after_balance,
transfer_randomness,
) = generate_default_test_data();
new_test_ext(
alice_address,
alice_private_key,
alice_balance,
alice_radomness,
bob_private_key,
bob_address,
bob_balance,
bob_radomness,
)
.execute_with(|| {
let alice_encrypted_balance = ConfidentialTransfer::total_balance(&alice_address);
let alice_raw_balance = alice_encrypted_balance.decrypt(alice_private_key);
let bob_encrypted_balance = ConfidentialTransfer::total_balance(&bob_address);
let bob_raw_balance = bob_encrypted_balance.decrypt(bob_private_key);
assert_eq!(alice_raw_balance.unwrap() as u16, alice_balance);
assert_eq!(bob_raw_balance.unwrap() as u16, bob_balance);
let result =
ConfidentialTransfer::trusted_setup(Origin::signed(alice_address), k, rng.clone());
assert_ok!(result);
let pp = Plonk::public_parameter().unwrap();
let alice_public_key = GENERATOR_EXTENDED * alice_private_key;
let bob_public_key = GENERATOR_EXTENDED * bob_private_key;
let transfer_amount_scalar = Fp::from(transfer_amount as u64);
let alice_after_balance_scalar = Fp::from(alice_after_balance as u64);
let alice_balance =
EncryptedNumber::encrypt(alice_private_key, alice_balance.into(), alice_radomness);
let alice_transfer_amount = EncryptedNumber::encrypt(
alice_private_key,
transfer_amount.into(),
transfer_randomness,
);
let bob_encrypted_transfer_amount =
(GENERATOR_EXTENDED * transfer_amount_scalar) + (bob_public_key * transfer_randomness);
let alice_public_key = JubJubAffine::from(alice_public_key);
let bob_public_key = JubJubAffine::from(bob_public_key);
let bob_encrypted_transfer_amount = JubJubAffine::from(bob_encrypted_transfer_amount);
let bob_encrypted_transfer_amount_other = (GENERATOR_EXTENDED * transfer_randomness).into();
let confidential_transfer_circuit = ConfidentialTransferCircuit::new(
alice_public_key,
bob_public_key,
alice_balance,
alice_transfer_amount,
bob_encrypted_transfer_amount,
alice_private_key,
transfer_amount_scalar,
alice_after_balance_scalar,
transfer_randomness,
);
let prover = Compiler::compile::<ConfidentialTransferCircuit>(&pp, label)
.expect("failed to compile circuit");
let proof = prover
.0
.prove(&mut rng, &confidential_transfer_circuit)
.expect("failed to prove");
let transaction_params = ConfidentialTransferTransaction::new(
alice_public_key,
bob_public_key,
alice_transfer_amount,
bob_encrypted_transfer_amount,
bob_encrypted_transfer_amount_other,
);
let result = ConfidentialTransfer::confidential_transfer(
Origin::signed(alice_address),
bob_address,
proof.0,
transaction_params,
);
assert_ok!(result);
let alice_balance = ConfidentialTransfer::total_balance(&alice_address);
let alice_raw_balance = alice_balance.decrypt(alice_private_key);
let bob_balance = ConfidentialTransfer::total_balance(&bob_address);
let bob_raw_balance = bob_balance.decrypt(bob_private_key);
assert_eq!(alice_raw_balance.unwrap() as u16, alice_after_balance);
assert_eq!(bob_raw_balance.unwrap() as u16, bob_after_balance);
})
}