WIP: check if there is a contract underneath specified hash

- increment/decrement codehash refcount

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.
@@ -156,7 +157,7 @@ pub trait Ext: sealing::Sealed {
 	) -> Result<WriteOutcome, DispatchError>;
 
 	/// Sets new code hash for existing contract
-	fn set_code_hash(&mut self, hash: CodeHash<Self::T>);
+	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>;
@@ -893,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>
@@ -1019,10 +1026,17 @@ where
 		)
 	}
 
-	fn set_code_hash(&mut self, code: CodeHash<Self::T>) {
+	fn set_code_hash(&mut self, hash: CodeHash<Self::T>) -> Result<(), DispatchError> {
 		let top_frame = &mut self.top_frame_mut();
-		top_frame.contract_info().code_hash = code;
+		let prev_hash = top_frame.contract_info().code_hash.clone();
+		top_frame.contract_info().code_hash = hash;
+		if !Self::code_exists(&hash) {
+			return Err(DispatchError::CannotLookup)
+		}
 		<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 {

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,8 +409,9 @@ mod tests {
 			}
 			Ok(result)
 		}
-		fn set_code_hash(&mut self, hash: CodeHash<Self::T>) {
+		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

frame/contracts/src/wasm/runtime.rs

@@ -831,7 +831,7 @@ define_env!(Env, <E: Ext>,
 	[__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);
+		ctx.ext.set_code_hash(code_hash)?;
 		Ok(ReturnCode::Success)
 	},