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
//! Concrete externalities implementation.
19

            
20
#[cfg(feature = "std")]
21
use crate::overlayed_changes::OverlayedExtensions;
22
use crate::{
23
	backend::Backend, IndexOperation, IterArgs, OverlayedChanges, StorageKey, StorageValue,
24
};
25
use codec::{Compact, CompactLen, Decode, Encode};
26
use hash_db::Hasher;
27
#[cfg(feature = "std")]
28
use sp_core::hexdisplay::HexDisplay;
29
use sp_core::storage::{
30
	well_known_keys::is_child_storage_key, ChildInfo, StateVersion, TrackedStorageKey,
31
};
32
use sp_externalities::{Extension, ExtensionStore, Externalities, MultiRemovalResults};
33

            
34
use crate::{trace, warn};
35
use alloc::{boxed::Box, vec::Vec};
36
use core::{
37
	any::{Any, TypeId},
38
	cmp::Ordering,
39
};
40
#[cfg(feature = "std")]
41
use std::error;
42

            
43
const EXT_NOT_ALLOWED_TO_FAIL: &str = "Externalities not allowed to fail within runtime";
44
const BENCHMARKING_FN: &str = "\
45
	This is a special fn only for benchmarking where a database commit happens from the runtime.
46
	For that reason client started transactions before calling into runtime are not allowed.
47
	Without client transactions the loop condition guarantees the success of the tx close.";
48

            
49
#[cfg(feature = "std")]
50
321822
fn guard() -> sp_panic_handler::AbortGuard {
51
321822
	sp_panic_handler::AbortGuard::force_abort()
52
321822
}
53

            
54
#[cfg(not(feature = "std"))]
55
fn guard() -> () {
56
	()
57
}
58

            
59
/// Errors that can occur when interacting with the externalities.
60
#[cfg(feature = "std")]
61
#[derive(Debug, Copy, Clone)]
62
pub enum Error<B, E> {
63
	/// Failure to load state data from the backend.
64
	#[allow(unused)]
65
	Backend(B),
66
	/// Failure to execute a function.
67
	#[allow(unused)]
68
	Executor(E),
69
}
70

            
71
#[cfg(feature = "std")]
72
impl<B: std::fmt::Display, E: std::fmt::Display> std::fmt::Display for Error<B, E> {
73
	fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
74
		match *self {
75
			Error::Backend(ref e) => write!(f, "Storage backend error: {}", e),
76
			Error::Executor(ref e) => write!(f, "Sub-call execution error: {}", e),
77
		}
78
	}
79
}
80

            
81
#[cfg(feature = "std")]
82
impl<B: error::Error, E: error::Error> error::Error for Error<B, E> {
83
	fn description(&self) -> &str {
84
		match *self {
85
			Error::Backend(..) => "backend error",
86
			Error::Executor(..) => "executor error",
87
		}
88
	}
89
}
90

            
91
/// Wraps a read-only backend, call executor, and current overlayed changes.
92
pub struct Ext<'a, H, B>
93
where
94
	H: Hasher,
95
	B: 'a + Backend<H>,
96
{
97
	/// The overlayed changes to write to.
98
	overlay: &'a mut OverlayedChanges<H>,
99
	/// The storage backend to read from.
100
	backend: &'a B,
101
	/// Pseudo-unique id used for tracing.
102
	pub id: u16,
103
	/// Extensions registered with this instance.
104
	#[cfg(feature = "std")]
105
	extensions: Option<OverlayedExtensions<'a>>,
106
}
107

            
108
impl<'a, H, B> Ext<'a, H, B>
109
where
110
	H: Hasher,
111
	B: Backend<H>,
112
{
113
	/// Create a new `Ext`.
114
	#[cfg(not(feature = "std"))]
115
	pub fn new(overlay: &'a mut OverlayedChanges<H>, backend: &'a B) -> Self {
116
		Ext { overlay, backend, id: 0 }
117
	}
118

            
119
	/// Create a new `Ext` from overlayed changes and read-only backend
120
	#[cfg(feature = "std")]
121
3540042
	pub fn new(
122
3540042
		overlay: &'a mut OverlayedChanges<H>,
123
3540042
		backend: &'a B,
124
3540042
		extensions: Option<&'a mut sp_externalities::Extensions>,
125
3540042
	) -> Self {
126
3540042
		Self {
127
3540042
			overlay,
128
3540042
			backend,
129
3540042
			id: rand::random(),
130
3540042
			extensions: extensions.map(OverlayedExtensions::new),
131
3540042
		}
132
3540042
	}
133
}
134

            
135
#[cfg(test)]
136
impl<'a, H, B> Ext<'a, H, B>
137
where
138
	H: Hasher,
139
	H::Out: Ord + 'static,
140
	B: 'a + Backend<H>,
