From fb8ebcf6fbaa25d2f80334019c308f21b9addd26 Mon Sep 17 00:00:00 2001
From: Kirill Fomichev <fanatid@ya.ru>
Date: Fri, 5 Jul 2024 23:59:33 -0500
Subject: [PATCH] rpc: add rollback to getLatestBlockhash (agave#2023)

---
 rpc-client-api/src/config.rs             |  9 +++++++++
 rpc-client/src/nonblocking/rpc_client.rs | 20 ++++++++++++++++++++
 rpc/src/rpc.rs                           | 20 ++++++++++++++++----
 3 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/rpc-client-api/src/config.rs b/rpc-client-api/src/config.rs
index 2b297226b9..81d5fa74e7 100644
--- a/rpc-client-api/src/config.rs
+++ b/rpc-client-api/src/config.rs
@@ -358,3 +358,12 @@ pub struct RpcContextConfig {
 pub struct RpcRecentPrioritizationFeesConfig {
     pub percentile: Option<u16>,
 }
+
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct RpcLatestBlockhashConfig {
+    #[serde(flatten)]
+    pub context: RpcContextConfig,
+    #[serde(default)]
+    pub rollback: usize,
+}
diff --git a/rpc-client/src/nonblocking/rpc_client.rs b/rpc-client/src/nonblocking/rpc_client.rs
index b60680165e..d22787dccb 100644
--- a/rpc-client/src/nonblocking/rpc_client.rs
+++ b/rpc-client/src/nonblocking/rpc_client.rs
@@ -5290,6 +5290,26 @@ impl RpcClient {
         Ok(blockhash)
     }
 
+    pub async fn get_latest_blockhash_with_config(
+        &self,
+        config: RpcLatestBlockhashConfig,
+    ) -> ClientResult<(Hash, u64)> {
+        let RpcBlockhash {
+            blockhash,
+            last_valid_block_height,
+        } = self
+            .send::<Response<RpcBlockhash>>(RpcRequest::GetLatestBlockhash, json!([config]))
+            .await?
+            .value;
+        let blockhash = blockhash.parse().map_err(|_| {
+            ClientError::new_with_request(
+                RpcError::ParseError("Hash".to_string()).into(),
+                RpcRequest::GetLatestBlockhash,
+            )
+        })?;
+        Ok((blockhash, last_valid_block_height))
+    }
+
     #[allow(deprecated)]
     pub async fn get_latest_blockhash_with_commitment(
         &self,
diff --git a/rpc/src/rpc.rs b/rpc/src/rpc.rs
index d348376ce4..0ca62e070b 100644
--- a/rpc/src/rpc.rs
+++ b/rpc/src/rpc.rs
@@ -2135,8 +2135,20 @@ impl JsonRpcRequestProcessor {
         }
     }
 
-    fn get_latest_blockhash(&self, config: RpcContextConfig) -> Result<RpcResponse<RpcBlockhash>> {
-        let bank = self.get_bank_with_config(config)?;
+    fn get_latest_blockhash(
+        &self,
+        config: RpcLatestBlockhashConfig,
+    ) -> Result<RpcResponse<RpcBlockhash>> {
+        let mut bank = self.get_bank_with_config(config.context)?;
+        if config.rollback > 300 {
+            return Err(Error::invalid_params("rollback exceeds 300"));
+        }
+        for _ in 0..config.rollback {
+            bank = match bank.parent() {
+                Some(bank) => bank,
+                None => return Err(Error::invalid_params("failed to rollback block")),
+            };
+        }
         let blockhash = bank.last_blockhash();
         let last_valid_block_height = bank
             .get_blockhash_last_valid_block_height(&blockhash)
@@ -3383,7 +3395,7 @@ pub mod rpc_full {
         fn get_latest_blockhash(
             &self,
             meta: Self::Metadata,
-            config: Option<RpcContextConfig>,
+            config: Option<RpcLatestBlockhashConfig>,
         ) -> Result<RpcResponse<RpcBlockhash>>;
 
         #[rpc(meta, name = "isBlockhashValid")]
@@ -4045,7 +4057,7 @@ pub mod rpc_full {
         fn get_latest_blockhash(
             &self,
             meta: Self::Metadata,
-            config: Option<RpcContextConfig>,
+            config: Option<RpcLatestBlockhashConfig>,
         ) -> Result<RpcResponse<RpcBlockhash>> {
             debug!("get_latest_blockhash rpc request received");
             meta.get_latest_blockhash(config.unwrap_or_default())