diff --git a/src/config.rs b/src/config.rs index b4a9509d3..409d1c09c 100644 --- a/src/config.rs +++ b/src/config.rs @@ -74,6 +74,7 @@ pub struct Config { pub utxos_limit: usize, pub electrum_txs_limit: usize, pub electrum_subscription_limit: usize, + pub electrum_checkpoint_proof_concurrency_limit: usize, pub electrum_banner: String, pub rpc_logging: RpcLogging, pub zmq_addr: Option, @@ -292,6 +293,12 @@ impl Config { .help("Maximum number of scripthash subscriptions a single Electrum connection may hold. Every subscription costs a history lookup on each new block, so an unbounded count lets one client impose unbounded recurring work. Re-subscribing to an already-tracked scripthash is always allowed. 0 = unlimited.") .default_value("10000") .takes_value(true) + ).arg( + Arg::with_name("electrum_checkpoint_proof_concurrency_limit") + .long("electrum-checkpoint-proof-concurrency-limit") + .help("Maximum number of blockchain.block.header(s) checkpoint Merkle proof builds (triggered by a non-zero cp_height) allowed to run at once, process-wide. Each build hashes every header from genesis up to cp_height, so an unbounded count lets concurrent cheap requests pin every CPU core. Requests past the cap fail immediately rather than queueing. 0 = reject all such requests.") + .default_value("2") + .takes_value(true) ).arg( Arg::with_name("electrum_banner") .long("electrum-banner") @@ -589,6 +596,11 @@ impl Config { electrum_rpc_max_request_num_bytes, electrum_txs_limit: value_t_or_exit!(m, "electrum_txs_limit", usize), electrum_subscription_limit: value_t_or_exit!(m, "electrum_subscription_limit", usize), + electrum_checkpoint_proof_concurrency_limit: value_t_or_exit!( + m, + "electrum_checkpoint_proof_concurrency_limit", + usize + ), electrum_banner, rpc_logging: { let params = RpcLogging { diff --git a/src/electrum/server.rs b/src/electrum/server.rs index 1ba8ea83b..c843cb522 100644 --- a/src/electrum/server.rs +++ b/src/electrum/server.rs @@ -183,6 +183,7 @@ struct Connection { txs_limit: usize, subscription_limit: usize, max_request_bytes: usize, + checkpoint_proof_concurrency_limit: usize, #[cfg(feature = "electrum-discovery")] discovery: Option>, rpc_logging: RpcLogging, @@ -206,6 +207,7 @@ impl Connection { txs_limit: usize, subscription_limit: usize, max_request_bytes: usize, + checkpoint_proof_concurrency_limit: usize, #[cfg(feature = "electrum-discovery")] discovery: Option>, rpc_logging: RpcLogging, salt: String, @@ -221,6 +223,7 @@ impl Connection { txs_limit, subscription_limit, max_request_bytes, + checkpoint_proof_concurrency_limit, #[cfg(feature = "electrum-discovery")] discovery, rpc_logging, @@ -309,7 +312,12 @@ impl Connection { if cp_height == 0 { return Ok(json!(raw_header_hex)); } - let (branch, root) = get_header_merkle_proof(self.query.chain(), height, cp_height)?; + let (branch, root) = get_header_merkle_proof( + self.query.chain(), + height, + cp_height, + self.checkpoint_proof_concurrency_limit, + )?; Ok(json!({ "header": raw_header_hex, @@ -341,8 +349,12 @@ impl Connection { })); } - let (branch, root) = - get_header_merkle_proof(self.query.chain(), start_height + (count - 1), cp_height)?; + let (branch, root) = get_header_merkle_proof( + self.query.chain(), + start_height + (count - 1), + cp_height, + self.checkpoint_proof_concurrency_limit, + )?; Ok(json!({ "count": headers.len(), @@ -1238,6 +1250,7 @@ impl RPC { let txs_limit = config.electrum_txs_limit; let subscription_limit = config.electrum_subscription_limit; let max_request_bytes = config.electrum_rpc_max_request_num_bytes; + let checkpoint_proof_concurrency_limit = config.electrum_checkpoint_proof_concurrency_limit; let conn_max_age = config.electrum_rpc_conn_max_age; RPC { @@ -1290,6 +1303,7 @@ impl RPC { txs_limit, subscription_limit, max_request_bytes, + checkpoint_proof_concurrency_limit, #[cfg(feature = "electrum-discovery")] discovery, rpc_logging, diff --git a/src/util/electrum_merkle.rs b/src/util/electrum_merkle.rs index 52e0a825a..7bf813afd 100644 --- a/src/util/electrum_merkle.rs +++ b/src/util/electrum_merkle.rs @@ -2,9 +2,40 @@ use crate::chain::{BlockHash, Txid}; use crate::errors::*; use crate::new_index::ChainQuery; use bitcoin::hashes::{sha256d::Hash as Sha256dHash, Hash}; +use std::sync::atomic::{AtomicUsize, Ordering}; use electrs_macros::trace; +static INFLIGHT_CHECKPOINT_PROOFS: AtomicUsize = AtomicUsize::new(0); + +struct InflightCheckpointProofGuard; + +impl InflightCheckpointProofGuard { + fn acquire(limit: usize) -> Result { + let mut current = INFLIGHT_CHECKPOINT_PROOFS.load(Ordering::Relaxed); + loop { + if current >= limit { + bail!("too many concurrent checkpoint merkle proof requests, try again later"); + } + match INFLIGHT_CHECKPOINT_PROOFS.compare_exchange_weak( + current, + current + 1, + Ordering::AcqRel, + Ordering::Relaxed, + ) { + Ok(_) => return Ok(Self), + Err(observed) => current = observed, + } + } + } +} + +impl Drop for InflightCheckpointProofGuard { + fn drop(&mut self) { + INFLIGHT_CHECKPOINT_PROOFS.fetch_sub(1, Ordering::AcqRel); + } +} + #[trace] pub fn get_tx_merkle_proof( chain: &ChainQuery, @@ -29,6 +60,7 @@ pub fn get_header_merkle_proof( chain: &ChainQuery, height: usize, cp_height: usize, + checkpoint_proof_concurrency_limit: usize, ) -> Result<(Vec, Sha256dHash)> { if cp_height < height { bail!("cp_height #{} < height #{}", cp_height, height); @@ -43,6 +75,8 @@ pub fn get_header_merkle_proof( ); } + let _guard = InflightCheckpointProofGuard::acquire(checkpoint_proof_concurrency_limit)?; + let heights: Vec = (0..=cp_height).collect(); let header_hashes: Vec = heights .into_iter() @@ -107,3 +141,34 @@ fn create_merkle_branch_and_root( } (merkle, hashes[0]) } + +#[cfg(test)] +mod tests { + use super::*; + + // Single test so the shared static isn't touched by other tests running + // concurrently in the same process. + #[test] + fn checkpoint_proof_guard_caps_concurrency_and_releases_on_drop() { + const LIMIT: usize = 2; + let mut guards = Vec::new(); + for _ in 0..LIMIT { + guards.push(InflightCheckpointProofGuard::acquire(LIMIT).unwrap()); + } + + // All permits are taken: the next caller must be rejected outright + // instead of blocking or panicking. + assert!(InflightCheckpointProofGuard::acquire(LIMIT).is_err()); + + // Releasing one permit (drop, e.g. on early return via `?`) must let + // the next acquire through. + guards.pop(); + let extra = InflightCheckpointProofGuard::acquire(LIMIT).unwrap(); + + // Back at the cap: still no free permits. + assert!(InflightCheckpointProofGuard::acquire(LIMIT).is_err()); + + drop(extra); + drop(guards); + } +} diff --git a/tests/common.rs b/tests/common.rs index 55cd457d4..7fddf1599 100644 --- a/tests/common.rs +++ b/tests/common.rs @@ -115,6 +115,7 @@ impl TestRunner { utxos_limit: 100, electrum_txs_limit: 100, electrum_subscription_limit: 10_000, + electrum_checkpoint_proof_concurrency_limit: 2, electrum_banner: "".into(), rpc_logging: RpcLogging::default(), zmq_addr: None,