[VM2] Add KeySpace for associated keys

executor/wasm/tests/integration.rs

@@ -304,6 +304,85 @@ fn traits() {
     );
 }
 
+#[test]
+fn assoc_keys() {
+    let chainspec_config = ChainspecConfig::from_chainspec_path(&*CHAINSPEC_SYMLINK)
+        .expect("must get chainspec config");
+
+    let mut executor = make_executor(&chainspec_config);
+    let (global_state, state_root_hash, _tempdir) = make_global_state_with_genesis();
+
+    let execute_request = base_execute_builder(&chainspec_config)
+        .with_target(ExecutionKind::SessionBytes(read_wasm(
+            "vm2_assoc_keys.wasm",
+        )))
+        .with_serialized_input(3u8)
+        .expect("expected serialized input to be correct")
+        .with_shared_address_generator(make_address_generator())
+        .build()
+        .expect("should build");
+
+    let add_assoc_keys_result = expect_successful_execution(
+        &mut executor,
+        &global_state,
+        state_root_hash,
+        execute_request,
+    );
+
+    let state_root_hash = global_state
+        .commit_effects(state_root_hash, add_assoc_keys_result.effects().clone())
+        .expect("Should commit");
+
+    let query_request = QueryRequest::new(
+        state_root_hash,
+        Key::AddressableEntity(EntityAddr::Account(DEFAULT_ACCOUNT_HASH.value())),
+        vec![],
+    );
+
+    let query_result = global_state.query(query_request);
+    if let QueryResult::Success { value, .. } = query_result {
+        let entity = value.as_addressable_entity().expect("must get entity");
+        assert!(entity.associated_keys().len() == 2)
+    } else {
+        panic!("Unexpected query result")
+    }
+
+    let execute_request = base_execute_builder(&chainspec_config)
+        .with_target(ExecutionKind::SessionBytes(read_wasm(
+            "vm2_assoc_keys.wasm",
+        )))
+        .with_serialized_input(0u8)
+        .expect("expected serialized input to be correct")
+        .with_shared_address_generator(make_address_generator())
+        .build()
+        .expect("should build");
+
+    let remove_assoc_key = expect_successful_execution(
+        &mut executor,
+        &global_state,
+        state_root_hash,
+        execute_request,
+    );
+
+    let state_root_hash = global_state
+        .commit_effects(state_root_hash, remove_assoc_key.effects().clone())
+        .expect("Should commit");
+
+    let query_request = QueryRequest::new(
+        state_root_hash,
+        Key::AddressableEntity(EntityAddr::Account(DEFAULT_ACCOUNT_HASH.value())),
+        vec![],
+    );
+
+    let query_result = global_state.query(query_request);
+    if let QueryResult::Success { value, .. } = query_result {
+        let entity = value.as_addressable_entity().expect("must get entity");
+        assert!(entity.associated_keys().len() == 1)
+    } else {
+        panic!("Unexpected query result")
+    }
+}
+
 #[test]
 fn upgradable() {
     let chainspec_config = ChainspecConfig::from_chainspec_path(&*CHAINSPEC_SYMLINK)

executor/wasm_common/src/keyspace.rs

@@ -15,6 +15,8 @@ pub enum KeyspaceTag {
     NamedKey = 2,
     /// Used for a payment info based storage which usually involves payment information.
     PaymentInfo = 3,
+    /// Used for updating the associated keys for a given entity.
+    AssociatedKeys = 5,
 }
 
 #[repr(u64)]
@@ -31,6 +33,8 @@ pub enum Keyspace<'a> {
     NamedKey(&'a str),
     /// Entry point payment info.
     PaymentInfo(&'a str),
+    /// Used for updating the associated keys for a given entity
+    AssociatedKeys(&'a [u8]),
 }
 
 impl Keyspace<'_> {
@@ -41,6 +45,7 @@ impl Keyspace<'_> {
             Keyspace::Context(_) => KeyspaceTag::Context,
             Keyspace::NamedKey(_) => KeyspaceTag::NamedKey,
             Keyspace::PaymentInfo(_) => KeyspaceTag::PaymentInfo,
+            Keyspace::AssociatedKeys(_) => KeyspaceTag::AssociatedKeys,
         }
     }
 

executor/wasm_host/src/host.rs

@@ -48,6 +48,7 @@ use blake2::{
     digest::{Update, VariableOutput},
     Blake2bVar,
 };
+use casper_types::addressable_entity::Weight;
 use keccak_asm::Digest as KeccakDigest;
 use sha2::Sha256;
 
@@ -188,6 +189,7 @@ pub fn casper_write<S: GlobalStateReader, E: Executor>(
 
             Keyspace::PaymentInfo(key_name)
         }
+        KeyspaceTag::AssociatedKeys => Keyspace::AssociatedKeys(&key_payload_bytes),
     };
 
     let global_state_key = match keyspace_to_global_state_key(caller.context(), keyspace) {
@@ -268,6 +270,41 @@ pub fn casper_write<S: GlobalStateReader, E: Executor>(
             let entry_point_value = EntryPointValue::V1CasperVm(entry_point);
             StoredValue::EntryPoint(entry_point_value)
         }
+        Keyspace::AssociatedKeys(bytes) => {
+            let account_hash = match AccountHash::from_bytes(bytes) {
+                Ok((account_hash, remainder)) => {
+                    if !remainder.is_empty() {
+                        return Ok(HOST_ERROR_INVALID_DATA);
+                    }
+                    account_hash
+                }
+                Err(_) => return Ok(HOST_ERROR_INVALID_INPUT),
+            };
+            let weight = match Weight::from_bytes(&value) {
+                Ok((weight, remainder)) => {
+                    if !remainder.is_empty() {
+                        return Ok(HOST_ERROR_INVALID_DATA);
+                    }
+                    weight
+                }
+                Err(_) => return Ok(HOST_ERROR_INVALID_DATA),
+            };
+
+            let mut entity = match caller.context_mut().tracking_copy.read(&global_state_key) {
+                Ok(Some(StoredValue::AddressableEntity(entity))) => entity,
+                Ok(_) | Err(_) => return Ok(HOST_ERROR_NOT_FOUND),
+            };
+
+            if entity.associated_keys().contains_key(&account_hash) {
+                if entity.update_associated_key(account_hash, weight).is_err() {
+                    return Ok(HOST_ERROR_INVALID_INPUT);
+                }
+            } else if entity.add_associated_key(account_hash, weight).is_err() {
+                return Ok(HOST_ERROR_INVALID_INPUT);
+            }
+
+            StoredValue::AddressableEntity(entity)
+        }
     };
 
     metered_write(&mut caller, global_state_key, stored_value)?;
