From 36440cc4d23d1d2fe1ae8a3751aab8da997f45d4 Mon Sep 17 00:00:00 2001
From: nothing0012 <rex@magiceden.io>
Date: Wed, 11 Oct 2023 11:35:52 -0700
Subject: [PATCH 1/3] Add timeout

---
 src/subcommand/server/rpc.rs | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/subcommand/server/rpc.rs b/src/subcommand/server/rpc.rs
index 849b084cc1..dac4f78d60 100644
--- a/src/subcommand/server/rpc.rs
+++ b/src/subcommand/server/rpc.rs
@@ -1,4 +1,5 @@
 use super::*;
+use tokio::time::{timeout, Duration};
 use crate::block_rarity::{
   is_palindrome, BLOCK78_BLOCK_HEIGHT, BLOCK9_BLOCK_HEIGHT, FIRST_TRANSACTION_SAT_RANGE,
   NAKAMOTO_BLOCK_HEIGHTS, PIZZA_RANGE_MAP, VINTAGE_BLOCK_HEIGHT,
@@ -24,7 +25,9 @@ pub(super) async fn handler(
 ) -> JrpcResult {
   match value.method.as_str() {
     "getHealth" => get_health(value).await,
-    "getSatRanges" => get_sat_ranges(value, index).await,
+    "getSatRanges" => {
+      timeout(Duration::from_secs(5), get_sat_ranges(value, index)).await?.map_err(timeout_error(answer_id))
+    }
     method => Ok(value.method_not_found(method)),
   }
 }
@@ -36,6 +39,13 @@ fn invalid_params(answer_id: i64, message: String) -> JrpcResult {
   ))
 }
 
