From d8de56920a7a7de169a88fd1b5fc505ec9d0685a Mon Sep 17 00:00:00 2001
From: elizabeth <elizabethjbinks@gmail.com>
Date: Fri, 18 Oct 2024 15:20:48 -0400
Subject: [PATCH] get height from grpc metadata for relevant ibc queries

---
 .../core/component/ibc/src/component/rpc.rs   |  1 +
 .../ibc/src/component/rpc/client_query.rs     | 41 ++++++++++++-
 .../ibc/src/component/rpc/connection_query.rs | 21 ++++++-
 .../ibc/src/component/rpc/consensus_query.rs  | 59 ++++++++++++++++++-
 .../component/ibc/src/component/rpc/utils.rs  | 27 +++++++++
 5 files changed, 143 insertions(+), 6 deletions(-)
 create mode 100644 crates/core/component/ibc/src/component/rpc/utils.rs

diff --git a/crates/core/component/ibc/src/component/rpc.rs b/crates/core/component/ibc/src/component/rpc.rs
index 4d125b1350..4b54f52105 100644
--- a/crates/core/component/ibc/src/component/rpc.rs
+++ b/crates/core/component/ibc/src/component/rpc.rs
@@ -6,6 +6,7 @@ use super::HostInterface;
 mod client_query;
 mod connection_query;
 mod consensus_query;
+mod utils;
 
 use std::marker::PhantomData;
 
diff --git a/crates/core/component/ibc/src/component/rpc/client_query.rs b/crates/core/component/ibc/src/component/rpc/client_query.rs
index 9b52f35097..82eaa1bdf0 100644
--- a/crates/core/component/ibc/src/component/rpc/client_query.rs
+++ b/crates/core/component/ibc/src/component/rpc/client_query.rs
@@ -13,6 +13,7 @@ use ibc_proto::ibc::core::client::v1::{
 };
 use prost::Message;
 
+use crate::component::rpc::utils::height_from_str;
 use ibc_types::core::client::ClientId;
 use ibc_types::core::client::Height;
 use ibc_types::path::ClientConsensusStatePath;
