From e1825fdc77855d58c6473cdd0b679e0ef909a063 Mon Sep 17 00:00:00 2001
From: Roman Zeyde <me@romanzey.de>
Date: Sun, 6 Jan 2019 09:11:15 +0200
Subject: [PATCH] Fail RPC queries for "too popular" addresses

Since electrs is not optimized for this use-case, the query may take a lot
of time (on HDDs, the transaction lookups take most of the time).

The limit can be disabled by passing --txid-limit=0 flag.
---
 src/bin/electrs.rs |  2 +-
 src/config.rs      |  8 ++++++++
 src/query.rs       | 18 +++++++++++++++++-
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/src/bin/electrs.rs b/src/bin/electrs.rs
index a598986..bbe7a85 100644
--- a/src/bin/electrs.rs
+++ b/src/bin/electrs.rs
@@ -56,7 +56,7 @@ fn run_server(config: &Config) -> Result<()> {
 
     let app = App::new(store, index, daemon, &config)?;
     let tx_cache = TransactionCache::new(config.tx_cache_size);
-    let query = Query::new(app.clone(), &metrics, tx_cache);
+    let query = Query::new(app.clone(), &metrics, tx_cache, config.txid_limit);
 
     let mut server = None; // Electrum RPC server
     loop {
diff --git a/src/config.rs b/src/config.rs
index c0ac79c..408651f 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -26,6 +26,7 @@ pub struct Config {
     pub index_batch_size: usize,
     pub bulk_index_threads: usize,
     pub tx_cache_size: usize,
+    pub txid_limit: usize,
     pub server_banner: String,
 }
 
@@ -109,6 +110,12 @@ impl Config {
                     .help("Number of transactions to keep in for query LRU cache")
                     .default_value("10000")  // should be enough for a small wallet.
             )
+            .arg(
+                Arg::with_name("txid_limit")
+                    .long("txid-limit")
+                    .help("Number of transactions to lookup before returning an error, to prevent \"too popular\" addresses from causing the RPC server to get stuck (0 - disable the limit)")
+                    .default_value("100")  // should take a few seconds on a HDD
+            )
             .arg(
                 Arg::with_name("server_banner")
                     .long("server-banner")
@@ -199,6 +206,7 @@ impl Config {
             index_batch_size: value_t_or_exit!(m, "index_batch_size", usize),
             bulk_index_threads,
             tx_cache_size: value_t_or_exit!(m, "tx_cache_size", usize),
+            txid_limit: value_t_or_exit!(m, "txid_limit", usize),
             server_banner: value_t_or_exit!(m, "server_banner", String),
         };
         eprintln!("{:?}", config);
diff --git a/src/query.rs b/src/query.rs
index b1f440c..ba1ba2e 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -203,14 +203,21 @@ pub struct Query {
     app: Arc<App>,
     tracker: RwLock<Tracker>,
     tx_cache: TransactionCache,
+    txid_limit: usize,
 }
 
 impl Query {
-    pub fn new(app: Arc<App>, metrics: &Metrics, tx_cache: TransactionCache) -> Arc<Query> {
+    pub fn new(
+        app: Arc<App>,
+        metrics: &Metrics,
+        tx_cache: TransactionCache,
+        txid_limit: usize,
+    ) -> Arc<Query> {
         Arc::new(Query {
             app,
             tracker: RwLock::new(Tracker::new(metrics)),
             tx_cache,
+            txid_limit,
         })
     }
 
@@ -291,6 +298,15 @@ impl Query {
         let mut spending = vec![];
         let read_store = self.app.read_store();
         let txid_prefixes = txids_by_script_hash(read_store, script_hash);
+        // if the limit is enabled
+        if self.txid_limit > 0 {
+            if txid_prefixes.len() > self.txid_limit {
+                bail!(
+                    "{}+ transactions found, query may take a long time",
+                    txid_prefixes.len()
+                );
+            }
+        }
         for t in self.load_txns_by_prefix(read_store, txid_prefixes)? {
             funding.extend(self.find_funding_outputs(&t, script_hash));
         }