Upgradable contracts using set_code function (#10690)

* poc logic

* set_code_hash impl, tests, benchmark

* Address @xgreenx's comments

* Move func defs closer to set_storage

* Check if code exists

- increment/decrement codehash refcount

* Document error for non-existing code hash

* Revert unrelated change

* Changes due to @athei's review

* Fix error handling

- comment errors: ReturnCodes
- update mock ext implementation
- return Error::CodeNotFound when no code for such hash

* Emit ContractCodeUpdated when setting new code_hash

* Address @athei's comments

* Move related defs to the bottom

* Minor comment update

Co-authored-by: Alexander Theißen <alex.theissen@me.com>

* Improve docs

* Improve docs

* Update frame/contracts/src/wasm/runtime.rs

Co-authored-by: Alexander Theißen <alex.theissen@me.com>

* Refactor set_code_hash test

* Minor change to benchmark

Co-authored-by: Alexander Theißen <alex.theissen@me.com>

* Minor change to benchmark

Co-authored-by: Alexander Theißen <alex.theissen@me.com>

* Minor comment refactor

Co-authored-by: Alexander Theißen <alex.theissen@me.com>

* Address @HCastano's comments

* Update seal_set_code_hash comment

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

* Move set_code_hash after delegate_call

* Move function to the bottom

* Moved and changed banchmark, added verify block

* Bring back previous benchmark

* Remove skip_meta for seal_set_code_hash

* Bring back skip_meta for seal_set_storage_per_new_kb

* Apply weights

Co-authored-by: Alexander Theißen <alex.theissen@me.com>
Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

frame/contracts/fixtures/new_set_code_hash_contract.wat

@@ -0,0 +1,13 @@
+(module
+	(import "seal0" "seal_return" (func $seal_return (param i32 i32 i32)))
+	(import "env" "memory" (memory 1 1))
+
+	;; [0, 32) return value
+	(data (i32.const 0) "\02")
+
+	(func (export "deploy"))
+
+	(func (export "call")
+		(call $seal_return (i32.const 0) (i32.const 0) (i32.const 4))
+	)
+)

frame/contracts/fixtures/set_code_hash.wat

@@ -0,0 +1,43 @@
+(module
+	(import "seal0" "seal_input" (func $seal_input (param i32 i32)))
+	(import "seal0" "seal_return" (func $seal_return (param i32 i32 i32)))
+	(import "__unstable__" "seal_set_code_hash" (func $seal_set_code_hash (param i32) (result i32)))
+
+	(import "env" "memory" (memory 1 1))
+
+	;; [0, 32) here we store input
+
+	;; [32, 36) input size
+	(data (i32.const 32) "\20")
+
+	;; [36, 40) return value
+	(data (i32.const 36) "\01")
+
+	(func $assert (param i32)
+		(block $ok
+			(br_if $ok
+				(get_local 0)
+			)
+			(unreachable)
+		)
+	)
+
+	(func (export "call")
+		(local $exit_code i32)
+
+		(call $seal_input (i32.const 0) (i32.const 32))
+
+		(set_local $exit_code
+			(call $seal_set_code_hash (i32.const 0)) ;; Pointer to the input data.
+		)
+		(call $assert
+			(i32.eq (get_local $exit_code) (i32.const 0)) ;; ReturnCode::Success
+		)
+
+		;; we return 1 after setting new code_hash
+		;; next `call` will NOT return this value, because contract code has been changed
+		(call $seal_return (i32.const 0) (i32.const 36) (i32.const 4))
+	)
+
+	(func (export "deploy"))
+)

frame/contracts/src/benchmarking/mod.rs

@@ -1971,6 +1971,46 @@ benchmarks! {
 		let origin = RawOrigin::Signed(instance.caller.clone());
 	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
 
+	seal_set_code_hash {
+		let r in 0 .. API_BENCHMARK_BATCHES;
+		let code_hashes = (0..r * API_BENCHMARK_BATCH_SIZE)
+			.map(|i| {
+				let new_code = WasmModule::<T>::dummy_with_bytes(i);
+				Contracts::<T>::store_code_raw(new_code.code, whitelisted_caller())?;
+				Ok(new_code.hash)
+			})
+			.collect::<Result<Vec<_>, &'static str>>()?;
+		let code_hash_len = code_hashes.get(0).map(|x| x.encode().len()).unwrap_or(0);
+		let code_hashes_bytes = code_hashes.iter().flat_map(|x| x.encode()).collect::<Vec<_>>();
+		let code_hashes_len = code_hashes_bytes.len();
+
+		let code = WasmModule::<T>::from(ModuleDefinition {
+			memory: Some(ImportedMemory::max::<T>()),
+			imported_functions: vec![ImportedFunction {
+				module: "__unstable__",
+				name: "seal_set_code_hash",
+				params: vec![
+					ValueType::I32,
+				],
+				return_type: Some(ValueType::I32),
+			}],
+			data_segments: vec![
+				DataSegment {
+					offset: 0,
+					value: code_hashes_bytes,
+				},
+			],
+			call_body: Some(body::repeated_dyn(r * API_BENCHMARK_BATCH_SIZE, vec![
+				Counter(0, code_hash_len as u32), // code_hash_ptr
+				Regular(Instruction::Call(0)),
+				Regular(Instruction::Drop),
+			])),
+			.. Default::default()
+		});
+		let instance = Contract::<T>::new(code, vec![])?;
+		let origin = RawOrigin::Signed(instance.caller.clone());
+	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
+
 	// We make the assumption that pushing a constant and dropping a value takes roughly
 	// the same amount of time. We follow that `t.load` and `drop` both have the weight
 	// of this benchmark / 2. We need to make this assumption because there is no way

frame/contracts/src/exec.rs

@@ -18,6 +18,7 @@
 use crate::{
 	gas::GasMeter,
 	storage::{self, Storage, WriteOutcome},
+	wasm::{decrement_refcount, increment_refcount},
 	AccountCounter, BalanceOf, CodeHash, Config, ContractInfo, ContractInfoOf, Error, Event,
 	Pallet as Contracts, Schedule,
 };
@@ -239,6 +240,9 @@ pub trait Ext: sealing::Sealed {
 	/// Tests sometimes need to modify and inspect the contract info directly.
 	#[cfg(test)]
 	fn contract_info(&mut self) -> &mut ContractInfo<Self::T>;
+
+	/// Sets new code hash for existing contract.
+	fn set_code_hash(&mut self, hash: CodeHash<Self::T>) -> Result<(), DispatchError>;
 }
 
 /// Describes the different functions that can be exported by an [`Executable`].
@@ -1182,6 +1186,20 @@ where
 	fn contract_info(&mut self) -> &mut ContractInfo<Self::T> {
 		self.top_frame_mut().contract_info()
 	}
+
+	fn set_code_hash(&mut self, hash: CodeHash<Self::T>) -> Result<(), DispatchError> {
+		increment_refcount::<Self::T>(hash)?;
+		let top_frame = self.top_frame_mut();
+		let prev_hash = top_frame.contract_info().code_hash.clone();
+		decrement_refcount::<Self::T>(prev_hash.clone())?;
+		top_frame.contract_info().code_hash = hash;
+		Contracts::<Self::T>::deposit_event(Event::ContractCodeUpdated {
+			contract: top_frame.account_id.clone(),
+			new_code_hash: hash,
+			old_code_hash: prev_hash,
+		});
+		Ok(())
+	}
 }
 
 fn deposit_event<T: Config>(topics: Vec<T::Hash>, event: Event<T>) {

frame/contracts/src/lib.rs

@@ -564,6 +564,16 @@ pub mod pallet {
 
 		/// A code with the specified hash was removed.
 		CodeRemoved { code_hash: T::Hash },
+
+		/// A contract's code was updated.
+		ContractCodeUpdated {
+			/// The contract that has been updated.
+			contract: T::AccountId,
+			/// New code hash that was set for the contract.
+			new_code_hash: T::Hash,
+			/// Previous code hash of the contract.
+			old_code_hash: T::Hash,
+		},
 	}
 
 	#[pallet::error]

frame/contracts/src/schedule.rs

@@ -328,6 +328,9 @@ pub struct HostFnWeights<T: Config> {
 	/// Weight per overwritten byte of an item stored with `seal_set_storage`.
 	pub set_storage_per_old_byte: Weight,
 
+	/// Weight of calling `seal_set_code_hash`.
+	pub set_code_hash: Weight,
+
 	/// Weight of calling `seal_clear_storage`.
 	pub clear_storage: Weight,
 
@@ -606,6 +609,7 @@ impl<T: Config> Default for HostFnWeights<T> {
 			),
 			debug_message: cost_batched!(seal_debug_message),
 			set_storage: cost_batched!(seal_set_storage),
+			set_code_hash: cost_batched!(seal_set_code_hash),
 			set_storage_per_new_byte: cost_byte_batched!(seal_set_storage_per_new_kb),
 			set_storage_per_old_byte: cost_byte_batched!(seal_set_storage_per_old_kb),
 			clear_storage: cost_batched!(seal_clear_storage),

frame/contracts/src/tests.rs

@@ -3023,3 +3023,64 @@ fn code_rejected_error_works() {
 		);
 	});
 }
+
+#[test]
+#[cfg(feature = "unstable-interface")]
+fn set_code_hash() {
+	let (wasm, code_hash) = compile_module::<Test>("set_code_hash").unwrap();
+	let (new_wasm, new_code_hash) = compile_module::<Test>("new_set_code_hash_contract").unwrap();
+
+	let contract_addr = Contracts::contract_address(&ALICE, &code_hash, &[]);
+
+	ExtBuilder::default().existential_deposit(100).build().execute_with(|| {
+		let _ = Balances::deposit_creating(&ALICE, 1_000_000);
+
+		// Instantiate the 'caller'
+		assert_ok!(Contracts::instantiate_with_code(
+			Origin::signed(ALICE),
+			300_000,
+			GAS_LIMIT,
+			None,
+			wasm,
+			vec![],
+			vec![],
+		));
+		// upload new code
+		assert_ok!(Contracts::upload_code(Origin::signed(ALICE), new_wasm.clone(), None));
+
+		// First call sets new code_hash and returns 1
+		let result = Contracts::bare_call(
+			ALICE,
+			contract_addr.clone(),
+			0,
+			GAS_LIMIT,
+			None,
+			new_code_hash.as_ref().to_vec(),
+			true,
+		)
+		.result
+		.unwrap();
+		assert_return_code!(result, 1);
+
+		// Second calls new contract code that returns 2
+		let result =
+			Contracts::bare_call(ALICE, contract_addr.clone(), 0, GAS_LIMIT, None, vec![], true)
+				.result
+				.unwrap();
+		assert_return_code!(result, 2);
+
+		// Checking for the last event only
+		assert_eq!(
+			System::events().pop().unwrap(),
+			EventRecord {
+				phase: Phase::Initialization,
+				event: Event::Contracts(crate::Event::ContractCodeUpdated {
+					contract: contract_addr.clone(),
+					new_code_hash: new_code_hash.clone(),
+					old_code_hash: code_hash.clone(),
+				}),
+				topics: vec![],
+			},
+		);
+	});
+}

frame/contracts/src/wasm/code_cache.rs

@@ -117,6 +117,22 @@ pub fn decrement_refcount<T: Config>(code_hash: CodeHash<T>) -> Result<(), Dispa
 	Ok(())
 }
 
+/// Increment the refcount of a code in-storage by one.
+///
+/// # Errors
+///
+/// [`Error::CodeNotFound`] is returned if the specified `code_hash` does not exist.
+pub fn increment_refcount<T: Config>(code_hash: CodeHash<T>) -> Result<(), DispatchError> {
+	<OwnerInfoOf<T>>::mutate(code_hash, |existing| -> Result<(), DispatchError> {
+		if let Some(info) = existing {
+			info.refcount = info.refcount.saturating_add(1);
+			Ok(())
+		} else {
+			Err(Error::<T>::CodeNotFound.into())
+		}
+	})
+}
+
 /// Try to remove code together with all associated information.
 pub fn try_remove<T: Config>(origin: &T::AccountId, code_hash: CodeHash<T>) -> DispatchResult {
 	<OwnerInfoOf<T>>::try_mutate_exists(&code_hash, |existing| {