Add add_dependency / remove_dependency

frame/contracts/fixtures/add_dependency.wat

@@ -0,0 +1,95 @@
+;; This contract tests the behavior of adding / removing dependencies when delegate calling into a contract.
+(module
+	(import "seal0" "add_dependency" (func $add_dependency (param i32) (result i32)))
+	(import "seal0" "remove_dependency" (func $remove_dependency (param i32) (result i32)))
+	(import "seal0" "seal_input" (func $seal_input (param i32 i32)))
+	(import "seal0" "seal_delegate_call" (func $seal_delegate_call (param i32 i32 i32 i32 i32 i32) (result i32)))
+	(import "env" "memory" (memory 1 1))
+
+	(func $assert (param i32)
+		(block $ok
+			(br_if $ok
+				(get_local 0)
+			)
+			(unreachable)
+		)
+	)
+
+	;; This function loads input data and performs the action specified.
+	;; The first 4 bytes of the input specify the action to perform.
+	;; if the action is 1 we call add_dependency, if the action is 2 we call remove_dependency.
+	;; The next 32 bytes specify the code hash to use when calling add_dependency or remove_dependency.
+	(func $load_input
+		(local $action i32)
+		(local $code_hash_ptr i32)
+
+	    ;; Store available input size at offset 0.
+        (i32.store (i32.const 0) (i32.const 512))
+
+		;; Read input data
+		(call $seal_input (i32.const 4) (i32.const 0))
+
+		;; Input data layout.
+		;; [0..4) - size of the call
+		;; [4..8) - action to perform before calling delegate_call (1: add_dependency, 2: remove_dependency, default: nothing)
+		;; [8..42) - code hash of the callee
+		(set_local $action (i32.load (i32.const 4)))
+		(set_local $code_hash_ptr (i32.const 8))
+
+		;; Assert input size == 36 (4 for action + 32 for code_hash).
+		(call $assert
+			(i32.eq
+				(i32.load (i32.const 0))
+				(i32.const 36)
+			)
+		)
+
+		;; Call add_dependency when action == 1.
+		(if (i32.eq (get_local $action) (i32.const 1))
+		    (then
+				(call $assert (i32.eqz
+					(call $add_dependency
+						(get_local $code_hash_ptr)
+					)
+				))
+			)
+			(else)
+		)
+
+		;; Call remove_dependency when action == 2.
+		(if (i32.eq (get_local $action) (i32.const 2))
+		    (then
+				(call $assert (i32.eqz
+					(call $remove_dependency
+						(get_local $code_hash_ptr)
+					)
+				))
+			)
+			(else)
+		)
+	)
+
+	(func (export "deploy")
+		(call $load_input)
+	)
+
+	(func (export "call")
+		(call $load_input)
+
+		;; Delegate call into passed code hash.
+		(call $assert
+			(i32.eq
+				(call $seal_delegate_call
+					(i32.const 0)	;; Set no call flags.
+					(i32.const 8)	;; Pointer to "callee" code_hash.
+					(i32.const 0)	;; Input is ignored.
+					(i32.const 0)	;; Length of the input.
+					(i32.const 4294967295)	;; u32 max sentinel value: do not copy output.
+					(i32.const 0)	;; Length is ignored in this case.
+				)
+				(i32.const 0)
+			)
+		)
+	)
+
+)

frame/contracts/src/exec.rs

