Add rate limiting

This commit is contained in:
Patryk Bęza 2024-09-03 12:28:47 +02:00
parent 6e77578695
commit 752802da0b
No known key found for this signature in database
GPG key ID: 9AD1B44D9F6258EC

View file

@ -8,10 +8,12 @@ use clap::Parser;
use reqwest::Client; use reqwest::Client;
use secp256k1::{constants::PUBLIC_KEY_SIZE, ecdsa::Signature, Message, PublicKey}; use secp256k1::{constants::PUBLIC_KEY_SIZE, ecdsa::Signature, Message, PublicKey};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::time::Duration;
use teepot::{ use teepot::{
client::TcbLevel, client::TcbLevel,
sgx::{tee_qv_get_collateral, verify_quote_with_collateral, QuoteVerificationResult}, sgx::{tee_qv_get_collateral, verify_quote_with_collateral, QuoteVerificationResult},
}; };
use tokio::time::sleep;
use url::Url; use url::Url;
use zksync_basic_types::{L1BatchNumber, H256}; use zksync_basic_types::{L1BatchNumber, H256};
use zksync_types::L2ChainId; use zksync_types::L2ChainId;
@ -28,11 +30,14 @@ struct Arguments {
#[clap(short = 'n', long = "batch-number", value_parser = parse_batch_range)] #[clap(short = 'n', long = "batch-number", value_parser = parse_batch_range)]
batch_range: (L1BatchNumber, L1BatchNumber), batch_range: (L1BatchNumber, L1BatchNumber),
/// URL of the RPC server to query for the batch attestation and signature. /// URL of the RPC server to query for the batch attestation and signature.
#[clap(short, long)] #[clap(short = 'u', long)]
rpc_url: Url, rpc_url: Url,
/// Chain ID of the network to query. /// Chain ID of the network to query.
#[clap(short, long, default_value_t = L2ChainId::default().as_u64())] #[clap(short = 'c', long, default_value_t = L2ChainId::default().as_u64())]
chain_id: u64, chain_id: u64,
/// Rate limit between requests in milliseconds.
#[clap(short = 'r', long, default_value_t = Duration::from_millis(0), value_parser = parse_duration)]
rate_limit: Duration,
} }
fn parse_batch_range(s: &str) -> Result<(L1BatchNumber, L1BatchNumber)> { fn parse_batch_range(s: &str) -> Result<(L1BatchNumber, L1BatchNumber)> {
@ -61,6 +66,11 @@ fn parse_batch_range(s: &str) -> Result<(L1BatchNumber, L1BatchNumber)> {
} }
} }
fn parse_duration(s: &str) -> Result<Duration> {
let millis = s.parse()?;
Ok(Duration::from_millis(millis))
}
trait JsonRpcClient { trait JsonRpcClient {
async fn get_root_hash(&self, batch_number: L1BatchNumber) -> Result<H256>; async fn get_root_hash(&self, batch_number: L1BatchNumber) -> Result<H256>;
// TODO implement get_tee_proofs(batch_number, tee_type) once https://crates.io/crates/zksync_web3_decl crate is updated // TODO implement get_tee_proofs(batch_number, tee_type) once https://crates.io/crates/zksync_web3_decl crate is updated
@ -163,6 +173,8 @@ async fn main() -> Result<()> {
verify_signature(&proof.signature, public_key, root_hash)?; verify_signature(&proof.signature, public_key, root_hash)?;
println!(); println!();
} }
sleep(args.rate_limit).await;
} }
Ok(()) Ok(())