Skip to content

Commit

Permalink
Simplify random enum member selection
Browse files Browse the repository at this point in the history
  • Loading branch information
@popzxc
popzxc committed Apr 7, 2021
1 parent 0872cd5 commit 1511ca3
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 74 deletions.
96 changes: 33 additions & 63 deletions core/tests/loadnext/src/command/tx_command.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
use std::iter;

use num::BigUint;
use rand::{seq::SliceRandom, Rng};

Expand All @@ -24,39 +22,22 @@ impl TxType {
/// Generates a random transaction type. Not all the variants have the equal chance to be generated;
/// specifically transfers are made more likely.
pub fn random(rng: &mut LoadtestRng) -> Self {
// All available options.
let mut options = vec![
Self::Deposit,
Self::TransferToNew,
Self::TransferToExisting,
Self::WithdrawToSelf,
Self::WithdrawToOther,
Self::FullExit,
Self::ChangePubKey,
// All available options together with their weight.
// `TransferToNew` and `TransferToExisting` the most likely options.
const DEFAULT_WEIGHT: usize = 1;
const HIGH_WEIGHT: usize = 3;
let options = vec![
(Self::Deposit, DEFAULT_WEIGHT),
(Self::TransferToNew, HIGH_WEIGHT),
(Self::TransferToExisting, HIGH_WEIGHT),
(Self::WithdrawToSelf, DEFAULT_WEIGHT),
(Self::WithdrawToOther, DEFAULT_WEIGHT),
(Self::FullExit, DEFAULT_WEIGHT),
(Self::ChangePubKey, DEFAULT_WEIGHT),
];

// Make `TransferToNew` and `TransferToExisting` the most likely options
// by adding them multiple times.
let transfer_to_new_likehood = 0.3f64;
let transfer_to_existing_likehood = 0.4f64;

// We are ignoring the fact that variables in fact rely on each other; it's not that important for our purposes.
let required_transfer_to_new_copies =
Self::required_amount_of_copies(&options, transfer_to_new_likehood);
let required_transfer_to_existing_copies =
Self::required_amount_of_copies(&options, transfer_to_existing_likehood);
let total_new_elements =
required_transfer_to_new_copies + required_transfer_to_existing_copies;

options.reserve(total_new_elements);

options.extend(iter::repeat(Self::TransferToNew).take(required_transfer_to_new_copies));
options.extend(
iter::repeat(Self::TransferToExisting).take(required_transfer_to_existing_copies),
);

// Now we can get weighted element by simply choosing the random value from the vector.
options.choose(rng).copied().unwrap()
options.choose_weighted(rng, |item| item.1).unwrap().0
}

/// Generates a random transaction type that can be a part of the batch.
Expand All @@ -70,12 +51,6 @@ impl TxType {
}
}
}

fn required_amount_of_copies(options: &[Self], likehood: f64) -> usize {
// This value will be truncated down, but it will be compensated by the fact
// that element is already inserted into `options`.
(options.len() as f64 * likehood) as usize
}
}