141
{
142
	pub fn storage_pairs(&mut self) -> Vec<(StorageKey, StorageValue)> {
143
		use std::collections::HashMap;
144

            
145
		self.backend
146
			.pairs(Default::default())
147
			.expect("never fails in tests; qed.")
148
			.map(|key_value| key_value.expect("never fails in tests; qed."))
149
			.map(|(k, v)| (k, Some(v)))
150
			.chain(self.overlay.changes_mut().map(|(k, v)| (k.clone(), v.value().cloned())))
151
			.collect::<HashMap<_, _>>()
152
			.into_iter()
153
			.filter_map(|(k, maybe_val)| maybe_val.map(|val| (k, val)))
154
			.collect()
155
	}
156
}
157

            
158
impl<'a, H, B> Externalities for Ext<'a, H, B>
159
where
160
	H: Hasher,
161
	H::Out: Ord + 'static + codec::Codec,
162
	B: Backend<H>,
163
{
164
	fn set_offchain_storage(&mut self, key: &[u8], value: Option<&[u8]>) {
165
		self.overlay.set_offchain_storage(key, value)
166
	}
167

            
168
107274
	fn storage(&mut self, key: &[u8]) -> Option<StorageValue> {
169
107274
		let _guard = guard();
170
107274
		let result = self
171
107274
			.overlay
172
107274
			.storage(key)
173
107274
			.map(|x| x.map(|x| x.to_vec()))
174
107274
			.unwrap_or_else(|| self.backend.storage(key).expect(EXT_NOT_ALLOWED_TO_FAIL));
175
107274

            
176
107274
		// NOTE: be careful about touching the key names – used outside substrate!
177
107274
		trace!(
178
			target: "state",
179
			method = "Get",
180
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
181
			key = %HexDisplay::from(&key),
182
			result = ?result.as_ref().map(HexDisplay::from),
183
			result_encoded = %HexDisplay::from(
184
				&result
185
					.as_ref()
186
					.map(|v| EncodeOpaqueValue(v.clone()))
187
					.encode()
188
			),
189
		);
190

            
191
107274
		result
192
107274
	}
193

            
194
	fn storage_hash(&mut self, key: &[u8]) -> Option<Vec<u8>> {
195
		let _guard = guard();
196
		let result = self
197
			.overlay
198
			.storage(key)
199
			.map(|x| x.map(|x| H::hash(x)))
200
			.unwrap_or_else(|| self.backend.storage_hash(key).expect(EXT_NOT_ALLOWED_TO_FAIL));
201

            
202
		trace!(
203
			target: "state",
204
			method = "Hash",
205
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
206
			key = %HexDisplay::from(&key),
207
			?result,
208
		);
209
		result.map(|r| r.encode())
210
	}
211

            
212
	fn child_storage(&mut self, child_info: &ChildInfo, key: &[u8]) -> Option<StorageValue> {
213
		let _guard = guard();
214
		let result = self
215
			.overlay
216
			.child_storage(child_info, key)
217
			.map(|x| x.map(|x| x.to_vec()))
218
			.unwrap_or_else(|| {
219
				self.backend.child_storage(child_info, key).expect(EXT_NOT_ALLOWED_TO_FAIL)
220
			});
221

            
222
		trace!(
223
			target: "state",
224
			method = "ChildGet",
225
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
226
			child_info = %HexDisplay::from(&child_info.storage_key()),
227
			key = %HexDisplay::from(&key),
228
			result = ?result.as_ref().map(HexDisplay::from)
229
		);
230

            
231
		result
232
	}
233

            
234
	fn child_storage_hash(&mut self, child_info: &ChildInfo, key: &[u8]) -> Option<Vec<u8>> {
235
		let _guard = guard();
236
		let result = self
237
			.overlay
238
			.child_storage(child_info, key)
239
			.map(|x| x.map(|x| H::hash(x)))
240
			.unwrap_or_else(|| {
241
				self.backend.child_storage_hash(child_info, key).expect(EXT_NOT_ALLOWED_TO_FAIL)
242
			});
243

            
244
		trace!(
245
			target: "state",
246
			method = "ChildHash",
247
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
248
			child_info = %HexDisplay::from(&child_info.storage_key()),
249
			key = %HexDisplay::from(&key),
250
			?result,
251
		);
252

            
253
		result.map(|r| r.encode())
254
	}
255

            
256
160911
	fn exists_storage(&mut self, key: &[u8]) -> bool {
257
160911
		let _guard = guard();
258
160911
		let result = match self.overlay.storage(key) {
259
53637
			Some(x) => x.is_some(),
260
107274
			_ => self.backend.exists_storage(key).expect(EXT_NOT_ALLOWED_TO_FAIL),
261
		};
262

            
263
160911
		trace!(
264
			target: "state",
265
			method = "Exists",
266
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
267
			key = %HexDisplay::from(&key),
268
			%result,
269
		);
270

            
271
160911
		result
272
160911
	}
273

            
274
	fn exists_child_storage(&mut self, child_info: &ChildInfo, key: &[u8]) -> bool {
275
		let _guard = guard();
276

            
277
		let result = match self.overlay.child_storage(child_info, key) {
278
			Some(x) => x.is_some(),
279
			_ => self
280
				.backend
281
				.exists_child_storage(child_info, key)
282
				.expect(EXT_NOT_ALLOWED_TO_FAIL),
283
		};
284

            
285
		trace!(
286
			target: "state",
287
			method = "ChildExists",
288
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
289
			child_info = %HexDisplay::from(&child_info.storage_key()),
290
			key = %HexDisplay::from(&key),
291
			%result,
292
		);
293
		result
294
	}
295

            
296
	fn next_storage_key(&mut self, key: &[u8]) -> Option<StorageKey> {
297
		let mut next_backend_key =
298
			self.backend.next_storage_key(key).expect(EXT_NOT_ALLOWED_TO_FAIL);
299
		let mut overlay_changes = self.overlay.iter_after(key).peekable();
300

            
301
		match (&next_backend_key, overlay_changes.peek()) {
302
			(_, None) => next_backend_key,
303
			(Some(_), Some(_)) => {
304
				for overlay_key in overlay_changes {
305
					let cmp = next_backend_key.as_deref().map(|v| v.cmp(overlay_key.0));
306

            
307
					// If `backend_key` is less than the `overlay_key`, we found out next key.
308
					if cmp == Some(Ordering::Less) {
309
						return next_backend_key
310
					} else if overlay_key.1.value().is_some() {
311
						// If there exists a value for the `overlay_key` in the overlay
312
						// (aka the key is still valid), it means we have found our next key.
313
						return Some(overlay_key.0.to_vec())
314
					} else if cmp == Some(Ordering::Equal) {
315
						// If the `backend_key` and `overlay_key` are equal, it means that we need
316
						// to search for the next backend key, because the overlay has overwritten
317
						// this key.
318
						next_backend_key = self
319
							.backend
320
							.next_storage_key(overlay_key.0)
321
							.expect(EXT_NOT_ALLOWED_TO_FAIL);
322
					}
323
				}
324

            
325
				next_backend_key
326
			},
327
			(None, Some(_)) => {
328
				// Find the next overlay key that has a value attached.
329
				overlay_changes.find_map(|k| k.1.value().as_ref().map(|_| k.0.to_vec()))
330
			},
331
		}
332
	}
333

            
334
	fn next_child_storage_key(&mut self, child_info: &ChildInfo, key: &[u8]) -> Option<StorageKey> {
335
		let mut next_backend_key = self
336
			.backend
337
			.next_child_storage_key(child_info, key)
338
			.expect(EXT_NOT_ALLOWED_TO_FAIL);
339
		let mut overlay_changes =
340
			self.overlay.child_iter_after(child_info.storage_key(), key).peekable();
341

            
342
		match (&next_backend_key, overlay_changes.peek()) {
343
			(_, None) => next_backend_key,
344
			(Some(_), Some(_)) => {
345
				for overlay_key in overlay_changes {
346
					let cmp = next_backend_key.as_deref().map(|v| v.cmp(overlay_key.0));
347

            
348
					// If `backend_key` is less than the `overlay_key`, we found out next key.
349
					if cmp == Some(Ordering::Less) {
350
						return next_backend_key
351
					} else if overlay_key.1.value().is_some() {
352
						// If there exists a value for the `overlay_key` in the overlay
353
						// (aka the key is still valid), it means we have found our next key.
354
						return Some(overlay_key.0.to_vec())
355
					} else if cmp == Some(Ordering::Equal) {
356
						// If the `backend_key` and `overlay_key` are equal, it means that we need
357
						// to search for the next backend key, because the overlay has overwritten
358
						// this key.
359
						next_backend_key = self
360
							.backend
361
							.next_child_storage_key(child_info, overlay_key.0)
362
							.expect(EXT_NOT_ALLOWED_TO_FAIL);
363
					}
364
				}
365

            
366
				next_backend_key
367
			},
368
			(None, Some(_)) => {
369
				// Find the next overlay key that has a value attached.
370
				overlay_changes.find_map(|k| k.1.value().as_ref().map(|_| k.0.to_vec()))
371
			},
372
		}
373
	}
374

            
375
214548
	fn place_storage(&mut self, key: StorageKey, value: Option<StorageValue>) {
376
214548
		let _guard = guard();
377
214548
		if is_child_storage_key(&key) {
378
			warn!(target: "trie", "Refuse to directly set child storage key");
379
			return
380
214548
		}
381
214548

            
382
214548
		// NOTE: be careful about touching the key names – used outside substrate!
383
214548
		trace!(
384
			target: "state",
385
			method = "Put",
386
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
387
			key = %HexDisplay::from(&key),
388
			value = ?value.as_ref().map(HexDisplay::from),
389
			value_encoded = %HexDisplay::from(
390
				&value
391
					.as_ref()
392
					.map(|v| EncodeOpaqueValue(v.clone()))
393
					.encode()
394
			),
395
		);
396

            
397
214548
		self.overlay.set_storage(key, value);
398
214548
	}
399

            
400
	fn place_child_storage(
401
		&mut self,
402
		child_info: &ChildInfo,
403
		key: StorageKey,
404
		value: Option<StorageValue>,
405
	) {
406
		trace!(
407
			target: "state",
408
			method = "ChildPut",
409
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
410
			child_info = %HexDisplay::from(&child_info.storage_key()),
411
			key = %HexDisplay::from(&key),
412
			value = ?value.as_ref().map(HexDisplay::from),
413
		);
414
		let _guard = guard();
415

            
416
		self.overlay.set_child_storage(child_info, key, value);
417
	}
418

            
419
	fn kill_child_storage(
420
		&mut self,
421
		child_info: &ChildInfo,
422
		maybe_limit: Option<u32>,
423
		maybe_cursor: Option<&[u8]>,
424
	) -> MultiRemovalResults {
425
		trace!(
426
			target: "state",
427
			method = "ChildKill",
428
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
429
			child_info = %HexDisplay::from(&child_info.storage_key()),
430
		);
431
		let _guard = guard();
432
		let overlay = self.overlay.clear_child_storage(child_info);
433
		let (maybe_cursor, backend, loops) =
434
			self.limit_remove_from_backend(Some(child_info), None, maybe_limit, maybe_cursor);
435
		MultiRemovalResults { maybe_cursor, backend, unique: overlay + backend, loops }
436
	}
437

            
438
	fn clear_prefix(
439
		&mut self,
440
		prefix: &[u8],
441
		maybe_limit: Option<u32>,
442
		maybe_cursor: Option<&[u8]>,
443
	) -> MultiRemovalResults {
444
		trace!(
445
			target: "state",
446
			method = "ClearPrefix",
447
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
448
			prefix = %HexDisplay::from(&prefix),
449
		);
450
		let _guard = guard();
451

            
452
		if sp_core::storage::well_known_keys::starts_with_child_storage_key(prefix) {
453
			warn!(
454
				target: "trie",
455
				"Refuse to directly clear prefix that is part or contains of child storage key",
456
			);
457
			return MultiRemovalResults { maybe_cursor: None, backend: 0, unique: 0, loops: 0 }
458
		}
459

            
460
		let overlay = self.overlay.clear_prefix(prefix);
461
		let (maybe_cursor, backend, loops) =
462
			self.limit_remove_from_backend(None, Some(prefix), maybe_limit, maybe_cursor);
463
		MultiRemovalResults { maybe_cursor, backend, unique: overlay + backend, loops }
464
	}
465

            
466
	fn clear_child_prefix(
467
		&mut self,
468
		child_info: &ChildInfo,
469
		prefix: &[u8],
470
		maybe_limit: Option<u32>,
471
		maybe_cursor: Option<&[u8]>,
472
	) -> MultiRemovalResults {
473
		trace!(
474
			target: "state",
475
			method = "ChildClearPrefix",
476
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
477
			child_info = %HexDisplay::from(&child_info.storage_key()),
478
			prefix = %HexDisplay::from(&prefix),
479
		);
480
		let _guard = guard();
481

            
482
		let overlay = self.overlay.clear_child_prefix(child_info, prefix);
483
		let (maybe_cursor, backend, loops) = self.limit_remove_from_backend(
484
			Some(child_info),
485
			Some(prefix),
486
			maybe_limit,
487
			maybe_cursor,
488
		);
489
		MultiRemovalResults { maybe_cursor, backend, unique: overlay + backend, loops }
490
	}
491

            
492
	fn storage_append(&mut self, key: Vec<u8>, value: Vec<u8>) {
493
		trace!(
494
			target: "state",
495
			method = "Append",
496
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
497
			key = %HexDisplay::from(&key),
498
			value = %HexDisplay::from(&value),
499
		);
500

            
501
		let _guard = guard();
502

            
503
		let backend = &mut self.backend;
504
		self.overlay.append_storage(key.clone(), value, || {
505
			backend.storage(&key).expect(EXT_NOT_ALLOWED_TO_FAIL).unwrap_or_default()
506
		});
507
	}
508

            
509
	fn storage_root(&mut self, state_version: StateVersion) -> Vec<u8> {
510
		let _guard = guard();
511

            
512
		let (root, _cached) = self.overlay.storage_root(self.backend, state_version);
513

            
514
		trace!(
515
			target: "state",
516
			method = "StorageRoot",
517
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
518
			storage_root = %HexDisplay::from(&root.as_ref()),
519
			cached = %_cached,
520
		);
521

            
522
		root.encode()
523
	}
524

            
525
	fn child_storage_root(
526
		&mut self,
527
		child_info: &ChildInfo,
528
		state_version: StateVersion,
529
	) -> Vec<u8> {
530
		let _guard = guard();
531

            
532
		let (root, _cached) = self
533
			.overlay
534
			.child_storage_root(child_info, self.backend, state_version)
535
			.expect(EXT_NOT_ALLOWED_TO_FAIL);
536

            
537
		trace!(
538
			target: "state",
539
			method = "ChildStorageRoot",
540
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
541
			child_info = %HexDisplay::from(&child_info.storage_key()),
542
			storage_root = %HexDisplay::from(&root.as_ref()),
543
			cached = %_cached,
544
		);
545

            
546
		root.encode()
547
	}
548

            
549
	fn storage_index_transaction(&mut self, index: u32, hash: &[u8], size: u32) {
550
		trace!(
551
			target: "state",
552
			method = "IndexTransaction",
553
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
554
			%index,
555
			tx_hash = %HexDisplay::from(&hash),
556
			%size,
557
		);
558

            
559
		self.overlay.add_transaction_index(IndexOperation::Insert {
560
			extrinsic: index,
561
			hash: hash.to_vec(),
562
			size,
563
		});
564
	}
565

            
566
	/// Renew existing piece of data storage.
567
	fn storage_renew_transaction_index(&mut self, index: u32, hash: &[u8]) {
568
		trace!(
569
			target: "state",
570
			method = "RenewTransactionIndex",
571
			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
572
			%index,
573
			tx_hash = %HexDisplay::from(&hash),
574
		);
575

            
576
		self.overlay
577
			.add_transaction_index(IndexOperation::Renew { extrinsic: index, hash: hash.to_vec() });
578
	}
579

            
580
	fn storage_start_transaction(&mut self) {
581
		self.overlay.start_transaction()
582
	}
583

            
584
	fn storage_rollback_transaction(&mut self) -> Result<(), ()> {
585
		self.overlay.rollback_transaction().map_err(|_| ())
586
	}
587

            
588
	fn storage_commit_transaction(&mut self) -> Result<(), ()> {
589
		self.overlay.commit_transaction().map_err(|_| ())
590
	}
591

            
592
	fn wipe(&mut self) {
593
		for _ in 0..self.overlay.transaction_depth() {
594
			self.overlay.rollback_transaction().expect(BENCHMARKING_FN);
595
		}
596
		self.overlay
597
			.drain_storage_changes(self.backend, Default::default())
598
			.expect(EXT_NOT_ALLOWED_TO_FAIL);
599
		self.backend.wipe().expect(EXT_NOT_ALLOWED_TO_FAIL);
600
		self.overlay
601
			.enter_runtime()
602
			.expect("We have reset the overlay above, so we can not be in the runtime; qed");
603
	}
604

            
605
	fn commit(&mut self) {
606
		// Bench always use latest state.
607
		let state_version = StateVersion::default();
608
		for _ in 0..self.overlay.transaction_depth() {
609
			self.overlay.commit_transaction().expect(BENCHMARKING_FN);
610
		}
611
		let changes = self
612
			.overlay
613
			.drain_storage_changes(self.backend, state_version)
614
			.expect(EXT_NOT_ALLOWED_TO_FAIL);
615
		self.backend
616
			.commit(
617
				changes.transaction_storage_root,
618
				changes.transaction,
619
				changes.main_storage_changes,
620
				changes.child_storage_changes,
621
			)
622
			.expect(EXT_NOT_ALLOWED_TO_FAIL);
623
		self.overlay
624
			.enter_runtime()
625
			.expect("We have reset the overlay above, so we can not be in the runtime; qed");
626
	}
627

            
628
	fn read_write_count(&self) -> (u32, u32, u32, u32) {
629
		self.backend.read_write_count()
630
	}
631

            
632
	fn reset_read_write_count(&mut self) {
633
		self.backend.reset_read_write_count()
634
	}
635

            
636
	fn get_whitelist(&self) -> Vec<TrackedStorageKey> {
637
		self.backend.get_whitelist()
638
	}
639

            
640
	fn set_whitelist(&mut self, new: Vec<TrackedStorageKey>) {
641
		self.backend.set_whitelist(new)
642
	}
643

            
644
	fn proof_size(&self) -> Option<u32> {
645
		self.backend.proof_size()
646
	}
647

            
648
	fn get_read_and_written_keys(&self) -> Vec<(Vec<u8>, u32, u32, bool)> {
649
		self.backend.get_read_and_written_keys()
650
	}
651
}
652

            
653
impl<'a, H, B> Ext<'a, H, B>
654
where
655
	H: Hasher,