+fn timeout_error(answer_id: i64) -> JrpcResult {
+  Err(JsonRpcResponse::error(
+    answer_id,
+    JsonRpcError::new(JsonRpcErrorReason::InternalError, "timeout".to_owned(), Value::default()),
+  ))
+}
+
 async fn get_health(value: JsonRpcExtractor) -> JrpcResult {
   let answer_id = value.get_answer_id();
   Ok(JsonRpcResponse::success(answer_id, "OK"))

From 1178400bfb23154ddbe53a6b556dd49c7446e1ae Mon Sep 17 00:00:00 2001
From: Michael Yu <michael@magiceden.io>
Date: Wed, 11 Oct 2023 12:01:18 -0700
Subject: [PATCH 2/3] use tower_http timeout

---
 Cargo.lock                   |  4 ++--
 Cargo.toml                   |  2 +-
 src/subcommand/server.rs     | 20 ++++++++++++++++++++
 src/subcommand/server/rpc.rs | 12 +-----------
 4 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 2cac7ab137..63204c12af 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3384,9 +3384,9 @@ dependencies = [
 
 [[package]]
 name = "tower-http"
-version = "0.4.3"
+version = "0.4.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "55ae70283aba8d2a8b411c695c437fe25b8b5e44e23e780662002fc72fb47a82"
+checksum = "61c5bb1d698276a2443e5ecfabc1008bf15a36c12e6a7176e7bf089ea9131140"
 dependencies = [
  "async-compression",
  "bitflags 2.3.3",
diff --git a/Cargo.toml b/Cargo.toml
index c751c131c5..fae1a344db 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -57,7 +57,7 @@ tempfile = "3.2.0"
 tokio = { version = "1.17.0", features = ["rt-multi-thread"] }
 tokio-stream = "0.1.9"
 tokio-util = {version = "0.7.3", features = ["compat"] }
-tower-http = { version = "0.4.0", features = ["compression-br", "compression-gzip", "cors", "set-header"] }
+tower-http = { version = "0.4.0", features = ["compression-br", "compression-gzip", "cors", "set-header", "timeout"] }
 
 rdkafka = { version = "0.33.2" }
 axum-jrpc = { version = "0.5.1", features = ["serde_json", "anyhow_error"] }
diff --git a/src/subcommand/server.rs b/src/subcommand/server.rs
index 863aca1e46..5ab1f3c46c 100644
--- a/src/subcommand/server.rs
+++ b/src/subcommand/server.rs
@@ -37,6 +37,7 @@ use {
     compression::CompressionLayer,
     cors::{Any, CorsLayer},
     set_header::SetResponseHeaderLayer,
+    timeout::TimeoutLayer,
   },
 };
 
@@ -132,6 +133,11 @@ pub(crate) struct Server {
   https: bool,
   #[clap(long, help = "Redirect HTTP traffic to HTTPS.")]
   redirect_http_to_https: bool,
+  #[clap(
+    long,
+    help = "Timeout requests after <SECONDS> seconds. Default: 30 seconds."
+  )]
+  timeout: Option<u64>,
 }
 
 impl Server {
@@ -214,6 +220,7 @@ impl Server {
             .allow_origin(Any),
         )
         .layer(CompressionLayer::new())
+        .layer(TimeoutLayer::new(Duration::from_secs(self.timeout.unwrap_or(30))))
         .with_state(server_config);
 
       match (self.http_port(), self.https_port()) {
@@ -1112,6 +1119,10 @@ mod tests {
       Self::new_with_args(&["--index-sats"], &[])
     }
 
+    fn new_with_timeout() -> Self {
+      Self::new_with_args(&[], &["--timeout", "1"])
+    }
+
     fn new_with_args(ord_args: &[&str], server_args: &[&str]) -> Self {
       Self::new_server(test_bitcoincore_rpc::spawn(), None, ord_args, server_args)
     }
@@ -1770,6 +1781,15 @@ mod tests {
     );
   }
 
+  #[test]
+  fn output_with_timeout() {
+    TestServer::new_with_timeout().assert_response_regex(
+      format!("/block/0"),
+      StatusCode::OK,
+      format!(".*<title>Block 0</title>.*<h1>Block 0</h1>.*"),
+    );
+  }
+
   #[test]
   fn output_without_sat_index() {
     let txid = "4a5e1e4baab89f3a32518a88c31bc87f618f76673e2cc77ab2127b7afdeda33b";
diff --git a/src/subcommand/server/rpc.rs b/src/subcommand/server/rpc.rs
index dac4f78d60..849b084cc1 100644
--- a/src/subcommand/server/rpc.rs
+++ b/src/subcommand/server/rpc.rs
@@ -1,5 +1,4 @@
 use super::*;
-use tokio::time::{timeout, Duration};
 use crate::block_rarity::{
   is_palindrome, BLOCK78_BLOCK_HEIGHT, BLOCK9_BLOCK_HEIGHT, FIRST_TRANSACTION_SAT_RANGE,
   NAKAMOTO_BLOCK_HEIGHTS, PIZZA_RANGE_MAP, VINTAGE_BLOCK_HEIGHT,
@@ -25,9 +24,7 @@ pub(super) async fn handler(
 ) -> JrpcResult {
   match value.method.as_str() {
     "getHealth" => get_health(value).await,
-    "getSatRanges" => {
-      timeout(Duration::from_secs(5), get_sat_ranges(value, index)).await?.map_err(timeout_error(answer_id))
-    }
+    "getSatRanges" => get_sat_ranges(value, index).await,
     method => Ok(value.method_not_found(method)),
   }
 }
@@ -39,13 +36,6 @@ fn invalid_params(answer_id: i64, message: String) -> JrpcResult {
   ))
 }
 
-fn timeout_error(answer_id: i64) -> JrpcResult {
-  Err(JsonRpcResponse::error(
-    answer_id,
-    JsonRpcError::new(JsonRpcErrorReason::InternalError, "timeout".to_owned(), Value::default()),
-  ))
-}
-
 async fn get_health(value: JsonRpcExtractor) -> JrpcResult {
   let answer_id = value.get_answer_id();
   Ok(JsonRpcResponse::success(answer_id, "OK"))

From 16d882bf7eff689fbf0b3e9efb7060b990767757 Mon Sep 17 00:00:00 2001
From: Michael Yu <michael@magiceden.io>
Date: Wed, 11 Oct 2023 12:05:41 -0700
Subject: [PATCH 3/3] fix lint

---
 src/block_rarity.rs      | 2 +-
 src/subcommand/server.rs | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/block_rarity.rs b/src/block_rarity.rs
index eb454971f3..c323791317 100644
--- a/src/block_rarity.rs
+++ b/src/block_rarity.rs
@@ -100,7 +100,7 @@ impl<'de> Deserialize<'de> for BlockRarity {
 
 pub(crate) fn is_palindrome(n: &u64) -> bool {
   let s = n.to_string();
-  if s.chars().nth(0) != s.chars().last() {
+  if s.chars().next() != s.chars().last() {
     return false;
   }
   let reversed = s.chars().rev().collect::<String>();
diff --git a/src/subcommand/server.rs b/src/subcommand/server.rs
index 5ab1f3c46c..af32409d81 100644
--- a/src/subcommand/server.rs
+++ b/src/subcommand/server.rs
@@ -1784,9 +1784,9 @@ mod tests {
   #[test]
   fn output_with_timeout() {
     TestServer::new_with_timeout().assert_response_regex(
-      format!("/block/0"),
+      "/block/0",
       StatusCode::OK,
-      format!(".*<title>Block 0</title>.*<h1>Block 0</h1>.*"),
+      ".*<title>Block 0</title>.*<h1>Block 0</h1>.*",
     );
   }