ViewState RPC supports prefix search

core/primitives/src/serialize.rs

@@ -56,6 +56,13 @@ pub fn from_base_buf(s: &str, buffer: &mut Vec<u8>) -> Result<(), Box<dyn std::e
     }
 }
 
+pub fn to_str_or_base(v: &[u8]) -> String {
+    match std::str::from_utf8(v) {
+        Ok(res) => res.to_string(),
+        _ => to_base(v),
+    }
+}
+
 pub trait BaseEncode {
     fn to_base(&self) -> String;
 }

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];
@@ -28,6 +30,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.append(&mut ACCOUNT_DATA_SEPARATOR.to_vec());
+    bytes.append(&mut data.clone().to_vec());
+    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);

near/src/runtime.rs

@@ -402,9 +402,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

@@ -8,15 +8,14 @@ use near_primitives::crypto::signature::PublicKey;
 use near_primitives::hash::CryptoHash;
 use near_primitives::rpc::{AccountViewCallResult, ViewStateResult};
 use near_primitives::types::AccountId;
-use near_primitives::utils::{is_valid_account_id, key_for_account};
+use near_primitives::utils::{is_valid_account_id, key_for_account, ACCOUNT_DATA_SEPARATOR};
 use near_store::{get_access_key, get_account, TrieUpdate};
 use wasm::executor;
 use wasm::types::{ReturnData, RuntimeContext};
 
 use crate::ethereum::EthashProvider;
 use crate::Runtime;
 
-use super::ext::ACCOUNT_DATA_SEPARATOR;
 use super::RuntimeExt;
 
 pub struct TrieViewer {
@@ -75,16 +74,19 @@ 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 mut prefix = key_for_account(account_id);
-        prefix.extend_from_slice(ACCOUNT_DATA_SEPARATOR);
-        state_update.for_keys_with_prefix(&prefix, |key| {
+        let mut query = key_for_account(account_id);
+        query.extend_from_slice(ACCOUNT_DATA_SEPARATOR);
+        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 })
@@ -185,7 +187,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,
     };
@@ -247,29 +249,29 @@ mod tests {
         assert_eq!(view_call_result.unwrap(), encode_int(3).to_vec());
     }
 
-    fn account_suffix(account_id: &AccountId, suffix: &[u8]) -> Vec<u8> {
-        let mut bytes = key_for_account(account_id);
-        bytes.append(&mut ACCOUNT_DATA_SEPARATOR.to_vec());
-        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

@@ -428,7 +428,7 @@ pub fn test_callback(node: RuntimeNode) {
     let transaction_result = node_user.get_transaction_result(&hash);
     assert_eq!(transaction_result.status, TransactionStatus::Completed);
     let callback: Option<Callback> = node_user
-        .view_state(account_id)
+        .view_state(account_id, b"")
         .unwrap()
         .values
         .get(&key_for_callback(&callback_id))
@@ -477,7 +477,7 @@ pub fn test_callback_failure(node: RuntimeNode) {
     let transaction_result = node_user.get_transaction_result(&hash);
     assert_eq!(transaction_result.status, TransactionStatus::Failed);
     let callback: Option<Callback> = node_user
-        .view_state(account_id)
+        .view_state(account_id, b"")
         .unwrap()
         .values
         .get(&key_for_callback(&callback_id))

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

@@ -25,7 +25,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, 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), &vec![])?.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,7 +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![])?
+        self.query(format!("access_key/{}/{}", account_id, public_key.to_base()), &vec![])?
             .try_into()
     }
 }

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

@@ -146,9 +146,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> {