From 1b6fe8dbbfc5ab6f4cd2f2a3eebfb65cd5a5d681 Mon Sep 17 00:00:00 2001
From: Henry de Valence <hdevalence@penumbralabs.xyz>
Date: Fri, 22 Sep 2023 16:12:21 -0700
Subject: [PATCH] add a penumbra_use_prehash_key_before_comparison to config

---
 crates/relayer-cli/src/chain_registry.rs       |  1 +
 crates/relayer/src/chain/penumbra/endpoint.rs  | 12 +++++++-----
 crates/relayer/src/chain/penumbra/proofspec.rs | 16 +++++++++++++---
 crates/relayer/src/config.rs                   |  4 ++++
 tools/test-framework/src/types/single/node.rs  |  1 +
 5 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/crates/relayer-cli/src/chain_registry.rs b/crates/relayer-cli/src/chain_registry.rs
index 305a8d615a..4912ed2f39 100644
--- a/crates/relayer-cli/src/chain_registry.rs
+++ b/crates/relayer-cli/src/chain_registry.rs
@@ -129,6 +129,7 @@ where
             url: websocket_address,
             batch_delay: default::batch_delay(),
         },
+        penumbra_use_prehash_key_before_comparison: false,
         rpc_timeout: default::rpc_timeout(),
         trusted_node: default::trusted_node(),
         genesis_restart: None,
diff --git a/crates/relayer/src/chain/penumbra/endpoint.rs b/crates/relayer/src/chain/penumbra/endpoint.rs
index d7b233292c..b7c5d6f3e9 100644
--- a/crates/relayer/src/chain/penumbra/endpoint.rs
+++ b/crates/relayer/src/chain/penumbra/endpoint.rs
@@ -1313,11 +1313,13 @@ impl ChainEndpoint for PenumbraChain {
             .trusting_period
             .unwrap_or_else(|| trusting_period_default);
 
-        let proof_specs = self
-            .config
-            .proof_specs
-            .clone()
-            .unwrap_or(crate::chain::penumbra::proofspec::penumbra_proof_spec());
+        let proof_specs = self.config.proof_specs.clone().unwrap_or_else(|| {
+            if self.config.penumbra_use_prehash_key_before_comparison {
+                crate::chain::penumbra::proofspec::penumbra_proof_spec_with_prehash()
+            } else {
+                crate::chain::penumbra::proofspec::penumbra_proof_spec_no_prehash()
+            }
+        });
 
         // Build the client state.
         TmClientState::new(
diff --git a/crates/relayer/src/chain/penumbra/proofspec.rs b/crates/relayer/src/chain/penumbra/proofspec.rs
index b2d56f4488..b32c4c3017 100644
--- a/crates/relayer/src/chain/penumbra/proofspec.rs
+++ b/crates/relayer/src/chain/penumbra/proofspec.rs
@@ -2,7 +2,7 @@ use ibc_relayer_types::core::ics23_commitment::specs::ProofSpecs;
 
 const SPARSE_MERKLE_PLACEHOLDER_HASH: [u8; 32] = *b"SPARSE_MERKLE_PLACEHOLDER_HASH__";
 
-fn jmt_spec() -> ics23::ProofSpec {
+fn jmt_spec_no_prehash() -> ics23::ProofSpec {
     ics23::ProofSpec {
         leaf_spec: Some(ics23::LeafOp {
             hash: ics23::HashOp::Sha256.into(),
@@ -25,6 +25,12 @@ fn jmt_spec() -> ics23::ProofSpec {
     }
 }
 
+fn jmt_spec_with_prehash() -> ics23::ProofSpec {
+    let mut spec = jmt_spec_no_prehash();
+    spec.prehash_key_before_comparison = true;
+    spec
+}
+
 /// this is a proof spec for computing Penumbra's AppHash, which is defined as
 /// SHA256("PenumbraAppHash" || jmt.root()). In ICS/IBC terms, this applies a single global prefix
 /// to Penumbra's state. Having a stable merkle prefix is currently required for our IBC
@@ -55,6 +61,10 @@ fn apphash_spec() -> ibc_proto::cosmos::ics23::v1::ProofSpec {
 }
 
 // TODO: this should re-export the proof specs from the Penumbra crate
-pub fn penumbra_proof_spec() -> ProofSpecs {
-    vec![jmt_spec(), apphash_spec()].into()
+pub fn penumbra_proof_spec_no_prehash() -> ProofSpecs {
+    vec![jmt_spec_no_prehash(), apphash_spec()].into()
+}
+
+pub fn penumbra_proof_spec_with_prehash() -> ProofSpecs {
+    vec![jmt_spec_with_prehash(), apphash_spec()].into()
 }
diff --git a/crates/relayer/src/config.rs b/crates/relayer/src/config.rs
index 92fe601cbf..1101c66437 100644
--- a/crates/relayer/src/config.rs
+++ b/crates/relayer/src/config.rs
@@ -579,6 +579,10 @@ pub struct ChainConfig {
     #[serde(default = "default::trusted_node")]
     pub trusted_node: bool,
 
+    /// Whether to use prehash_key_before_comparison in Penumbra proof specs.
+    #[serde(default)]
+    pub penumbra_use_prehash_key_before_comparison: bool,
+
     pub account_prefix: String,
     pub key_name: String,
     #[serde(default)]
diff --git a/tools/test-framework/src/types/single/node.rs b/tools/test-framework/src/types/single/node.rs
index 5aeced1f4b..1c1bac94a3 100644
--- a/tools/test-framework/src/types/single/node.rs
+++ b/tools/test-framework/src/types/single/node.rs
@@ -135,6 +135,7 @@ impl FullNode {
             .to_string();
 
         Ok(config::ChainConfig {
+            penumbra_use_prehash_key_before_comparison: false,
             id: self.chain_driver.chain_id.clone(),
             r#type: ChainType::CosmosSdk,
             rpc_addr: Url::from_str(&self.chain_driver.rpc_address())?,