Upgradable contracts using set_code function

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

@@ -808,6 +808,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 as u32,
+					value: code_hashes_bytes,
+				},
+			],
+			call_body: Some(body::repeated_dyn(r * API_BENCHMARK_BATCH_SIZE, vec![
+				Counter(0 as u32, 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![])
+
 	seal_set_storage_per_kb {
 		let n in 0 .. T::Schedule::get().limits.payload_len / 1024;
 		let key = T::Hashing::hash_of(&1u32).as_ref().to_vec();

frame/contracts/src/exec.rs

@@ -18,8 +18,9 @@
 use crate::{
 	gas::GasMeter,
 	storage::{self, Storage, WriteOutcome},
-	AccountCounter, BalanceOf, CodeHash, Config, ContractInfo, ContractInfoOf, Error, Event,
-	Pallet as Contracts, Schedule,
+	wasm::{decrement_refcount, increment_refcount},
+	AccountCounter, BalanceOf, CodeHash, CodeStorage, Config, ContractInfo, ContractInfoOf, Error,
+	Event, OwnerInfoOf, Pallet as Contracts, PristineCode, Schedule,
 };
 use frame_support::{
 	dispatch::{DispatchError, DispatchResult, DispatchResultWithPostInfo, Dispatchable},
@@ -140,7 +141,7 @@ pub trait Ext: sealing::Sealed {
 	/// was deleted.
 	fn get_storage(&mut self, key: &StorageKey) -> Option<Vec<u8>>;
 
-	/// Returns true iff some storage entry exists under the supplied `key`
+	/// Returns true if some storage entry exists under the supplied `key`
 	///
 	/// Returns `false` if the `key` wasn't previously set by `set_storage` or
 	/// was deleted.
@@ -155,6 +156,9 @@ pub trait Ext: sealing::Sealed {
 		take_old: bool,
 	) -> Result<WriteOutcome, DispatchError>;
 
+	/// Sets new code hash for existing contract
+	fn set_code_hash(&mut self, hash: CodeHash<Self::T>) -> Result<(), DispatchError>;
+
 	/// Returns a reference to the account id of the caller.
 	fn caller(&self) -> &AccountIdOf<Self::T>;
 
@@ -890,6 +894,12 @@ where
 	fn initial_trie_seed() -> u64 {
 		<AccountCounter<T>>::get().wrapping_add(1)
 	}
+
+	fn code_exists(hash: &CodeHash<T>) -> bool {
+		<PristineCode<T>>::contains_key(hash) &&
+			<CodeStorage<T>>::contains_key(&hash) &&
+			<OwnerInfoOf<T>>::contains_key(&hash)
+	}
 }
 
 impl<'a, T, E> Ext for Stack<'a, T, E>
@@ -1016,6 +1026,19 @@ where
 		)
 	}
 
+	fn set_code_hash(&mut self, hash: CodeHash<Self::T>) -> Result<(), DispatchError> {
+		let top_frame = &mut self.top_frame_mut();
+		let prev_hash = top_frame.contract_info().code_hash.clone();
+		if !Self::code_exists(&hash) {
+			return Err(DispatchError::CannotLookup)
+		}
+		top_frame.contract_info().code_hash = hash;
+		<ContractInfoOf<T>>::insert(top_frame.account_id.clone(), top_frame.contract_info());
+		increment_refcount::<Self::T>(hash)?;
+		decrement_refcount::<Self::T>(prev_hash)?;
+		Ok(())
+	}
+
 	fn address(&self) -> &T::AccountId {
 		&self.top_frame().account_id
 	}

frame/contracts/src/schedule.rs

@@ -319,6 +319,9 @@ pub struct HostFnWeights<T: Config> {
 	/// Weight per byte of an item stored with `seal_set_storage`.
 	pub set_storage_per_byte: Weight,
 
+	/// Weight of calling `seal_set_code_hash`.
+	pub set_code_hash: Weight,
+
 	/// Weight of calling `seal_clear_storage`.
 	pub clear_storage: Weight,
 
@@ -587,6 +590,7 @@ impl<T: Config> Default for HostFnWeights<T> {
 			debug_message: cost_batched!(seal_debug_message),
 			set_storage: cost_batched!(seal_set_storage),
 			set_storage_per_byte: cost_byte_batched!(seal_set_storage_per_kb),
+			set_code_hash: cost_batched!(seal_set_code_hash),
 			clear_storage: cost_batched!(seal_clear_storage),
 			contains_storage: cost_batched!(seal_contains_storage),
 			get_storage: cost_batched!(seal_get_storage),

frame/contracts/src/tests.rs

@@ -691,6 +691,53 @@ fn deploy_and_call_other_contract() {
 	});
 }
 
+#[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);
+	});
+}
+
 #[test]
 fn cannot_self_destruct_through_draning() {
 	let (wasm, code_hash) = compile_module::<Test>("drain").unwrap();

frame/contracts/src/wasm/code_cache.rs

@@ -118,6 +118,16 @@ pub fn decrement_refcount<T: Config>(code_hash: CodeHash<T>) -> Result<(), Dispa
 	Ok(())
 }
 
