Skip to content

Commit

Permalink
use deteministic rng optionally
Browse files Browse the repository at this point in the history
  • Loading branch information
@contrun
contrun committed Dec 22, 2023
1 parent 7882493 commit 2ce0e5a
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 17 deletions.
47 changes: 36 additions & 11 deletions tests/auth-c-tests/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ use ckb_types::{
use dyn_clone::{clone_trait_object, DynClone};
use hex;
use log::{Metadata, Record};
use rand::{distributions::Standard, thread_rng, Rng};
use rand::{distributions::Standard, Rng};
use secp256k1;
use serde::{Deserialize, Serialize};
use sha3::{Digest, Keccak256};
Expand All @@ -39,7 +39,6 @@ type BtcNetwork = bitcoin::Network;

pub const MAX_CYCLES: u64 = std::u64::MAX;
pub const SIGNATURE_SIZE: usize = 65;
pub const RNG_SEED: u64 = 42;
pub const SOLANA_MAXIMUM_UNWRAPPED_SIGNATURE_SIZE: usize = 510;
pub const SOLANA_MAXIMUM_WRAPPED_SIGNATURE_SIZE: usize =
SOLANA_MAXIMUM_UNWRAPPED_SIGNATURE_SIZE + 2;
Expand Down Expand Up @@ -97,6 +96,32 @@ pub mod auth_program {
}
}

pub use rng::get_rng;
pub mod rng {
use rand::{thread_rng, RngCore};
use ref_thread_local::ref_thread_local;
use ref_thread_local::RefThreadLocal;

ref_thread_local! {
static managed RNG_SEED: Option<u64> = None;
}

pub fn get_rng() -> rand::rngs::SmallRng {
let seed = RNG_SEED.borrow().unwrap_or(thread_rng().next_u64());
rand::SeedableRng::seed_from_u64(seed)
}

pub fn set_seed(seed: u64) {
let mut p = RNG_SEED.borrow_mut();
*p = Some(seed);
}

pub fn clear_seed() {
let mut p = RNG_SEED.borrow_mut();
*p = None;
}
}

