contracts: Add set_code root dispatchable (#11451)

* Add `set_code` dispatchable

* cargo run --quiet --profile=production  --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark pallet --chain=dev --steps=50 --repeat=20 --pallet=pallet_contracts --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/contracts/src/weights.rs --template=./.maintain/frame-weight-template.hbs

Co-authored-by: Parity Bot <admin@parity.io>

frame/contracts/src/benchmarking/mod.rs

@@ -391,6 +391,20 @@ benchmarks! {
 		assert!(<Contract<T>>::code_removed(&hash));
 	}
 
+	set_code {
+		let instance = <Contract<T>>::with_caller(
+			whitelisted_caller(), WasmModule::dummy(), vec![],
+		)?;
+		// we just add some bytes so that the code hash is different
+		let WasmModule { code, hash, .. } = <WasmModule<T>>::dummy_with_bytes(128);
+		<Contracts<T>>::store_code_raw(code, instance.caller.clone())?;
+		let callee = instance.addr.clone();
+		assert_ne!(instance.info()?.code_hash, hash);
+	}: _(RawOrigin::Root, callee, hash)
+	verify {
+		assert_eq!(instance.info()?.code_hash, hash);
+	}
+
 	seal_caller {
 		let r in 0 .. API_BENCHMARK_BATCHES;
 		let instance = Contract::<T>::new(WasmModule::getter(

frame/contracts/src/lib.rs

@@ -586,6 +586,42 @@ pub mod pallet {
 			// we waive the fee because removing unused code is beneficial
 			Ok(Pays::No.into())
 		}
+
+		/// Privileged function that changes the code of an existing contract.
+		///
+		/// This takes care of updating refcounts and all other necessary operations. Returns
+		/// an error if either the `code_hash` or `dest` do not exist.
+		///
+		/// # Note
+		///
+		/// This does **not** change the address of the contract in question. This means
+		/// that the contract address is no longer derived from its code hash after calling
+		/// this dispatchable.
+		#[pallet::weight(T::WeightInfo::set_code())]
+		pub fn set_code(
+			origin: OriginFor<T>,
+			dest: <T::Lookup as StaticLookup>::Source,
+			code_hash: CodeHash<T>,
+		) -> DispatchResult {
+			ensure_root(origin)?;
+			let dest = T::Lookup::lookup(dest)?;
+			<ContractInfoOf<T>>::try_mutate(&dest, |contract| {
+				let contract = if let Some(contract) = contract {
+					contract
+				} else {
+					return Err(<Error<T>>::ContractNotFound.into())
+				};
+				<PrefabWasmModule<T>>::add_user(code_hash)?;
+				<PrefabWasmModule<T>>::remove_user(contract.code_hash);
+				Self::deposit_event(Event::ContractCodeUpdated {
+					contract: dest.clone(),
+					new_code_hash: code_hash,
+					old_code_hash: contract.code_hash,
+				});
+				contract.code_hash = code_hash;
+				Ok(())
+			})
+		}
 	}
 
 	#[pallet::event]

frame/contracts/src/tests.rs

@@ -40,7 +40,7 @@ use frame_support::{
 	weights::{constants::WEIGHT_PER_SECOND, DispatchClass, PostDispatchInfo, Weight},
 };
 use frame_system::{self as system, EventRecord, Phase};
-use pretty_assertions::assert_eq;
+use pretty_assertions::{assert_eq, assert_ne};
 use sp_core::Bytes;
 use sp_io::hashing::blake2_256;
 use sp_keystore::{testing::KeyStore, KeystoreExt};
@@ -2862,6 +2862,86 @@ fn storage_deposit_works() {
 	});
 }
 
+#[test]
+fn set_code_extrinsic() {
+	let (wasm, code_hash) = compile_module::<Test>("dummy").unwrap();
+	let (new_wasm, new_code_hash) = compile_module::<Test>("crypto_hashes").unwrap();
+
+	assert_ne!(code_hash, new_code_hash);
+
+	ExtBuilder::default().existential_deposit(100).build().execute_with(|| {
+		let _ = Balances::deposit_creating(&ALICE, 1_000_000);
+
+		assert_ok!(Contracts::instantiate_with_code(
+			Origin::signed(ALICE),
+			0,
+			GAS_LIMIT,
+			None,
+			wasm,
+			vec![],
+			vec![],
+		));
+		let addr = Contracts::contract_address(&ALICE, &code_hash, &[]);
+
+		assert_ok!(Contracts::upload_code(Origin::signed(ALICE), new_wasm, None,));
+
+		// Drop previous events
+		initialize_block(2);
+
+		assert_eq!(<ContractInfoOf<Test>>::get(&addr).unwrap().code_hash, code_hash);
+		assert_refcount!(&code_hash, 1);
+		assert_refcount!(&new_code_hash, 0);
+
+		// only root can execute this extrinsic
+		assert_noop!(
+			Contracts::set_code(Origin::signed(ALICE), addr.clone(), new_code_hash),
+			sp_runtime::traits::BadOrigin,
+		);
+		assert_eq!(<ContractInfoOf<Test>>::get(&addr).unwrap().code_hash, code_hash);
+		assert_refcount!(&code_hash, 1);
+		assert_refcount!(&new_code_hash, 0);
+		assert_eq!(System::events(), vec![],);
+
+		// contract must exist
+		assert_noop!(
+			Contracts::set_code(Origin::root(), BOB, new_code_hash),
+			<Error<Test>>::ContractNotFound,
+		);
+		assert_eq!(<ContractInfoOf<Test>>::get(&addr).unwrap().code_hash, code_hash);
+		assert_refcount!(&code_hash, 1);
+		assert_refcount!(&new_code_hash, 0);
+		assert_eq!(System::events(), vec![],);
+
+		// new code hash must exist
+		assert_noop!(
+			Contracts::set_code(Origin::root(), addr.clone(), Default::default()),
+			<Error<Test>>::CodeNotFound,
+		);
+		assert_eq!(<ContractInfoOf<Test>>::get(&addr).unwrap().code_hash, code_hash);
+		assert_refcount!(&code_hash, 1);
+		assert_refcount!(&new_code_hash, 0);
+		assert_eq!(System::events(), vec![],);
+
+		// successful call
+		assert_ok!(Contracts::set_code(Origin::root(), addr.clone(), new_code_hash));
+		assert_eq!(<ContractInfoOf<Test>>::get(&addr).unwrap().code_hash, new_code_hash);
+		assert_refcount!(&code_hash, 0);
+		assert_refcount!(&new_code_hash, 1);
+		assert_eq!(
+			System::events(),
+			vec![EventRecord {
+				phase: Phase::Initialization,
+				event: Event::Contracts(pallet_contracts::Event::ContractCodeUpdated {
+					contract: addr,
+					new_code_hash,
+					old_code_hash: code_hash,
+				}),
+				topics: vec![],
+			},]
+		);
+	});
+}
+
 #[test]
 fn call_after_killed_account_needs_funding() {
 	let (wasm, code_hash) = compile_module::<Test>("dummy").unwrap();