656
	H::Out: Ord + 'static + codec::Codec,
657
	B: Backend<H>,
658
{
659
	fn limit_remove_from_backend(
660
		&mut self,
661
		child_info: Option<&ChildInfo>,
662
		prefix: Option<&[u8]>,
663
		maybe_limit: Option<u32>,
664
		start_at: Option<&[u8]>,
665
	) -> (Option<Vec<u8>>, u32, u32) {
666
		let iter = match self.backend.keys(IterArgs {
667
			child_info: child_info.cloned(),
668
			prefix,
669
			start_at,
670
			..IterArgs::default()
671
		}) {
672
			Ok(iter) => iter,
673
			Err(error) => {
674
				log::debug!(target: "trie", "Error while iterating the storage: {}", error);
675
				return (None, 0, 0)
676
			},
677
		};
678

            
679
		let mut delete_count: u32 = 0;
680
		let mut loop_count: u32 = 0;
681
		let mut maybe_next_key = None;
682
		for key in iter {
683
			let key = match key {
684
				Ok(key) => key,
685
				Err(error) => {
686
					log::debug!(target: "trie", "Error while iterating the storage: {}", error);
687
					break
688
				},
689
			};
690

            
691
			if maybe_limit.map_or(false, |limit| loop_count == limit) {
692
				maybe_next_key = Some(key);
693
				break
694
			}
695
			let overlay = match child_info {
696
				Some(child_info) => self.overlay.child_storage(child_info, &key),
697
				None => self.overlay.storage(&key),
698
			};
699
			if !matches!(overlay, Some(None)) {
700
				// not pending deletion from the backend - delete it.
701
				if let Some(child_info) = child_info {
702
					self.overlay.set_child_storage(child_info, key, None);
703
				} else {
704
					self.overlay.set_storage(key, None);
705
				}
706
				delete_count = delete_count.saturating_add(1);
707
			}
708
			loop_count = loop_count.saturating_add(1);
709
		}
710

            
711
		(maybe_next_key, delete_count, loop_count)
712
	}
713
}
714

            
715
/// Implement `Encode` by forwarding the stored raw vec.
716
struct EncodeOpaqueValue(Vec<u8>);
717

            
718
impl Encode for EncodeOpaqueValue {
719
	fn using_encoded<R, F: FnOnce(&[u8]) -> R>(&self, f: F) -> R {
720
		f(&self.0)
721
	}
722
}
723

            
724
/// Auxiliary structure for appending a value to a storage item.
725
pub(crate) struct StorageAppend<'a>(&'a mut Vec<u8>);
726

            
727
impl<'a> StorageAppend<'a> {
728
	/// Create a new instance using the given `storage` reference.
729
1109889
	pub fn new(storage: &'a mut Vec<u8>) -> Self {
730
1109889
		Self(storage)
731
1109889
	}
732

            
733
	/// Extract the length of the list like data structure.
734
129842
	pub fn extract_length(&self) -> Option<u32> {
735
129842
		Compact::<u32>::decode(&mut &self.0[..]).map(|c| c.0).ok()
736
129842
	}
737

            
738
	/// Replace the length in the encoded data.
739
	///
740
	/// If `old_length` is `None`, the previous length will be assumed to be `0`.
741
295458
	pub fn replace_length(&mut self, old_length: Option<u32>, new_length: u32) {
742
295458
		let old_len_encoded_len = old_length.map(|l| Compact::<u32>::compact_len(&l)).unwrap_or(0);
743
295458
		let new_len_encoded = Compact::<u32>(new_length).encode();
744
295458
		self.0.splice(0..old_len_encoded_len, new_len_encoded);
745
295458
	}
746

            
747
	/// Append the given `value` to the storage item.
748
	///
749
	/// If appending fails, `[value]` is stored in the storage item and we return false.
750
	#[cfg(any(test, feature = "fuzzing"))]
751
	pub fn append(&mut self, value: Vec<u8>) -> bool {
752
		use codec::EncodeAppend;
753
		let mut result = true;
754
		let value = vec![EncodeOpaqueValue(value)];
755

            
756
		let item = core::mem::take(self.0);
757

            
758
		*self.0 = match Vec::<EncodeOpaqueValue>::append_or_new(item, &value) {
759
			Ok(item) => item,
760
			Err(_) => {
761
				result = false;
762
				crate::log_error!(
763
					target: "runtime",
764
					"Failed to append value, resetting storage item to `[value]`.",
765
				);
766
				value.encode()
767
			},
768
		};
769
		result
770
	}
771

            
772
	/// Append to current buffer, do not touch the prefixed length.
773
444468
	pub fn append_raw(&mut self, mut value: Vec<u8>) {
774
444468
		self.0.append(&mut value)
775
444468
	}
776
}
777

            
778
#[cfg(not(feature = "std"))]
779
impl<'a, H, B> ExtensionStore for Ext<'a, H, B>
780
where
781
	H: Hasher,