@@ -340,6 +377,7 @@ pub fn casper_remove<S: GlobalStateReader, E: Executor>(
 
             Keyspace::PaymentInfo(key_name)
         }
+        KeyspaceTag::AssociatedKeys => Keyspace::AssociatedKeys(&key_payload_bytes),
     };
 
     let global_state_key = match keyspace_to_global_state_key(caller.context(), keyspace) {
@@ -352,9 +390,33 @@ pub fn casper_remove<S: GlobalStateReader, E: Executor>(
 
     let global_state_read_result = caller.context_mut().tracking_copy.read(&global_state_key);
     match global_state_read_result {
+        // Produce a prune transform for the named key
+        Ok(Some(StoredValue::AddressableEntity(mut entity))) => {
+            if let Keyspace::AssociatedKeys(account_hash_bytes) = keyspace {
+                let account_hash = match AccountHash::from_bytes(account_hash_bytes) {
+                    Ok((account_hash, remainder)) => {
+                        if !remainder.is_empty() {
+                            return Ok(HOST_ERROR_INVALID_INPUT);
+                        }
+                        account_hash
+                    }
+                    Err(_) => return Ok(HOST_ERROR_INVALID_DATA),
+                };
+
+                if entity.remove_associated_key(account_hash).is_err() {
+                    return Ok(HOST_ERROR_INVALID_INPUT);
+                }
+                caller
+                    .context_mut()
+                    .tracking_copy
+                    .write(global_state_key, StoredValue::AddressableEntity(entity))
+            } else {
+                return Ok(HOST_ERROR_INVALID_INPUT);
+            }
+        }
         Ok(Some(_)) => {
             // If it's a named key pointing to a URef, prune both the named key and the URef.
-            if let Keyspace::NamedKey(_) = keyspace {
+            if let Keyspace::NamedKey(_) | Keyspace::Context(_) = keyspace {
                 if let Ok(Some(StoredValue::NamedKey(named_key_value))) =
                     caller.context_mut().tracking_copy.read(&global_state_key)
                 {
@@ -363,9 +425,7 @@ pub fn casper_remove<S: GlobalStateReader, E: Executor>(
                     }
                 }
             }
-
-            // Produce a prune transform for the named key
-            caller.context_mut().tracking_copy.prune(global_state_key);
+            caller.context_mut().tracking_copy.prune(global_state_key)
         }
         Ok(None) => {
             // Entry does not exist, and we can't proceed with the prune operation
@@ -469,6 +529,7 @@ pub fn casper_read<S: GlobalStateReader, E: Executor>(
             }
             Keyspace::PaymentInfo(key_name)
         }
+        KeyspaceTag::AssociatedKeys => Keyspace::AssociatedKeys(&key_payload_bytes),
     };
 
     let global_state_key = match keyspace_to_global_state_key(caller.context(), keyspace) {
@@ -481,6 +542,13 @@ pub fn casper_read<S: GlobalStateReader, E: Executor>(
     let global_state_read_result = caller.context_mut().tracking_copy.read(&global_state_key);
 
     let global_state_raw_bytes: Cow<[u8]> = match global_state_read_result {
+        Ok(Some(StoredValue::AddressableEntity(entity))) => {
+            let entity_bytes = match entity.to_bytes() {
+                Ok(bytes) => bytes,
+                Err(_) => return Ok(HOST_ERROR_INVALID_DATA),
+            };
+            Cow::Owned(entity_bytes)
+        }
         Ok(Some(StoredValue::CLValue(cl_value))) => {
             let CLType::Any = cl_value.cl_type() else {
                 return Err(InternalHostError::TypeConversion)?;
@@ -588,6 +656,7 @@ fn keyspace_to_global_state_key<S: GlobalStateReader, E: Executor>(
                 EntryPointAddr::new_v1_entry_point_addr(entity_addr, payload).ok()?;
             Some(Key::EntryPoint(entry_point_addr))
         }
+        Keyspace::AssociatedKeys(_) => Some(Key::AddressableEntity(entity_addr)),
     }
 }
 

smart_contracts/contracts/vm2/vm2-assoc-keys/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "vm2-assoc-keys"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+[lib]
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+casper-contract-sdk = { path = "../../../sdk" }

smart_contracts/contracts/vm2/vm2-assoc-keys/build.rs

@@ -0,0 +1,7 @@
+fn main() {
+    // Check if target arch is wasm32 and set link flags accordingly
+    if std::env::var("TARGET").unwrap() == "wasm32-unknown-unknown" {
+        println!("cargo:rustc-link-arg=--import-memory");
+        println!("cargo:rustc-link-arg=--export-table");
+    }
+}

smart_contracts/contracts/vm2/vm2-assoc-keys/src/lib.rs

@@ -0,0 +1,18 @@
+#![cfg_attr(target_arch = "wasm32", no_main)]
+#![cfg_attr(target_arch = "wasm32", no_std)]
+
+pub mod exports {
+    use casper_contract_sdk::{casper_executor_wasm_common::keyspace::Keyspace, prelude::*};
+
+    #[casper(export)]
+    pub fn call(weight: u8) {
+        let account_hash_bytes = [10u8; 32];
+        if weight == 0 {
+            let keyspace = Keyspace::AssociatedKeys(&account_hash_bytes);
+            casper::remove(keyspace).unwrap()
+        } else {
+            let keyspace = Keyspace::AssociatedKeys(&account_hash_bytes);
+            casper::write(keyspace, &[weight]).unwrap();
+        }
+    }
+}

smart_contracts/sdk/src/casper.rs

@@ -114,6 +114,9 @@ pub fn read<F: FnOnce(usize) -> Option<ptr::NonNull<u8>>>(
         Keyspace::Context(key_bytes) => (KeyspaceTag::Context as u64, key_bytes),
         Keyspace::NamedKey(key_bytes) => (KeyspaceTag::NamedKey as u64, key_bytes.as_bytes()),
         Keyspace::PaymentInfo(payload) => (KeyspaceTag::PaymentInfo as u64, payload.as_bytes()),
+        Keyspace::AssociatedKeys(account_hash_bytes) => {
+            (KeyspaceTag::AssociatedKeys as u64, account_hash_bytes)
+        }
     };
 
     let mut info = casper_contract_sdk_sys::ReadInfo {
@@ -160,6 +163,9 @@ pub fn write(key: Keyspace, value: &[u8]) -> Result<(), HostResult> {
         Keyspace::Context(key_bytes) => (KeyspaceTag::Context as u64, key_bytes),
         Keyspace::NamedKey(key_bytes) => (KeyspaceTag::NamedKey as u64, key_bytes.as_bytes()),
         Keyspace::PaymentInfo(payload) => (KeyspaceTag::PaymentInfo as u64, payload.as_bytes()),
+        Keyspace::AssociatedKeys(account_hash_bytes) => {
+            (KeyspaceTag::AssociatedKeys as u64, account_hash_bytes)
+        }
     };
     let ret = unsafe {
         casper_contract_sdk_sys::casper_write(
@@ -180,6 +186,9 @@ pub fn remove(key: Keyspace) -> Result<(), HostResult> {
         Keyspace::Context(key_bytes) => (KeyspaceTag::Context as u64, key_bytes),
         Keyspace::NamedKey(key_bytes) => (KeyspaceTag::NamedKey as u64, key_bytes.as_bytes()),
         Keyspace::PaymentInfo(payload) => (KeyspaceTag::PaymentInfo as u64, payload.as_bytes()),
+        Keyspace::AssociatedKeys(account_hash_bytes) => {
+            (KeyspaceTag::AssociatedKeys as u64, account_hash_bytes)
+        }
     };
     let ret = unsafe {
         casper_contract_sdk_sys::casper_remove(key_space, key_bytes.as_ptr(), key_bytes.len())

vm2-build-contracts.sh

@@ -4,6 +4,7 @@ set -e
 VM2_BINS=(
   "vm2-harness"
   "vm2-cep18-caller"
+  "vm2-assoc-keys"
 )
 
 VM2_LIBS=(