/// Modifier to be applied to the transaction in order to make it incorrect.
Expand Down Expand Up @@ -112,28 +87,30 @@ pub enum ExpectedOutcome {

impl IncorrectnessModifier {
pub fn random(rng: &mut LoadtestRng) -> Self {
if Self::no_modifier(rng) {
return Self::None;
}

let options = &[
Self::ZeroFee,
Self::IncorrectZkSyncSignature,
Self::IncorrectEthSignature,
Self::NonExistentToken,
Self::TooBigAmount,
Self::NotPackableAmount,
Self::NotPackableFeeAmount,
];

options.choose(rng).copied().unwrap()
}

fn no_modifier(rng: &mut LoadtestRng) -> bool {
// 90% of transactions should be correct.
const NO_MODIFIER_PROBABILITY: f32 = 0.9f32;
// Amount of elements in the enum.
const MODIFIERS_AMOUNT: usize = 7;

let chance = rng.gen_range(0f32, 1f32);
if chance <= NO_MODIFIER_PROBABILITY {
return Self::None;
}

let modifier_type = rng.gen_range(0, MODIFIERS_AMOUNT);

match modifier_type {
0 => Self::ZeroFee,
1 => Self::IncorrectZkSyncSignature,
2 => Self::IncorrectEthSignature,
3 => Self::NonExistentToken,
4 => Self::TooBigAmount,
5 => Self::NotPackableAmount,
6 => Self::NotPackableFeeAmount,
_ => unreachable!("Unexpected modifier type number"),
}
chance <= NO_MODIFIER_PROBABILITY
}

pub fn expected_outcome(self) -> ExpectedOutcome {
Expand Down Expand Up @@ -219,9 +196,6 @@ impl TxCommand {
command.to = own_address;
}

// `ChangePubKey` does not have a 2FA signature.
let cpk_incorrect_signature = command.command_type == TxType::ChangePubKey
&& command.modifier == IncorrectnessModifier::IncorrectEthSignature;
// Transactions that have no amount field.
let no_amount_field = matches!(command.command_type, TxType::ChangePubKey)
&& matches!(
Expand All @@ -239,11 +213,7 @@ impl TxCommand {
== IncorrectnessModifier::NotPackableAmount;

// Check whether generator modifier does not make sense.
if cpk_incorrect_signature
|| no_amount_field
|| incorrect_priority_op
|| unpackable_withdrawal
{
if no_amount_field || incorrect_priority_op || unpackable_withdrawal {
command.modifier = IncorrectnessModifier::None;
}

Expand Down
32 changes: 21 additions & 11 deletions core/tests/loadnext/src/corrupted_tx.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use zksync::utils::{
private_key_from_seed,
};
use zksync_types::{
tx::{PackedEthSignature, TxSignature},
tx::{ChangePubKeyECDSAData, ChangePubKeyEthAuthData, PackedEthSignature, TxSignature},
TokenId, ZkSyncTx, H256,
};

Expand All @@ -19,8 +19,8 @@ pub trait Corrupted: Sized {
/// Replaces the zkSync signature with an incorrect one.
fn bad_zksync_signature(self) -> Self;
/// Replaces the zkSync 2FA ECDSA signature with an incorrect one.
/// Note that it only affects 2FA signatures; signature in ChangePubKey transaction will not be affected.
fn bad_eth_signature(self) -> Self;
/// In case of ChangePubKey, Ethereum signature inside of the transaction will be affected.
fn bad_eth_signature(self, eth_pk: H256, token_symbol: &str, decimals: u8) -> Self;
/// Replaces the transaction token with the non-existing one.
/// In case of `ChangePubKey` transaction it affects the `fee_token`.
fn nonexistent_token(self, eth_pk: H256, token_symbol: &str, decimals: u8) -> Self;
Expand Down Expand Up @@ -48,7 +48,9 @@ pub trait Corrupted: Sized {
) -> Self {
match modifier {
IncorrectnessModifier::None => self,
IncorrectnessModifier::IncorrectEthSignature => self.bad_eth_signature(),
IncorrectnessModifier::IncorrectEthSignature => {
self.bad_eth_signature(eth_pk, token_symbol, decimals)
}
IncorrectnessModifier::IncorrectZkSyncSignature => self.bad_zksync_signature(),
IncorrectnessModifier::NonExistentToken => {
self.nonexistent_token(eth_pk, token_symbol, decimals)
Expand Down Expand Up @@ -103,15 +105,22 @@ impl Corrupted for (ZkSyncTx, Option<PackedEthSignature>) {
}
}

fn bad_eth_signature(self) -> Self {
fn bad_eth_signature(mut self, eth_pk: H256, token_symbol: &str, decimals: u8) -> Self {
let private_key = H256::random();
let message = b"bad message";
let (tx, eth_signature) = self;
let bad_signature = PackedEthSignature::sign(&private_key, message).ok();

if let ZkSyncTx::ChangePubKey(cpk_tx) = &mut self.0 {
let signature_data = ChangePubKeyECDSAData {
eth_signature: bad_signature.clone().unwrap(),
batch_hash: Default::default(),
};
cpk_tx.eth_auth_data = Some(ChangePubKeyEthAuthData::ECDSA(signature_data))
}

(
tx,
eth_signature.and_then(|_| PackedEthSignature::sign(&private_key, message).ok()),
)
self.resign(eth_pk, token_symbol, decimals);
let (tx, eth_signature) = self;
(tx, eth_signature.and(bad_signature))
}

fn bad_zksync_signature(mut self) -> Self {
Expand Down Expand Up @@ -362,7 +371,8 @@ mod tests {
let transfer = create_transfer(&account);
let current_eth_signature = transfer.1.clone();

let (_modified_transfer, new_eth_signature) = transfer.bad_eth_signature();
let (_modified_transfer, new_eth_signature) =
transfer.bad_eth_signature(account.eth_account_data.unwrap_eoa_pk(), "ETH", 18);

assert_ne!(current_eth_signature, new_eth_signature);
}
Expand Down

0 comments on commit 1511ca3

Please sign in to comment.