fn _dbg_print_mem(data: &Vec<u8>, name: &str) {
print!("rustdbg {}: (size:{})\n", name, data.len());
let mut count = 0;
Expand Down Expand Up @@ -274,7 +299,7 @@ pub fn sign_tx_by_input_group(
begin_index: usize,
len: usize,
) -> TransactionView {
let mut rng = thread_rng();
let mut rng = get_rng();
let tx_hash = tx.hash();
let mut signed_witnesses: Vec<packed::Bytes> = tx
.inputs()
Expand Down Expand Up @@ -454,7 +479,7 @@ pub fn gen_tx_with_pub_key_hash(
let lock_args = gen_args_with_pub_key_hash(&config, hash);
// Note that we use deterministic here to ensure the same transaction structure
// is generated.
let mut rng: rand::rngs::SmallRng = rand::SeedableRng::seed_from_u64(RNG_SEED);
let mut rng = get_rng();

gen_tx_with_grouped_args(
dummy,
Expand All @@ -467,7 +492,7 @@ pub fn gen_tx_with_pub_key_hash(
pub fn gen_tx(dummy: &mut DummyDataLoader, config: &TestConfig) -> TransactionView {
let lock_args = gen_args(&config);

let mut rng = thread_rng();
let mut rng = get_rng();
gen_tx_with_grouped_args(
dummy,
vec![(lock_args, config.sign_size as usize)],
Expand Down Expand Up @@ -625,7 +650,7 @@ pub fn do_gen_args(config: &TestConfig, pub_key_hash: Option<Vec<u8>>) -> Bytes
.pubkey_hash
.copy_from_slice(pub_hash.as_slice());
} else {
let mut rng = thread_rng();
let mut rng = get_rng();
let incorrect_pubkey = {
let mut buf = [0u8; 32];
rng.fill(&mut buf);
Expand Down Expand Up @@ -884,7 +909,7 @@ pub struct EthereumAuth {
impl EthereumAuth {
fn new() -> Box<EthereumAuth> {
let generator: secp256k1::Secp256k1<secp256k1::All> = secp256k1::Secp256k1::new();
let mut rng = thread_rng();
let mut rng = get_rng();
let (privkey, pubkey) = generator.generate_keypair(&mut rng);
Box::new(EthereumAuth {
privkey,
Expand Down Expand Up @@ -993,7 +1018,7 @@ pub struct TronAuth {
impl TronAuth {
fn new() -> Box<dyn Auth> {
let generator: secp256k1::Secp256k1<secp256k1::All> = secp256k1::Secp256k1::new();
let mut rng = thread_rng();
let mut rng = get_rng();
let (privkey, pubkey) = generator.generate_keypair(&mut rng);
Box::new(TronAuth { privkey, pubkey })
}
Expand Down Expand Up @@ -1051,7 +1076,7 @@ impl BitcoinAuth {
}

pub fn new_rng_key(v_type: BitcoinSignVType, btc_network: BtcNetwork) -> Self {
let mut rng = thread_rng();
let mut rng = get_rng();
let mut secret_key = [0u8; 32];
rng.fill(&mut secret_key);

Expand Down Expand Up @@ -1446,7 +1471,7 @@ pub struct MoneroAuth {
impl MoneroAuth {
pub fn new() -> Box<MoneroAuth> {
fn get_random_key_pair() -> monero::KeyPair {
let mut rng = thread_rng();
let mut rng = get_rng();
let mut seed = vec![0; 32];
let spend_key = loop {
rng.fill(seed.as_mut_slice());
Expand Down Expand Up @@ -2054,7 +2079,7 @@ pub struct SchnorrAuth {
impl SchnorrAuth {
pub fn new() -> Box<dyn Auth> {
let generator: secp256k1::Secp256k1<secp256k1::All> = secp256k1::Secp256k1::new();
let mut rng = thread_rng();
let mut rng = get_rng();
let (privkey, pubkey) = generator.generate_keypair(&mut rng);
Box::new(SchnorrAuth { privkey, pubkey })
}
Expand Down
13 changes: 7 additions & 6 deletions tests/auth-c-tests/src/tests/mod.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#![allow(unused_imports)]
#![allow(dead_code)]

use crate::get_rng;
use ckb_auth_rs::EntryCategoryType;
use ckb_chain_spec::consensus::{Consensus, ConsensusBuilder};
use ckb_crypto::secp::{Generator, Privkey, Pubkey};
Expand All @@ -11,7 +12,7 @@ use ckb_types::{
H256,
};
use log::{Level, LevelFilter, Metadata, Record};
use rand::{thread_rng, Rng};
use rand::Rng;
use sha3::{digest::generic_array::typenum::private::IsEqualPrivate, Digest, Keccak256};
use std::sync::Arc;

Expand Down Expand Up @@ -81,7 +82,7 @@ fn unit_test_multiple_group(auth: &Box<dyn Auth>, run_type: EntryCategoryType) {

let config = TestConfig::new(auth, run_type, 1);

let mut rng = thread_rng();
let mut rng = get_rng();
let tx = gen_tx_with_grouped_args(
&mut data_loader,
vec![
Expand Down Expand Up @@ -264,7 +265,7 @@ fn bitcoin_pubkey_recid_verify() {
let sign = priv_key.sign_recoverable(&msg).expect("sign").serialize();
assert_eq!(sign.len(), 65);

let mut rng = rand::thread_rng();
let mut rng = get_rng();
let mut recid: u8 = rng.gen_range(0, 4);
while recid == sign[64] && recid < 31 {
recid = rng.gen_range(0, 4);
Expand Down Expand Up @@ -354,7 +355,7 @@ fn secp256r1_raw_verify() {
// is the data I used to verify this.
//
// TODO: fix this.
//
//
// unit_test_common(AuthAlgorithmIdType::Secp256r1Raw);
}

Expand Down Expand Up @@ -384,7 +385,7 @@ fn convert_eth_error() {
}

let generator: secp256k1::Secp256k1<secp256k1::All> = secp256k1::Secp256k1::new();
let mut rng = thread_rng();
let mut rng = get_rng();
let (privkey, pubkey) = generator.generate_keypair(&mut rng);

let auth: Box<dyn Auth> = Box::new(EthConverFaileAuth {
Expand Down Expand Up @@ -430,7 +431,7 @@ fn convert_tron_error() {
}

let generator: secp256k1::Secp256k1<secp256k1::All> = secp256k1::Secp256k1::new();
let mut rng = thread_rng();
let mut rng = get_rng();
let (privkey, pubkey) = generator.generate_keypair(&mut rng);
let auth: Box<dyn Auth> = Box::new(TronConverFaileAuth {
0: TronAuth { privkey, pubkey },
Expand Down

0 comments on commit 2ce0e5a

Please sign in to comment.