+/// Increment the refcount of a code in-storage by one.
+pub fn increment_refcount<T: Config>(code_hash: CodeHash<T>) -> Result<(), DispatchError> {
+	<OwnerInfoOf<T>>::mutate(code_hash, |existing| {
+		if let Some(info) = existing {
+			info.refcount = info.refcount.saturating_add(1);
+		}
+	});
+	Ok(())
+}
+
 /// 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| {

frame/contracts/src/wasm/mod.rs

@@ -26,7 +26,10 @@ mod runtime;
 
 #[cfg(feature = "runtime-benchmarks")]
 pub use self::code_cache::reinstrument;
-pub use self::runtime::{ReturnCode, Runtime, RuntimeCosts};
+pub use self::{
+	code_cache::{decrement_refcount, increment_refcount},
+	runtime::{ReturnCode, Runtime, RuntimeCosts},
+};
 use crate::{
 	exec::{ExecResult, Executable, ExportedFunction, Ext},
 	gas::GasMeter,
@@ -406,6 +409,10 @@ mod tests {
 			}
 			Ok(result)
 		}
+		fn set_code_hash(&mut self, hash: CodeHash<Self::T>) -> Result<(), DispatchError> {
+			self.storage.insert(StorageKey::from(hash), vec![1, 2, 3]);
+			Ok(())
+		}
 		fn caller(&self) -> &AccountIdOf<Self::T> {
 			&ALICE
 		}
@@ -2180,6 +2187,47 @@ mod tests {
 		assert_eq!(&result.data.0[4..], &[0u8; 0]);
 	}
 
+	#[test]
+	#[cfg(feature = "unstable-interface")]
+	fn set_code_hash() {
+		const CODE: &str = r#"
+(module
+	(import "__unstable__" "seal_set_code_hash" (func $seal_set_code_hash (param i32) (result i32)))
+	(import "env" "memory" (memory 1 1))
+	(func $assert (param i32)
+		(block $ok
+			(br_if $ok
+				(get_local 0)
+			)
+			(unreachable)
+		)
+	)
+	(func (export "call")
+		(local $exit_code i32)
+		(set_local $exit_code
+			(call $seal_set_code_hash (i32.const 0))
+		)
+		(call $assert
+			(i32.eq (get_local $exit_code) (i32.const 0)) ;; ReturnCode::Success
+		)
+	)
+
+	(func (export "deploy"))
+
+	;; Hash of code.
+	(data (i32.const 0)
+		"\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11"
+		"\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11"
+	)
+)
+"#;
+
+		let mut mock_ext = MockExt::default();
+		execute(CODE, [0u8; 32].encode(), &mut mock_ext).unwrap();
+
+		assert_eq!(mock_ext.storage.get(&[17u8; 32]), Some(&vec![1, 2, 3]));
+	}
+
 	#[test]
 	#[cfg(feature = "unstable-interface")]
 	fn contains_storage_works() {

frame/contracts/src/wasm/runtime.rs

@@ -171,6 +171,9 @@ pub enum RuntimeCosts {
 	DebugMessage,
 	/// Weight of calling `seal_set_storage` for the given storage item size.
 	SetStorage(u32),
+	/// Weight of calling `seal_set_code_hash`
+	#[cfg(feature = "unstable-interface")]
+	SetCodeHash,
 	/// Weight of calling `seal_clear_storage`.
 	ClearStorage,
 	/// Weight of calling `seal_contains_storage`.
@@ -251,6 +254,8 @@ impl RuntimeCosts {
 			DebugMessage => s.debug_message,
 			SetStorage(len) =>
 				s.set_storage.saturating_add(s.set_storage_per_byte.saturating_mul(len.into())),
+			#[cfg(feature = "unstable-interface")]
+			SetCodeHash => s.set_code_hash,
 			ClearStorage => s.clear_storage,
 			#[cfg(feature = "unstable-interface")]
 			ContainsStorage => s.contains_storage,
@@ -813,6 +818,24 @@ define_env!(Env, <E: Ext>,
 		ctx.set_storage(key_ptr, value_ptr, value_len).map(|_| ())
 	},
 
+	// Replace contract code
+	//
+	// # Parameters
+	//
+	// - code_hash_ptr: a pointer to the buffer that contains new contract code
+	//
+	// # Errors
+	//
+	// - requested buffer is not within the bounds of the sandbox memory.
+	// - the buffer contents cannot be decoded as the required type.
+	// - code for specified hash does not exist
+	[__unstable__] seal_set_code_hash(ctx, code_hash_ptr: u32) -> ReturnCode => {
+		ctx.charge_gas(RuntimeCosts::SetCodeHash)?;
+		let code_hash: CodeHash<<E as Ext>::T> = ctx.read_sandbox_memory_as(code_hash_ptr)?;
+		ctx.ext.set_code_hash(code_hash)?;
+		Ok(ReturnCode::Success)
+	},
+
 	// Set the value at the given key in the contract storage.
 	//
 	// The value length must not exceed the maximum defined by the contracts module parameters.