@@ -18,8 +18,9 @@
 use crate::{
 	gas::GasMeter,
 	storage::{self, DepositAccount, WriteOutcome},
+	wasm::{decrement_refcount, increment_refcount},
 	BalanceOf, CodeHash, Config, ContractInfo, ContractInfoOf, DebugBufferVec, Determinism, Error,
-	Event, Nonce, Origin, Pallet as Contracts, Schedule, System, LOG_TARGET,
+	Event, Nonce, Origin, OwnerInfoOf, Pallet as Contracts, Schedule, System, LOG_TARGET,
 };
 use frame_support::{
 	crypto::ecdsa::ECDSAExt,
@@ -35,14 +36,17 @@ use frame_support::{
 	Blake2_128Concat, BoundedVec, StorageHasher,
 };
 use frame_system::RawOrigin;
-use pallet_contracts_primitives::ExecReturnValue;
+use pallet_contracts_primitives::{ExecReturnValue, StorageDeposit};
 use smallvec::{Array, SmallVec};
 use sp_core::{
 	ecdsa::Public as ECDSAPublic,
 	sr25519::{Public as SR25519Public, Signature as SR25519Signature},
 };
 use sp_io::{crypto::secp256k1_ecdsa_recover_compressed, hashing::blake2_256};
-use sp_runtime::traits::{Convert, Hash, Zero};
+use sp_runtime::{
+	traits::{Convert, Hash, Zero},
+	Perbill,
+};
 use sp_std::{marker::PhantomData, mem, prelude::*, vec::Vec};
 
 pub type AccountIdOf<T> = <T as frame_system::Config>::AccountId;
@@ -306,6 +310,22 @@ pub trait Ext: sealing::Sealed {
 
 	/// Returns a nonce that is incremented for every instantiated contract.
 	fn nonce(&mut self) -> u64;
+
+	/// Add a contract dependency, This is useful for contract delegation, to make sure that the
+	/// delegated contract is not removed while it is still in use.
+	/// This will increase the reference count of the code hash, and charge a deposit of 30% of the
+	/// code deposit.
+	///
+	/// Returns an error if we have reached the maximum number of dependencies, or the dependency
+	/// already exists.
+	fn add_dependency(&mut self, _code: CodeHash<Self::T>) -> Result<(), DispatchError>;
+
+	/// Remove a contract dependency.
+	/// This is the counterpart of `add_dependency`. This method will decrease the reference count
+	/// And refund the deposit that was charged by `add_dependency`.
+	///
+	/// Returns an error if the dependency does not exist.
+	fn remove_dependency(&mut self, _code: CodeHash<Self::T>) -> Result<(), DispatchError>;
 }
 
 /// Describes the different functions that can be exported by an [`Executable`].
@@ -1461,6 +1481,30 @@ where
 			current
 		}
 	}
+	fn add_dependency(&mut self, code_hash: CodeHash<Self::T>) -> Result<(), DispatchError> {
+		let frame = self.top_frame_mut();
+		let info = frame.contract_info.get(&frame.account_id);
+
+		increment_refcount::<T>(code_hash)?;
+		let owner_info = OwnerInfoOf::<T>::get(code_hash).ok_or(Error::<T>::ContractNotFound)?;
+		let deposit = Perbill::from_percent(30).mul_ceil(owner_info.deposit());
+
+		frame
+			.nested_storage
+			.charge_dependency(info.deposit_account(), &StorageDeposit::Charge(deposit));
+		info.add_dependency(code_hash, deposit)
+	}
+
+	fn remove_dependency(&mut self, code_hash: CodeHash<Self::T>) -> Result<(), DispatchError> {
+		let frame = self.top_frame_mut();
+		let info = frame.contract_info.get(&frame.account_id);
+		let deposit = info.remove_dependency(code_hash)?;
+		decrement_refcount::<T>(code_hash);
+		frame
+			.nested_storage
+			.charge_dependency(info.deposit_account(), &StorageDeposit::Refund(deposit));
+		Ok(())
+	}
 }
 
 mod sealing {

frame/contracts/src/lib.rs

@@ -861,6 +861,12 @@ pub mod pallet {
 		CodeRejected,
 		/// An indetermistic code was used in a context where this is not permitted.
 		Indeterministic,
+		/// The contract has reached its maximum number of dependencies.
+		MaxDependenciesReached,
+		/// The dependency was not found in the contract's dependencies.
+		DependencyNotFound,
+		/// The contract already depends on the given dependency.
+		DependencyAlreadyExists,
 	}
 
 	/// A mapping from an original code hash to the original code, untouched by instrumentation.

frame/contracts/src/schedule.rs

@@ -431,6 +431,12 @@ pub struct HostFnWeights<T: Config> {
 	/// Weight of calling `instantiation_nonce`.
 	pub instantiation_nonce: Weight,
 
+	/// Weight of calling `add_dependency`.
+	pub add_dependency: Weight,
+
+	/// Weight of calling `remove_dependency`.
+	pub remove_dependency: Weight,
+
 	/// The type parameter is used in the default implementation.
 	#[codec(skip)]
 	pub _phantom: PhantomData<T>,
@@ -637,6 +643,8 @@ impl<T: Config> Default for HostFnWeights<T> {
 			reentrance_count: cost!(seal_reentrance_count),
 			account_reentrance_count: cost!(seal_account_reentrance_count),
 			instantiation_nonce: cost!(seal_instantiation_nonce),
+			add_dependency: cost!(add_dependency),
+			remove_dependency: cost!(remove_dependency),
 			_phantom: PhantomData,
 		}
 	}

frame/contracts/src/storage.rs

@@ -33,10 +33,11 @@ use frame_support::{
 	DefaultNoBound, RuntimeDebugNoBound,
 };
 use scale_info::TypeInfo;
+use sp_core::ConstU32;
 use sp_io::KillStorageResult;
 use sp_runtime::{
 	traits::{Hash, Saturating, Zero},
-	RuntimeDebug,
+	BoundedBTreeMap, DispatchResult, RuntimeDebug,
 };
 use sp_std::{marker::PhantomData, ops::Deref, prelude::*};
 
@@ -66,6 +67,10 @@ pub struct ContractInfo<T: Config> {
 	/// We need to store this information separately so it is not used when calculating any refunds
 	/// since the base deposit can only ever be refunded on contract termination.
 	storage_base_deposit: BalanceOf<T>,
+
+	/// This tracks the code hash and storage deposit paid to ensure that these dependencies don't
+	/// get removed and thus can be used for delegate calls.
+	dependencies: BoundedBTreeMap<CodeHash<T>, BalanceOf<T>, ConstU32<32>>,
 }
 
 impl<T: Config> ContractInfo<T> {
@@ -101,6 +106,7 @@ impl<T: Config> ContractInfo<T> {
 			storage_byte_deposit: Zero::zero(),
 			storage_item_deposit: Zero::zero(),
 			storage_base_deposit: Zero::zero(),
+			dependencies: Default::default(),
 		};
 
 		Ok(contract)
@@ -201,6 +207,35 @@ impl<T: Config> ContractInfo<T> {
 		})
 	}
 
