more tests

crates/core/component/ibc/src/component/rpc/client_query.rs

@@ -26,7 +26,7 @@ use crate::component::{ClientStateReadExt, HostInterface};
 use crate::prefix::MerklePrefixExt;
 use crate::IBC_COMMITMENT_PREFIX;
 
-use super::utils::determine_snapshot_from_height_header;
+use super::utils::determine_snapshot_from_metadata;
 use super::IbcQuery;
 
 #[async_trait]
@@ -35,7 +35,7 @@ impl<HI: HostInterface + Send + Sync + 'static> ClientQuery for IbcQuery<HI> {
         &self,
         request: tonic::Request<QueryClientStateRequest>,
     ) -> std::result::Result<Response<QueryClientStateResponse>, Status> {
-        let snapshot = match determine_snapshot_from_height_header(self.storage.clone(), &request) {
+        let snapshot = match determine_snapshot_from_metadata(self.storage.clone(), request.metadata()) {
             Err(err) => return Err(tonic::Status::aborted(
                 format!("could not determine the correct snapshot to open given the `\"height\"` header of the request: {err:#}")
             )),
@@ -116,7 +116,7 @@ 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 = match determine_snapshot_from_height_header(self.storage.clone(), &request) {
+        let snapshot = match determine_snapshot_from_metadata(self.storage.clone(), request.metadata()) {
             Err(err) => return Err(tonic::Status::aborted(
                 format!("could not determine the correct snapshot to open given the `\"height\"` header of the request: {err:#}")
             )),

crates/core/component/ibc/src/component/rpc/connection_query.rs

@@ -19,7 +19,7 @@ use ibc_types::DomainType;
 use prost::Message;
 use std::str::FromStr;
 
-use crate::component::rpc::utils::determine_snapshot_from_height_header;
+use crate::component::rpc::utils::determine_snapshot_from_metadata;
 use crate::component::{ConnectionStateReadExt, HostInterface};
 use crate::prefix::MerklePrefixExt;
 use crate::IBC_COMMITMENT_PREFIX;
@@ -35,7 +35,7 @@ impl<HI: HostInterface + Send + Sync + 'static> ConnectionQuery for IbcQuery<HI>
     ) -> std::result::Result<tonic::Response<QueryConnectionResponse>, tonic::Status> {
         tracing::debug!("querying connection {:?}", request);
 
-        let snapshot = match determine_snapshot_from_height_header(self.storage.clone(), &request) {
+        let snapshot = match determine_snapshot_from_metadata(self.storage.clone(), request.metadata()) {
             Err(err) => return Err(tonic::Status::aborted(
                 format!("could not determine the correct snapshot to open given the `\"height\"` header of the request: {err:#}")
             )),

crates/core/component/ibc/src/component/rpc/consensus_query.rs

@@ -31,7 +31,7 @@ use std::str::FromStr;
 
 use crate::component::{ChannelStateReadExt, ConnectionStateReadExt, HostInterface};
 
-use super::utils::determine_snapshot_from_height_header;
+use super::utils::determine_snapshot_from_metadata;
 use super::IbcQuery;
 
 #[async_trait]
@@ -350,7 +350,7 @@ 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 = match determine_snapshot_from_height_header(self.storage.clone(), &request) {
+        let snapshot = match determine_snapshot_from_metadata(self.storage.clone(), request.metadata()) {
             Err(err) => return Err(tonic::Status::aborted(
                 format!("could not determine the correct snapshot to open given the `\"height\"` header of the request: {err:#}")
             )),
@@ -459,7 +459,7 @@ 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 = match determine_snapshot_from_height_header(self.storage.clone(), &request) {
+        let snapshot = match determine_snapshot_from_metadata(self.storage.clone(), request.metadata()) {
             Err(err) => return Err(tonic::Status::aborted(
                 format!("could not determine the correct snapshot to open given the `\"height\"` header of the request: {err:#}")
             )),
@@ -499,7 +499,7 @@ impl<HI: HostInterface + Send + Sync + 'static> ConsensusQuery for IbcQuery<HI>
         request: tonic::Request<QueryPacketAcknowledgementRequest>,
     ) -> std::result::Result<tonic::Response<QueryPacketAcknowledgementResponse>, tonic::Status>
     {
-        let snapshot = match determine_snapshot_from_height_header(self.storage.clone(), &request) {
+        let snapshot = match determine_snapshot_from_metadata(self.storage.clone(), request.metadata()) {
             Err(err) => return Err(tonic::Status::aborted(
                 format!("could not determine the correct snapshot to open given the `\"height\"` header of the request: {err:#}")
             )),

crates/core/component/ibc/src/component/rpc/utils.rs

@@ -8,32 +8,45 @@ use ibc_proto::ibc::core::client::v1::Height;
 use tracing::debug;
 use tracing::instrument;
 
+type Type = tonic::metadata::MetadataMap;
+
+/// Determines which state snapshot to open given the height header in a [`MetadataMap`].
+///
+/// Returns the latest snapshot if the height header is 0, 0-0, or absent.
 #[instrument(skip_all, level = "debug")]
-pub(in crate::component::rpc) fn determine_snapshot_from_height_header<T>(
+pub(in crate::component::rpc) fn determine_snapshot_from_metadata(
     storage: Storage,
-    request: &tonic::Request<T>,
+    metadata: &Type,
 ) -> anyhow::Result<Snapshot> {
-    let height = match request.metadata().get("height") {
+    let height = determine_height_from_metadata(metadata)
+        .context("failed to determine height from metadata")?;
+    if height.revision_height == 0 {
+        Ok(storage.latest_snapshot())
+    } else {
+        storage
+            .snapshot(height.revision_height)
+            .context("failed to create state snapshot from IBC height in height header")
+    }
+}
+
+#[instrument(skip_all, level = "debug")]
+fn determine_height_from_metadata(
+    metadata: &tonic::metadata::MetadataMap,
+) -> anyhow::Result<Height> {
+    match metadata.get("height") {
         None => {
             debug!("height header was missing; assuming a height of 0");
-            TheHeight::zero().into_inner()
+            Ok(TheHeight::zero().into_inner())
         }
         Some(entry) => entry
             .to_str()
             .context("height header was present but its entry was not ASCII")
             .and_then(parse_as_ibc_height)
-            .context("failed to height header as IBC height")?,
-    };
-    if height.revision_height == 0 {
-        Ok(storage.latest_snapshot())
-    } else {
-        storage
-            .snapshot(height.revision_height)
-            .context("failed to create state snapshot from IBC height in height header")
+            .context("failed to parse height header as IBC height"),
     }
 }
 
-/// Utility to implement
+/// Newtype wrapper around [`Height`] to implement [`FromStr`].
 #[derive(Debug)]
 struct TheHeight(Height);
 
@@ -90,6 +103,9 @@ fn parse_as_ibc_height(input: &str) -> anyhow::Result<Height> {
 #[cfg(test)]
 mod tests {
     use ibc_proto::ibc::core::client::v1::Height;
+    use tonic::metadata::MetadataMap;
+
+    use crate::component::rpc::utils::determine_height_from_metadata;
 
     use super::TheHeight;
 
@@ -121,4 +137,24 @@ mod tests {
         assert_ibc_height_is_parsed_correctly("1-0", height(1, 0));
         assert_ibc_height_is_parsed_correctly("1-1", height(1, 1));
     }
+
+    #[track_caller]
+    fn assert_ibc_height_is_determined_correctly(input: Option<&str>, expected: Height) {
+        let mut metadata = MetadataMap::new();
+        if let Some(input) = input {
+            metadata.insert("height", input.parse().unwrap());
+        }
+        let actual = determine_height_from_metadata(&metadata).unwrap();
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn determine_ibc_height_from_metadata() {
+        assert_ibc_height_is_determined_correctly(None, zero());
+        assert_ibc_height_is_determined_correctly(Some("0"), zero());
+        assert_ibc_height_is_determined_correctly(Some("0-0"), zero());
+        assert_ibc_height_is_determined_correctly(Some("0-1"), height(0, 1));
+        assert_ibc_height_is_determined_correctly(Some("1-0"), height(1, 0));
+        assert_ibc_height_is_determined_correctly(Some("1-1"), height(1, 1));
+    }
 }