782
	H::Out: Ord + 'static + codec::Codec,
783
	B: Backend<H>,
784
{
785
	fn extension_by_type_id(&mut self, _type_id: TypeId) -> Option<&mut dyn Any> {
786
		None
787
	}
788

            
789
	fn register_extension_with_type_id(
790
		&mut self,
791
		_type_id: TypeId,
792
		_extension: Box<dyn Extension>,
793
	) -> Result<(), sp_externalities::Error> {
794
		Err(sp_externalities::Error::ExtensionsAreNotSupported)
795
	}
796

            
797
	fn deregister_extension_by_type_id(
798
		&mut self,
799
		_type_id: TypeId,
800
	) -> Result<(), sp_externalities::Error> {
801
		Err(sp_externalities::Error::ExtensionsAreNotSupported)
802
	}
803
}
804

            
805
#[cfg(feature = "std")]
806
impl<'a, H, B> ExtensionStore for Ext<'a, H, B>
807
where
808
	H: Hasher,
809
	B: 'a + Backend<H>,
810
{
811
	fn extension_by_type_id(&mut self, type_id: TypeId) -> Option<&mut dyn Any> {
812
		self.extensions.as_mut().and_then(|exts| exts.get_mut(type_id))
813
	}
814

            
815
	fn register_extension_with_type_id(
816
		&mut self,
817
		type_id: TypeId,
818
		extension: Box<dyn Extension>,
819
	) -> Result<(), sp_externalities::Error> {
820
		if let Some(ref mut extensions) = self.extensions {
821
			extensions.register(type_id, extension)
822
		} else {
823
			Err(sp_externalities::Error::ExtensionsAreNotSupported)
824
		}
825
	}
826

            
827
	fn deregister_extension_by_type_id(
828
		&mut self,
829
		type_id: TypeId,
830
	) -> Result<(), sp_externalities::Error> {
831
		if let Some(ref mut extensions) = self.extensions {
832
			if extensions.deregister(type_id) {
833
				Ok(())
834
			} else {
835
				Err(sp_externalities::Error::ExtensionIsNotRegistered(type_id))
836
			}
837
		} else {
838
			Err(sp_externalities::Error::ExtensionsAreNotSupported)
839
		}
840
	}
841
}
842

            
843
#[cfg(test)]
844
mod tests {
845
	use super::*;
846
	use crate::InMemoryBackend;
847
	use codec::{Decode, Encode};
848
	use sp_core::{
849
		map,
850
		storage::{Storage, StorageChild},
851
		Blake2Hasher,
852
	};
853

            
854
	type TestBackend = InMemoryBackend<Blake2Hasher>;
855
	type TestExt<'a> = Ext<'a, Blake2Hasher, TestBackend>;
856

            
857
	#[test]
858
	fn next_storage_key_works() {
859
		let mut overlay = OverlayedChanges::default();
860
		overlay.set_storage(vec![20], None);
861
		overlay.set_storage(vec![30], Some(vec![31]));
862
		let backend = (
863
			Storage {
864
				top: map![
865
					vec![10] => vec![10],
866
					vec![20] => vec![20],
867
					vec![40] => vec![40]
868
				],
869
				children_default: map![],
870
			},
871
			StateVersion::default(),
872
		)
873
			.into();
874

            
875
		let mut ext = TestExt::new(&mut overlay, &backend, None);
876

            
877
		// next_backend < next_overlay
878
		assert_eq!(ext.next_storage_key(&[5]), Some(vec![10]));
879

            
880
		// next_backend == next_overlay but next_overlay is a delete
881
		assert_eq!(ext.next_storage_key(&[10]), Some(vec![30]));
882

            
883
		// next_overlay < next_backend
884
		assert_eq!(ext.next_storage_key(&[20]), Some(vec![30]));
885

            
886
		// next_backend exist but next_overlay doesn't exist
887
		assert_eq!(ext.next_storage_key(&[30]), Some(vec![40]));
888

            
889
		drop(ext);
890
		overlay.set_storage(vec![50], Some(vec![50]));
891
		let mut ext = TestExt::new(&mut overlay, &backend, None);
892

            
893
		// next_overlay exist but next_backend doesn't exist
894
		assert_eq!(ext.next_storage_key(&[40]), Some(vec![50]));
895
	}
896

            
897
	#[test]
898
	fn next_storage_key_works_with_a_lot_empty_values_in_overlay() {
899
		let mut overlay = OverlayedChanges::default();
900
		overlay.set_storage(vec![20], None);
901
		overlay.set_storage(vec![21], None);
902
		overlay.set_storage(vec![22], None);
903
		overlay.set_storage(vec![23], None);
904
		overlay.set_storage(vec![24], None);
905
		overlay.set_storage(vec![25], None);
906
		overlay.set_storage(vec![26], None);
907
		overlay.set_storage(vec![27], None);
908
		overlay.set_storage(vec![28], None);
909
		overlay.set_storage(vec![29], None);
910
		let backend = (
911
			Storage {
912
				top: map![
913
					vec![30] => vec![30]
914
				],
915
				children_default: map![],
916
			},
917
			StateVersion::default(),
918
		)
919
			.into();
920

            
921
		let mut ext = TestExt::new(&mut overlay, &backend, None);
922

            
923
		assert_eq!(ext.next_storage_key(&[5]), Some(vec![30]));
924

            
925
		drop(ext);
926
	}
927

            
928
	#[test]
929
	fn next_child_storage_key_works() {
930
		let child_info = ChildInfo::new_default(b"Child1");
931
		let child_info = &child_info;
932

            
933
		let mut overlay = OverlayedChanges::default();
934
		overlay.set_child_storage(child_info, vec![20], None);
935
		overlay.set_child_storage(child_info, vec![30], Some(vec![31]));
936
		let backend = (
937
			Storage {
938
				top: map![],
939
				children_default: map![
940
					child_info.storage_key().to_vec() => StorageChild {
941
						data: map![
942
							vec![10] => vec![10],
943
							vec![20] => vec![20],
944
							vec![40] => vec![40]
945
						],
946
						child_info: child_info.to_owned(),
947
					}
948
				],
949
			},
950
			StateVersion::default(),
951
		)
952
			.into();
953

            
954
		let mut ext = TestExt::new(&mut overlay, &backend, None);
955

            
956
		// next_backend < next_overlay
957
		assert_eq!(ext.next_child_storage_key(child_info, &[5]), Some(vec![10]));
958

            
959
		// next_backend == next_overlay but next_overlay is a delete
960
		assert_eq!(ext.next_child_storage_key(child_info, &[10]), Some(vec![30]));
961

            
962
		// next_overlay < next_backend
963
		assert_eq!(ext.next_child_storage_key(child_info, &[20]), Some(vec![30]));
964

            
965
		// next_backend exist but next_overlay doesn't exist
966
		assert_eq!(ext.next_child_storage_key(child_info, &[30]), Some(vec![40]));
967

            
968
		drop(ext);
969
		overlay.set_child_storage(child_info, vec![50], Some(vec![50]));
970
		let mut ext = TestExt::new(&mut overlay, &backend, None);
971

            
972
		// next_overlay exist but next_backend doesn't exist
973
		assert_eq!(ext.next_child_storage_key(child_info, &[40]), Some(vec![50]));
974
	}
975

            
976
	#[test]
977
	fn child_storage_works() {
978
		let child_info = ChildInfo::new_default(b"Child1");
979
		let child_info = &child_info;
980
		let mut overlay = OverlayedChanges::default();
981
		overlay.set_child_storage(child_info, vec![20], None);
982
		overlay.set_child_storage(child_info, vec![30], Some(vec![31]));
983
		let backend = (
984
			Storage {
985
				top: map![],
986
				children_default: map![
987
					child_info.storage_key().to_vec() => StorageChild {
988
						data: map![
989
							vec![10] => vec![10],
990
							vec![20] => vec![20],
991
							vec![30] => vec![40]
992
						],
993
						child_info: child_info.to_owned(),
994
					}
995
				],
996
			},
997
			StateVersion::default(),
998
		)
999
			.into();
		let mut ext = TestExt::new(&mut overlay, &backend, None);
		assert_eq!(ext.child_storage(child_info, &[10]), Some(vec![10]));
		assert_eq!(
			ext.child_storage_hash(child_info, &[10]),
			Some(Blake2Hasher::hash(&[10]).as_ref().to_vec()),
		);
		assert_eq!(ext.child_storage(child_info, &[20]), None);
		assert_eq!(ext.child_storage_hash(child_info, &[20]), None);
		assert_eq!(ext.child_storage(child_info, &[30]), Some(vec![31]));
		assert_eq!(
			ext.child_storage_hash(child_info, &[30]),
			Some(Blake2Hasher::hash(&[31]).as_ref().to_vec()),
		);
	}
	#[test]
	fn clear_prefix_cannot_delete_a_child_root() {
		let child_info = ChildInfo::new_default(b"Child1");
		let child_info = &child_info;
		let mut overlay = OverlayedChanges::default();
		let backend = (
			Storage {
				top: map![],
				children_default: map![
					child_info.storage_key().to_vec() => StorageChild {
						data: map![
							vec![30] => vec![40]
						],
						child_info: child_info.to_owned(),
					}
				],
			},
			StateVersion::default(),
		)
			.into();
		let ext = TestExt::new(&mut overlay, &backend, None);
		use sp_core::storage::well_known_keys;
		let mut ext = ext;
		let mut not_under_prefix = well_known_keys::CHILD_STORAGE_KEY_PREFIX.to_vec();
		not_under_prefix[4] = 88;
		not_under_prefix.extend(b"path");
		ext.set_storage(not_under_prefix.clone(), vec![10]);
		let _ = ext.clear_prefix(&[], None, None);
		let _ = ext.clear_prefix(&well_known_keys::CHILD_STORAGE_KEY_PREFIX[..4], None, None);
		let mut under_prefix = well_known_keys::CHILD_STORAGE_KEY_PREFIX.to_vec();
		under_prefix.extend(b"path");
		let _ = ext.clear_prefix(&well_known_keys::CHILD_STORAGE_KEY_PREFIX[..4], None, None);
		assert_eq!(ext.child_storage(child_info, &[30]), Some(vec![40]));
		assert_eq!(ext.storage(not_under_prefix.as_slice()), Some(vec![10]));
		let _ = ext.clear_prefix(&not_under_prefix[..5], None, None);
		assert_eq!(ext.storage(not_under_prefix.as_slice()), None);
	}
	#[test]
	fn storage_append_works() {
		let mut data = Vec::new();
		let mut append = StorageAppend::new(&mut data);
		append.append(1u32.encode());
		append.append(2u32.encode());
		drop(append);
		assert_eq!(Vec::<u32>::decode(&mut &data[..]).unwrap(), vec![1, 2]);
		// Initialize with some invalid data
		let mut data = vec![1];
		let mut append = StorageAppend::new(&mut data);
		append.append(1u32.encode());
		append.append(2u32.encode());
		drop(append);
		assert_eq!(Vec::<u32>::decode(&mut &data[..]).unwrap(), vec![1, 2]);
	}
}