+	/// Add a new `code_hash` dependency to the contract.
+	/// The `amount` is the amount of funds that will be reserved for the dependency.
+	///
+	/// Returns an error if the maximum number of dependencies is reached or if
+	/// the dependency already exists.
+	pub fn add_dependency(
+		&mut self,
+		code_hash: CodeHash<T>,
+		amount: BalanceOf<T>,
+	) -> DispatchResult {
+		self.dependencies
+			.try_insert(code_hash, amount)
+			.map_err(|_| Error::<T>::MaxDependenciesReached)?
+			.map_or(Ok(()), |_| Err(Error::<T>::DependencyAlreadyExists))
+			.map_err(Into::into)
+	}
+
+	/// Remove a `code_hash` dependency from the contract.
+	///
+	/// Returns an error if the dependency doesn't exist.
+	pub fn remove_dependency(&mut self, dep: CodeHash<T>) -> Result<BalanceOf<T>, DispatchError> {
+		self.dependencies.remove(&dep).ok_or(Error::<T>::DependencyNotFound.into())
+	}
+
+	#[cfg(test)]
+	pub fn dependencies(&self) -> &BoundedBTreeMap<CodeHash<T>, BalanceOf<T>, ConstU32<32>> {
+		&self.dependencies
+	}
+
 	/// Push a contract's trie to the deletion queue for lazy removal.
 	///
 	/// You must make sure that the contract is also removed when queuing the trie for deletion.

frame/contracts/src/storage/meter.rs

@@ -406,6 +406,22 @@ where
 		};
 	}
 
+	/// Add a charge for updating the contract's dependencies.
+	pub fn charge_dependency(
+		&mut self,
+		deposit_account: &DepositAccount<T>,
+		amount: &DepositOf<T>,
+	) {
+		let charge = Charge {
+			deposit_account: deposit_account.clone(),
+			amount: amount.clone(),
+			terminated: false,
+		};
+
+		self.total_deposit = self.total_deposit.saturating_add(amount);
+		self.charges.push(charge);
+	}
+
 	/// Charge from `origin` a storage deposit for contract instantiation.
 	///
 	/// This immediately transfers the balance in order to create the account.
@@ -664,6 +680,7 @@ mod tests {
 			storage_byte_deposit: info.bytes_deposit,
 			storage_item_deposit: info.items_deposit,
 			storage_base_deposit: Default::default(),
+			dependencies: Default::default(),
 		}
 	}