ViewState RPC supports prefix search (#1144)

* Adding a prefix to view_state RPC to allow to sub-query subset of state from the account

* Remove not needed function

* Fix comments

* Remove not needed function

core/primitives/src/utils.rs

@@ -11,6 +11,8 @@ use crate::crypto::signature::PublicKey;
 use crate::hash::{hash, CryptoHash};
 use crate::types::{AccountId, ShardId};
 
+pub const ACCOUNT_DATA_SEPARATOR: &[u8; 1] = b",";
+
 pub mod col {
     pub const ACCOUNT: &[u8] = &[0];
     pub const CALLBACK: &[u8] = &[1];
@@ -22,8 +24,6 @@ pub mod col {
     pub const POSTPONED_RECEIPT: &[u8] = &[7];
 }
 
-pub const ACCOUNT_DATA_SEPARATOR: &[u8; 1] = b",";
-
 fn key_for_column_account_id(column: &[u8], account_key: &AccountId) -> Vec<u8> {
     let mut key = column.to_vec();
     key.append(&mut account_key.clone().into_bytes());
@@ -34,6 +34,13 @@ pub fn key_for_account(account_key: &AccountId) -> Vec<u8> {
     key_for_column_account_id(col::ACCOUNT, account_key)
 }
 
+pub fn key_for_data(account_id: &AccountId, data: &[u8]) -> Vec<u8> {
+    let mut bytes = key_for_account(account_id);
+    bytes.extend(ACCOUNT_DATA_SEPARATOR);
+    bytes.extend(data);
+    bytes
+}
+
 pub fn prefix_for_access_key(account_id: &AccountId) -> Vec<u8> {
     let mut key = key_for_column_account_id(col::ACCESS_KEY, account_id);
     key.extend_from_slice(col::ACCESS_KEY);
@@ -57,12 +64,6 @@ pub fn key_for_code(account_key: &AccountId) -> Vec<u8> {
     key_for_column_account_id(col::CODE, account_key)
 }
 
-pub fn key_for_data(account_id: &AccountId, key: &[u8]) -> Vec<u8> {
-    let mut prefix = prefix_for_data(account_id);
-    prefix.extend_from_slice(key);
-    prefix
-}
-
 pub fn key_for_received_data(account_id: &AccountId, data_id: &CryptoHash) -> Vec<u8> {
     let mut key = key_for_column_account_id(col::RECEIVED_DATA, account_id);
     key.append(&mut ACCOUNT_DATA_SEPARATOR.to_vec());

near/src/runtime.rs

@@ -446,9 +446,10 @@ impl node_runtime::adapter::ViewRuntimeAdapter for NightshadeRuntime {
         &self,
         state_root: MerkleHash,
         account_id: &AccountId,
+        prefix: &[u8],
     ) -> Result<ViewStateResult, Box<dyn std::error::Error>> {
         let state_update = TrieUpdate::new(self.trie.clone(), state_root);
-        self.trie_viewer.view_state(&state_update, account_id)
+        self.trie_viewer.view_state(&state_update, account_id, prefix)
     }
 }
 

runtime/runtime/src/adapter.rs

@@ -41,6 +41,7 @@ pub trait ViewRuntimeAdapter {
         &self,
         state_root: MerkleHash,
         account_id: &AccountId,
+        prefix: &[u8],
     ) -> Result<ViewStateResult, Box<dyn std::error::Error>>;
 }
 
@@ -76,7 +77,7 @@ pub fn query_client(
                 Err(err) => Ok(QueryResponse::Error(QueryError { error: err.to_string(), logs })),
             }
         }
-        "contract" => match adapter.view_state(state_root, &AccountId::from(path_parts[1])) {
+        "contract" => match adapter.view_state(state_root, &AccountId::from(path_parts[1]), data) {
             Ok(result) => Ok(QueryResponse::ViewState(result)),
             Err(err) => {
                 Ok(QueryResponse::Error(QueryError { error: err.to_string(), logs: vec![] }))

runtime/runtime/src/state_viewer.rs

@@ -73,15 +73,18 @@ impl TrieViewer {
         &self,
         state_update: &TrieUpdate,
         account_id: &AccountId,
+        prefix: &[u8],
     ) -> Result<ViewStateResult, Box<dyn std::error::Error>> {
         if !is_valid_account_id(account_id) {
             return Err(format!("Account ID '{}' is not valid", account_id).into());
         }
         let mut values = HashMap::default();
-        let prefix = prefix_for_data(account_id);
-        state_update.for_keys_with_prefix(&prefix, |key| {
+        let mut query = prefix_for_data(account_id);
+        let acc_sep_len = query.len();
+        query.extend_from_slice(prefix);
+        state_update.for_keys_with_prefix(&query, |key| {
             if let Some(value) = state_update.get(key) {
-                values.insert(key[prefix.len()..].to_vec(), value.to_vec());
+                values.insert(key[acc_sep_len..].to_vec(), value.to_vec());
             }
         });
         Ok(ViewStateResult { values })
@@ -176,7 +179,7 @@ mod tests {
     use kvdb::DBValue;
     use tempdir::TempDir;
 
-    use near_primitives::types::AccountId;
+    use near_primitives::utils::key_for_data;
     use testlib::runtime_utils::{
         alice_account, encode_int, get_runtime_and_trie, get_test_trie_viewer,
     };
@@ -188,8 +191,7 @@ mod tests {
         let (viewer, root) = get_test_trie_viewer();
 
         let mut logs = vec![];
-        let result =
-            viewer.call_function(root, 1, &alice_account(), "run_test", &vec![], &mut logs);
+        let result = viewer.call_function(root, 1, &alice_account(), "run_test", &[], &mut logs);
 
         assert_eq!(result.unwrap(), encode_int(10));
     }
@@ -199,14 +201,8 @@ mod tests {
         let (viewer, root) = get_test_trie_viewer();
 
         let mut logs = vec![];
-        let result = viewer.call_function(
-            root,
-            1,
-            &"bad!contract".to_string(),
-            "run_test",
-            &vec![],
-            &mut logs,
-        );
+        let result =
+            viewer.call_function(root, 1, &"bad!contract".to_string(), "run_test", &[], &mut logs);
 
         assert!(result.is_err());
     }
@@ -221,7 +217,7 @@ mod tests {
             1,
             &alice_account(),
             "run_test_with_storage_change",
-            &vec![],
+            &[],
             &mut logs,
         );
         // run_test tries to change storage, so it should fail
@@ -238,28 +234,29 @@ mod tests {
         assert_eq!(view_call_result.unwrap(), 3u64.to_le_bytes().to_vec());
     }
 
-    fn account_suffix(account_id: &AccountId, suffix: &[u8]) -> Vec<u8> {
-        let mut bytes = prefix_for_data(account_id);
-        bytes.append(&mut suffix.clone().to_vec());
-        bytes
-    }
-
     #[test]
     fn test_view_state() {
         let (_, trie, root) = get_runtime_and_trie();
         let mut state_update = TrieUpdate::new(trie.clone(), root);
-        state_update.set(account_suffix(&alice_account(), b"test123"), DBValue::from_slice(b"123"));
+        state_update.set(key_for_data(&alice_account(), b"test123"), DBValue::from_slice(b"123"));
         let (db_changes, new_root) = state_update.finalize().unwrap().into(trie.clone()).unwrap();
         db_changes.commit().unwrap();
 
         let state_update = TrieUpdate::new(trie, new_root);
         let ethash_provider =
             EthashProvider::new(TempDir::new("runtime_user_test_ethash").unwrap().path());
         let trie_viewer = TrieViewer::new(Arc::new(Mutex::new(ethash_provider)));
-        let result = trie_viewer.view_state(&state_update, &alice_account()).unwrap();
+        let result = trie_viewer.view_state(&state_update, &alice_account(), b"").unwrap();
         assert_eq!(
             result.values,
             [(b"test123".to_vec(), b"123".to_vec())].iter().cloned().collect()
         );
+        let result = trie_viewer.view_state(&state_update, &alice_account(), b"test321").unwrap();
+        assert_eq!(result.values, [].iter().cloned().collect());
+        let result = trie_viewer.view_state(&state_update, &alice_account(), b"test123").unwrap();
+        assert_eq!(
+            result.values,
+            [(b"test123".to_vec(), b"123".to_vec())].iter().cloned().collect()
+        )
     }
 }

test-utils/testlib/src/standard_test_cases.rs

@@ -868,15 +868,15 @@ pub fn test_stake_fail_not_enough_rent(node: impl Node) {
 pub fn test_delete_account(node: impl Node) {
     let node_user = node.user();
     // There is some data attached to the account.
-    assert!(node_user.view_state(&bob_account()).unwrap().values.len() > 0);
+    assert!(node_user.view_state(&bob_account(), b"").unwrap().values.len() > 0);
     let initial_amount = node_user.view_account(&node.account_id().unwrap()).unwrap().amount;
     let bobs_amount = node_user.view_account(&bob_account()).unwrap().amount;
     let transaction_result = node_user.delete_account(alice_account(), bob_account());
     assert_eq!(transaction_result.status, FinalTransactionStatus::Completed);
     assert_eq!(transaction_result.transactions.len(), 3);
     assert!(node_user.view_account(&bob_account()).is_err());
     // No data left.
-    assert_eq!(node_user.view_state(&bob_account()).unwrap().values.len(), 0);
+    assert_eq!(node_user.view_state(&bob_account(), b"").unwrap().values.len(), 0);
     // Receive back reward the balance of the bob's account.
     assert_eq!(
         node_user.view_account(&node.account_id().unwrap()).unwrap().amount,

test-utils/testlib/src/user/mod.rs

@@ -30,7 +30,7 @@ pub trait User {
         Ok(self.view_account(account_id)?.amount)
     }
 
-    fn view_state(&self, account_id: &AccountId) -> Result<ViewStateResult, String>;
+    fn view_state(&self, account_id: &AccountId, prefix: &[u8]) -> Result<ViewStateResult, String>;
 
     fn add_transaction(&self, signed_transaction: SignedTransaction) -> Result<(), String>;
 

test-utils/testlib/src/user/rpc_user.rs

@@ -34,18 +34,18 @@ impl RpcUser {
         System::new("actix").block_on(self.client.write().unwrap().status()).ok()
     }
 
-    pub fn query(&self, path: String, data: Vec<u8>) -> Result<QueryResponse, String> {
-        System::new("actix").block_on(self.client.write().unwrap().query(path, to_base(&data)))
+    pub fn query(&self, path: String, data: &[u8]) -> Result<QueryResponse, String> {
+        System::new("actix").block_on(self.client.write().unwrap().query(path, to_base(data)))
     }
 }
 
 impl User for RpcUser {
     fn view_account(&self, account_id: &AccountId) -> Result<AccountViewCallResult, String> {
-        self.query(format!("account/{}", account_id), vec![])?.try_into()
+        self.query(format!("account/{}", account_id), &[])?.try_into()
     }
 
-    fn view_state(&self, account_id: &AccountId) -> Result<ViewStateResult, String> {
-        self.query(format!("contract/{}", account_id), vec![])?.try_into()
+    fn view_state(&self, account_id: &AccountId, prefix: &[u8]) -> Result<ViewStateResult, String> {
+        self.query(format!("contract/{}", account_id), prefix)?.try_into()
     }
 
     fn add_transaction(&self, transaction: SignedTransaction) -> Result<(), String> {
@@ -104,8 +104,7 @@ impl User for RpcUser {
         account_id: &AccountId,
         public_key: &PublicKey,
     ) -> Result<Option<AccessKey>, String> {
-        self.query(format!("access_key/{}/{}", account_id, public_key.to_base()), vec![])?
-            .try_into()
+        self.query(format!("access_key/{}/{}", account_id, public_key.to_base()), &[])?.try_into()
     }
 
     fn signer(&self) -> Arc<dyn EDSigner> {

test-utils/testlib/src/user/runtime_user.rs

@@ -163,9 +163,11 @@ impl User for RuntimeUser {
         self.trie_viewer.view_account(&state_update, account_id).map_err(|err| err.to_string())
     }
 
-    fn view_state(&self, account_id: &AccountId) -> Result<ViewStateResult, String> {
+    fn view_state(&self, account_id: &AccountId, prefix: &[u8]) -> Result<ViewStateResult, String> {
         let state_update = self.client.read().expect(POISONED_LOCK_ERR).get_state_update();
-        self.trie_viewer.view_state(&state_update, account_id).map_err(|err| err.to_string())
+        self.trie_viewer
+            .view_state(&state_update, account_id, prefix)
+            .map_err(|err| err.to_string())
     }
 
     fn add_transaction(&self, transaction: SignedTransaction) -> Result<(), String> {