1
// This file is part of Substrate.
2

            
3
// Copyright (C) Parity Technologies (UK) Ltd.
4
// SPDX-License-Identifier: Apache-2.0
5

            
6
// Licensed under the Apache License, Version 2.0 (the "License");
7
// you may not use this file except in compliance with the License.
8
// You may obtain a copy of the License at
9
//
10
// 	http://www.apache.org/licenses/LICENSE-2.0
11
//
12
// Unless required by applicable law or agreed to in writing, software
13
// distributed under the License is distributed on an "AS IS" BASIS,
14
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
// See the License for the specific language governing permissions and
16
// limitations under the License.
17

            
18
//! Storage migrations for the preimage pallet.
19

            
20
use super::*;
21
use alloc::collections::btree_map::BTreeMap;
22
use frame_support::{
23
	storage_alias,
24
	traits::{ConstU32, OnRuntimeUpgrade},
25
};
26

            
27
#[cfg(feature = "try-runtime")]
28
use frame_support::ensure;
29
#[cfg(feature = "try-runtime")]
30
use sp_runtime::TryRuntimeError;
31

            
32
/// The log target.
33
const TARGET: &'static str = "runtime::preimage::migration::v1";
34

            
35
/// The original data layout of the preimage pallet without a specific version number.
36
mod v0 {
37
	use super::*;
38

            
39
	#[derive(Clone, Eq, PartialEq, Encode, Decode, TypeInfo, MaxEncodedLen, RuntimeDebug)]
40
	pub enum OldRequestStatus<AccountId, Balance> {
41
		Unrequested(Option<(AccountId, Balance)>),
42
		Requested(u32),
43
	}
44

            
45
	#[storage_alias]
46
	pub type PreimageFor<T: Config> = StorageMap<
47
		Pallet<T>,
48
		Identity,
49
		<T as frame_system::Config>::Hash,
50
		BoundedVec<u8, ConstU32<MAX_SIZE>>,
51
	>;
52

            
53
	#[storage_alias]
54
	pub type StatusFor<T: Config> = StorageMap<
55
		Pallet<T>,
56
		Identity,
57
		<T as frame_system::Config>::Hash,
58
		OldRequestStatus<<T as frame_system::Config>::AccountId, BalanceOf<T>>,
59
	>;
60

            
61
	/// Returns the number of images or `None` if the storage is corrupted.
62
	#[cfg(feature = "try-runtime")]
63
	pub fn image_count<T: Config>() -> Option<u32> {
64
		let images = v0::PreimageFor::<T>::iter_values().count() as u32;
65
		let status = v0::StatusFor::<T>::iter_values().count() as u32;
66

            
67
		if images == status {
68
			Some(images)
69
		} else {
70
			None
71
		}
72
	}
73
}
74

            
75
pub mod v1 {
76
	use super::*;
77

            
78
	/// Migration for moving preimage from V0 to V1 storage.
79
	///
80
	/// Note: This needs to be run with the same hashing algorithm as before
81
	/// since it is not re-hashing the preimages.
82
	pub struct Migration<T>(core::marker::PhantomData<T>);
83

            
84
	impl<T: Config> OnRuntimeUpgrade for Migration<T> {
85
		#[cfg(feature = "try-runtime")]
86
		fn pre_upgrade() -> Result<Vec<u8>, TryRuntimeError> {
87
			ensure!(StorageVersion::get::<Pallet<T>>() == 0, "can only upgrade from version 0");
88

            
89
			let images = v0::image_count::<T>().expect("v0 storage corrupted");
90
			log::info!(target: TARGET, "Migrating {} images", &images);
91
			Ok((images as u32).encode())
92
		}
93

            
94
		fn on_runtime_upgrade() -> Weight {
95
			let mut weight = T::DbWeight::get().reads(1);
96
			if StorageVersion::get::<Pallet<T>>() != 0 {
97
				log::warn!(
98
					target: TARGET,
99
					"skipping MovePreimagesIntoBuckets: executed on wrong storage version.\
100
				Expected version 0"
101
				);
102
				return weight
103
			}
104

            
105
			let status = v0::StatusFor::<T>::drain().collect::<Vec<_>>();
106
			weight.saturating_accrue(T::DbWeight::get().reads(status.len() as u64));
107

            
108
			let preimages = v0::PreimageFor::<T>::drain().collect::<BTreeMap<_, _>>();
109
			weight.saturating_accrue(T::DbWeight::get().reads(preimages.len() as u64));
110

            
111
			for (hash, status) in status.into_iter() {
112
				let preimage = if let Some(preimage) = preimages.get(&hash) {
113
					preimage
114
				} else {
115
					log::error!(target: TARGET, "preimage not found for hash {:?}", &hash);
116
					continue
117
				};
118
				let len = preimage.len() as u32;
119
				if len > MAX_SIZE {
120
					log::error!(
121
						target: TARGET,
122
						"preimage too large for hash {:?}, len: {}",
123
						&hash,
124
						len
125
					);
126
					continue
127
				}
128

            
129
				let status = match status {
130
					v0::OldRequestStatus::Unrequested(deposit) => match deposit {
131
						Some(deposit) => OldRequestStatus::Unrequested { deposit, len },
132
						// `None` depositor becomes system-requested.
133
						None =>
134
							OldRequestStatus::Requested { deposit: None, count: 1, len: Some(len) },
135
					},
136
					v0::OldRequestStatus::Requested(0) => {
137
						log::error!(target: TARGET, "preimage has counter of zero: {:?}", hash);
138
						continue
139
					},
140
					v0::OldRequestStatus::Requested(count) =>
141
						OldRequestStatus::Requested { deposit: None, count, len: Some(len) },
142
				};
143
				log::trace!(target: TARGET, "Moving preimage {:?} with len {}", hash, len);
144

            
145
				#[allow(deprecated)]
146
				crate::StatusFor::<T>::insert(hash, status);
147
				crate::PreimageFor::<T>::insert(&(hash, len), preimage);
148

            
149
				weight.saturating_accrue(T::DbWeight::get().writes(2));
150
			}
151
			StorageVersion::new(1).put::<Pallet<T>>();
152

            
153
			weight.saturating_add(T::DbWeight::get().writes(1))
154
		}
155

            
156
		#[cfg(feature = "try-runtime")]
157
		fn post_upgrade(state: Vec<u8>) -> DispatchResult {
158
			let old_images: u32 =
159
				Decode::decode(&mut &state[..]).expect("pre_upgrade provides a valid state; qed");
160
			let new_images = image_count::<T>().expect("V1 storage corrupted");
161

            
162
			if new_images != old_images {
163
				log::error!(
164
					target: TARGET,
165
					"migrated {} images, expected {}",
166
					new_images,
167
					old_images
168
				);
169
			}
170
			ensure!(StorageVersion::get::<Pallet<T>>() == 1, "must upgrade");
171
			Ok(())
172
		}
173
	}
174

            
175
	/// Returns the number of images or `None` if the storage is corrupted.
176
	#[cfg(feature = "try-runtime")]
177
	pub fn image_count<T: Config>() -> Option<u32> {
178
		// Use iter_values() to ensure that the values are decodable.
179
		let images = crate::PreimageFor::<T>::iter_values().count() as u32;
180
		#[allow(deprecated)]
181
		let status = crate::StatusFor::<T>::iter_values().count() as u32;
182

            
183
		if images == status {
184
			Some(images)
185
		} else {
186
			None
187
		}
188
	}
189
}
190

            
191
#[cfg(test)]
192
#[cfg(feature = "try-runtime")]
193
mod test {
194
	#![allow(deprecated)]
195
	use super::*;
196
	use crate::mock::{Test as T, *};
197

            
198
	use sp_runtime::bounded_vec;
199

            
200
	#[test]
201
	fn migration_works() {
202
		new_test_ext().execute_with(|| {
203
			assert_eq!(StorageVersion::get::<Pallet<T>>(), 0);
204
			// Insert some preimages into the v0 storage:
205

            
206
			// Case 1: Unrequested without deposit
207
			let (p, h) = preimage::<T>(128);
208
			v0::PreimageFor::<T>::insert(h, p);
209
			v0::StatusFor::<T>::insert(h, v0::OldRequestStatus::Unrequested(None));
210
			// Case 2: Unrequested with deposit
211
			let (p, h) = preimage::<T>(1024);
212
			v0::PreimageFor::<T>::insert(h, p);
213
			v0::StatusFor::<T>::insert(h, v0::OldRequestStatus::Unrequested(Some((1, 1))));
214
			// Case 3: Requested by 0 (invalid)
215
			let (p, h) = preimage::<T>(8192);
216
			v0::PreimageFor::<T>::insert(h, p);
217
			v0::StatusFor::<T>::insert(h, v0::OldRequestStatus::Requested(0));
218
			// Case 4: Requested by 10
219
			let (p, h) = preimage::<T>(65536);
220
			v0::PreimageFor::<T>::insert(h, p);
221
			v0::StatusFor::<T>::insert(h, v0::OldRequestStatus::Requested(10));
222

            
223
			assert_eq!(v0::image_count::<T>(), Some(4));
224
			assert_eq!(v1::image_count::<T>(), None, "V1 storage should be corrupted");
225

            
226
			let state = v1::Migration::<T>::pre_upgrade().unwrap();
227
			let _w = v1::Migration::<T>::on_runtime_upgrade();
228
			v1::Migration::<T>::post_upgrade(state).unwrap();
229

            
230
			// V0 and V1 share the same prefix, so `iter_values` still counts the same.
231
			assert_eq!(v0::image_count::<T>(), Some(3));
232
			assert_eq!(v1::image_count::<T>(), Some(3)); // One gets skipped therefore 3.
233
			assert_eq!(StorageVersion::get::<Pallet<T>>(), 1);
234

            
235
			// Case 1: Unrequested without deposit becomes system-requested
236
			let (p, h) = preimage::<T>(128);
237
			assert_eq!(crate::PreimageFor::<T>::get(&(h, 128)), Some(p));
238
			assert_eq!(
239
				crate::StatusFor::<T>::get(h),
240
				Some(OldRequestStatus::Requested { deposit: None, count: 1, len: Some(128) })
241
			);
242
			// Case 2: Unrequested with deposit becomes unrequested
243
			let (p, h) = preimage::<T>(1024);
244
			assert_eq!(crate::PreimageFor::<T>::get(&(h, 1024)), Some(p));
245
			assert_eq!(
246
				crate::StatusFor::<T>::get(h),
247
				Some(OldRequestStatus::Unrequested { deposit: (1, 1), len: 1024 })
248
			);
249
			// Case 3: Requested by 0 should be skipped
250
			let (_, h) = preimage::<T>(8192);
251
			assert_eq!(crate::PreimageFor::<T>::get(&(h, 8192)), None);
252
			assert_eq!(crate::StatusFor::<T>::get(h), None);
253
			// Case 4: Requested by 10 becomes requested by 10
254
			let (p, h) = preimage::<T>(65536);
255
			assert_eq!(crate::PreimageFor::<T>::get(&(h, 65536)), Some(p));
256
			assert_eq!(
257
				crate::StatusFor::<T>::get(h),
258
				Some(OldRequestStatus::Requested { deposit: None, count: 10, len: Some(65536) })
259
			);
260
		});
261
	}
262

            
263
	/// Returns a preimage with a given size and its hash.
264
	fn preimage<T: Config>(
265
		len: usize,
266
	) -> (BoundedVec<u8, ConstU32<MAX_SIZE>>, <T as frame_system::Config>::Hash) {
267
		let p = bounded_vec![1; len];
268
		let h = <T as frame_system::Config>::Hashing::hash_of(&p);
269
		(p, h)
270
	}
271
}