@@ -34,7 +35,25 @@ impl<HI: HostInterface + Send + Sync + 'static> ClientQuery for IbcQuery<HI> {
         &self,
         request: tonic::Request<QueryClientStateRequest>,
     ) -> std::result::Result<Response<QueryClientStateResponse>, Status> {
-        let snapshot = self.storage.latest_snapshot();
+        let Some(height_val) = request.metadata().get("height") else {
+            return Err(tonic::Status::aborted("missing height"));
+        };
+
+        let height_str: &str = height_val
+            .to_str()
+            .map_err(|e| tonic::Status::aborted(format!("invalid height: {e}")))?;
+
+        let snapshot = if height_str == "0" {
+            self.storage.latest_snapshot()
+        } else {
+            let height = height_from_str(height_str)
+                .map_err(|e| tonic::Status::aborted(format!("couldn't get snapshot: {e}")))?;
+
+            self.storage
+                .snapshot(height.revision_height as u64)
+                .ok_or(tonic::Status::aborted(format!("invalid height")))?
+        };
+
         let client_id = ClientId::from_str(&request.get_ref().client_id)
             .map_err(|e| tonic::Status::invalid_argument(format!("invalid client id: {e}")))?;
         let height = Height {
@@ -109,7 +128,25 @@ impl<HI: HostInterface + Send + Sync + 'static> ClientQuery for IbcQuery<HI> {
         &self,
         request: tonic::Request<QueryConsensusStateRequest>,
     ) -> std::result::Result<tonic::Response<QueryConsensusStateResponse>, tonic::Status> {
-        let snapshot = self.storage.latest_snapshot();
+        let Some(height_val) = request.metadata().get("height") else {
+            return Err(tonic::Status::aborted("missing height"));
+        };
+
+        let height_str: &str = height_val
+            .to_str()
+            .map_err(|e| tonic::Status::aborted(format!("invalid height: {e}")))?;
+
+        let snapshot = if height_str == "0" {
+            self.storage.latest_snapshot()
+        } else {
+            let height = height_from_str(height_str)
+                .map_err(|e| tonic::Status::aborted(format!("couldn't get snapshot: {e}")))?;
+
+            self.storage
+                .snapshot(height.revision_height as u64)
+                .ok_or(tonic::Status::aborted(format!("invalid height")))?
+        };
+
         let client_id = ClientId::from_str(&request.get_ref().client_id)
             .map_err(|e| tonic::Status::invalid_argument(format!("invalid client id: {e}")))?;
         let height = if request.get_ref().latest_height {
diff --git a/crates/core/component/ibc/src/component/rpc/connection_query.rs b/crates/core/component/ibc/src/component/rpc/connection_query.rs
index ae98d9f381..4cf9c8009f 100644
--- a/crates/core/component/ibc/src/component/rpc/connection_query.rs
+++ b/crates/core/component/ibc/src/component/rpc/connection_query.rs
@@ -19,6 +19,7 @@ use ibc_types::DomainType;
 use prost::Message;
 use std::str::FromStr;
 
+use crate::component::rpc::utils::height_from_str;
 use crate::component::{ConnectionStateReadExt, HostInterface};
 use crate::prefix::MerklePrefixExt;
 use crate::IBC_COMMITMENT_PREFIX;
@@ -33,7 +34,25 @@ impl<HI: HostInterface + Send + Sync + 'static> ConnectionQuery for IbcQuery<HI>
         request: tonic::Request<QueryConnectionRequest>,
     ) -> std::result::Result<tonic::Response<QueryConnectionResponse>, tonic::Status> {
         tracing::debug!("querying connection {:?}", request);
-        let snapshot = self.storage.latest_snapshot();
+        let Some(height_val) = request.metadata().get("height") else {
+            return Err(tonic::Status::aborted("missing height"));
+        };
+
+        let height_str: &str = height_val
+            .to_str()
+            .map_err(|e| tonic::Status::aborted(format!("invalid height: {e}")))?;
+
+        let snapshot = if height_str == "0" {
+            self.storage.latest_snapshot()
+        } else {
+            let height = height_from_str(height_str)
+                .map_err(|e| tonic::Status::aborted(format!("couldn't get snapshot: {e}")))?;
+
+            self.storage
+                .snapshot(height.revision_height as u64)
+                .ok_or(tonic::Status::aborted(format!("invalid height")))?
+        };
+
         let connection_id = &ConnectionId::from_str(&request.get_ref().connection_id)
             .map_err(|e| tonic::Status::aborted(format!("invalid connection id: {e}")))?;
 
diff --git a/crates/core/component/ibc/src/component/rpc/consensus_query.rs b/crates/core/component/ibc/src/component/rpc/consensus_query.rs
index 39221e2384..74de835849 100644
--- a/crates/core/component/ibc/src/component/rpc/consensus_query.rs
+++ b/crates/core/component/ibc/src/component/rpc/consensus_query.rs
@@ -29,6 +29,7 @@ use prost::Message;
 
 use std::str::FromStr;
 
+use crate::component::rpc::utils::height_from_str;
 use crate::component::{ChannelStateReadExt, ConnectionStateReadExt, HostInterface};
 
 use super::IbcQuery;
@@ -349,7 +350,24 @@ impl<HI: HostInterface + Send + Sync + 'static> ConsensusQuery for IbcQuery<HI>
         &self,
         request: tonic::Request<QueryPacketCommitmentRequest>,
     ) -> std::result::Result<tonic::Response<QueryPacketCommitmentResponse>, tonic::Status> {
-        let snapshot = self.storage.latest_snapshot();
+        let Some(height_val) = request.metadata().get("height") else {
+            return Err(tonic::Status::aborted("missing height"));
+        };
+
+        let height_str: &str = height_val
+            .to_str()
+            .map_err(|e| tonic::Status::aborted(format!("invalid height: {e}")))?;
+
+        let snapshot = if height_str == "0" {
+            self.storage.latest_snapshot()
+        } else {
+            let height = height_from_str(height_str)
+                .map_err(|e| tonic::Status::aborted(format!("couldn't get snapshot: {e}")))?;
+
+            self.storage
+                .snapshot(height.revision_height as u64)
+                .ok_or(tonic::Status::aborted(format!("invalid height")))?
+        };
 
         let port_id = PortId::from_str(&request.get_ref().port_id)
             .map_err(|e| tonic::Status::aborted(format!("invalid port id: {e}")))?;
@@ -453,7 +471,24 @@ impl<HI: HostInterface + Send + Sync + 'static> ConsensusQuery for IbcQuery<HI>
         &self,
         request: tonic::Request<QueryPacketReceiptRequest>,
     ) -> std::result::Result<tonic::Response<QueryPacketReceiptResponse>, tonic::Status> {
-        let snapshot = self.storage.latest_snapshot();
+        let Some(height_val) = request.metadata().get("height") else {
+            return Err(tonic::Status::aborted("missing height"));
+        };
+
+        let height_str: &str = height_val
+            .to_str()
+            .map_err(|e| tonic::Status::aborted(format!("invalid height: {e}")))?;
+
+        let snapshot = if height_str == "0" {
+            self.storage.latest_snapshot()
+        } else {
+            let height = height_from_str(height_str)
+                .map_err(|e| tonic::Status::aborted(format!("couldn't get snapshot: {e}")))?;
+
+            self.storage
+                .snapshot(height.revision_height as u64)
+                .ok_or(tonic::Status::aborted(format!("invalid height")))?
+        };
 
         let port_id = PortId::from_str(&request.get_ref().port_id)
             .map_err(|e| tonic::Status::aborted(format!("invalid port id: {e}")))?;
@@ -488,7 +523,25 @@ impl<HI: HostInterface + Send + Sync + 'static> ConsensusQuery for IbcQuery<HI>
         request: tonic::Request<QueryPacketAcknowledgementRequest>,
     ) -> std::result::Result<tonic::Response<QueryPacketAcknowledgementResponse>, tonic::Status>
     {
-        let snapshot = self.storage.latest_snapshot();
+        let Some(height_val) = request.metadata().get("height") else {
+            return Err(tonic::Status::aborted("missing height"));
+        };
+
+        let height_str: &str = height_val
+            .to_str()
+            .map_err(|e| tonic::Status::aborted(format!("invalid height: {e}")))?;
+
+        let snapshot = if height_str == "0" {
+            self.storage.latest_snapshot()
+        } else {
+            let height = height_from_str(height_str)
+                .map_err(|e| tonic::Status::aborted(format!("couldn't get snapshot: {e}")))?;
+
+            self.storage
+                .snapshot(height.revision_height as u64)
+                .ok_or(tonic::Status::aborted(format!("invalid height")))?
+        };
+
         let channel_id = ChannelId::from_str(request.get_ref().channel_id.as_str())
             .map_err(|e| tonic::Status::aborted(format!("invalid channel id: {e}")))?;
         let port_id = PortId::from_str(request.get_ref().port_id.as_str())
diff --git a/crates/core/component/ibc/src/component/rpc/utils.rs b/crates/core/component/ibc/src/component/rpc/utils.rs
new file mode 100644
index 0000000000..108a14248a
--- /dev/null
+++ b/crates/core/component/ibc/src/component/rpc/utils.rs
@@ -0,0 +1,27 @@
+use anyhow::anyhow;
+use ibc_proto::ibc::core::client::v1::Height;
+
+pub(crate) fn height_from_str(value: &str) -> anyhow::Result<Height> {
+    let split: Vec<&str> = value.split('-').collect();
+
+    if split.len() != 2 {
+        return Err(anyhow!("invalid height string"));
+    }
+
+    let revision_number = split[0]
+        .parse::<u64>()
+        .map_err(|e| anyhow!("failed to parse revision number"))?;
+
+    let revision_height = split[1]
+        .parse::<u64>()
+        .map_err(|e| anyhow!("failed to parse revision height"))?;
+
+    if revision_number == 0 && revision_height == 0 {
+        return Err(anyhow!("height is zero"));
+    }
+
+    Ok(Height {
+        revision_number,
+        revision_height,
+    })
+}