[Storage] Verify proof. (aptos-labs#3022)

storage/storage-interface/src/async_proof_fetcher.rs

@@ -4,7 +4,7 @@
 use crate::{proof_fetcher::ProofFetcher, DbReader};
 
 use crate::metrics::TIMER;
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use aptos_crypto::{hash::CryptoHash, HashValue};
 use aptos_types::{
     proof::SparseMerkleProofExt,
@@ -77,7 +77,13 @@ impl AsyncProofFetcher {
     }
 
     // Schedules proof reading work in a background running thread pool.
-    fn schedule_proof_read(&self, state_key: StateKey, version: Version) {
+    fn schedule_proof_read(
+        &self,
+        state_key: StateKey,
+        version: Version,
+        root_hash: Option<HashValue>,
+        value_hash: Option<HashValue>,
+    ) {
         let _timer = TIMER
             .with_label_values(&["schedule_async_proof_read"])
             .start_timer();
@@ -88,6 +94,19 @@ impl AsyncProofFetcher {
             let proof = reader
                 .get_state_proof_by_version_ext(&state_key, version)
                 .expect("Proof reading should succeed.");
+            if let Some(root_hash) = root_hash {
+                proof
+                    .verify_by_hash(root_hash, state_key.hash(), value_hash)
+                    .map_err(|err| {
+                        anyhow!(
+                            "Proof is invalid for key {:?} with state root hash {:?}: {}.",
+                            state_key,
+                            root_hash,
+                            err
+                        )
+                    })
+                    .expect("Failed to verify proof.");
+            }
             data_sender
                 .send(Proof {
                     state_key_hash: state_key.hash(),
@@ -103,12 +122,18 @@ impl ProofFetcher for AsyncProofFetcher {
         &self,
         state_key: &StateKey,
         version: Version,
+        root_hash: Option<HashValue>,
     ) -> Result<(Option<StateValue>, Option<SparseMerkleProofExt>)> {
         let _timer = TIMER
             .with_label_values(&["async_proof_fetcher_fetch"])
             .start_timer();
-        self.schedule_proof_read(state_key.clone(), version);
         let value = self.reader.get_state_value_by_version(state_key, version)?;
+        self.schedule_proof_read(
+            state_key.clone(),
+            version,
+            root_hash,
+            value.as_ref().map(|v| v.hash()),
+        );
         Ok((value, None))
     }
 
@@ -131,7 +156,7 @@ mod tests {
             let state_key = StateKey::Raw(format!("test_key_{}", i).into_bytes());
             expected_key_hashes.push(state_key.hash());
             let result = fetcher
-                .fetch_state_value_and_proof(&state_key, 0)
+                .fetch_state_value_and_proof(&state_key, 0, None)
                 .expect("Should not fail.");
             let expected_value = StateValue::from(match state_key {
                 StateKey::Raw(key) => key,

storage/storage-interface/src/cached_state_view.rs

@@ -150,9 +150,11 @@ impl CachedStateView {
             StateStoreStatus::ExistsInDB | StateStoreStatus::Unknown => {
                 match self.snapshot {
                     Some((version, root_hash)) => {
-                        let (value, proof) = self
-                            .proof_fetcher
-                            .fetch_state_value_and_proof(state_key, version)?;
+                        let (value, proof) = self.proof_fetcher.fetch_state_value_and_proof(
+                            state_key,
+                            version,
+                            Some(root_hash),
+                        )?;
                         // TODO: proof verification can be opted out, for performance
                         if let Some(proof) = proof {
                             proof

storage/storage-interface/src/proof_fetcher.rs

@@ -11,11 +11,13 @@ use std::collections::HashMap;
 
 /// Defines the trait for fetching proof from the DB
 pub trait ProofFetcher: Sync + Send {
-    /// API to fetch the state value along with proof
+    /// Fetches the state value along with its proof. If root_hash is provided, will also verify
+    /// its proof.
     fn fetch_state_value_and_proof(
         &self,
         state_key: &StateKey,
         version: Version,
+        root_hash: Option<HashValue>,
     ) -> anyhow::Result<(Option<StateValue>, Option<SparseMerkleProofExt>)>;
 
     /// API to return all the proofs fetched by the proof fetcher so far.

storage/storage-interface/src/sync_proof_fetcher.rs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{proof_fetcher::ProofFetcher, DbReader};
+use anyhow::format_err;
 use aptos_crypto::{hash::CryptoHash, HashValue};
 use aptos_types::{
     proof::SparseMerkleProofExt,
@@ -32,10 +33,23 @@ impl ProofFetcher for SyncProofFetcher {
         &self,
         state_key: &StateKey,
         version: Version,
+        root_hash: Option<HashValue>,
     ) -> anyhow::Result<(Option<StateValue>, Option<SparseMerkleProofExt>)> {
         let (state_value, proof) = self
             .reader
             .get_state_value_with_proof_by_version_ext(state_key, version)?;
+        if let Some(root_hash) = root_hash {
+            proof
+                .verify(root_hash, state_key.hash(), state_value.as_ref())
+                .map_err(|err| {
+                    format_err!(
+                        "Proof is invalid for key {:?} with state root hash {:?}: {}.",
+                        state_key,
+                        root_hash,
+                        err
+                    )
+                })?;
+        }
         // multiple threads may enter this code, and another thread might add
         // an address before this one. Thus the insertion might return a None here.
         self.state_proof_cache

types/src/proof/definition.rs

@@ -20,8 +20,7 @@ use aptos_crypto::{
     HashValue,
 };
 use serde::{Deserialize, Serialize};
-use std::any::type_name;
-use std::marker::PhantomData;
+use std::{any::type_name, marker::PhantomData};
 
 #[cfg(any(test, feature = "fuzzing"))]
 use aptos_crypto::hash::TestOnlyHasher;
@@ -212,6 +211,19 @@ impl SparseMerkleProofExt {
     ) -> Result<()> {
         SparseMerkleProof::from(self.clone()).verify(expected_root_hash, element_key, element_value)
     }
+
+    pub fn verify_by_hash(
+        &self,
+        expected_root_hash: HashValue,
+        element_key: HashValue,
+        element_hash: Option<HashValue>,
+    ) -> Result<()> {
+        SparseMerkleProof::from(self.clone()).verify_by_hash(
+            expected_root_hash,
+            element_key,
+            element_hash,
+        )
+    }
 }
 
 impl From<SparseMerkleProofExt> for SparseMerkleProof {