1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
// This file is part of Substrate.

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

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

//! Provides functionality around the transaction storage.
//!
//! Transactional storage provides functionality to run an entire code block
//! in a storage transaction. This means that either the entire changes to the
//! storage are committed or everything is thrown away. This simplifies the
//! writing of functionality that may bail at any point of operation. Otherwise
//! you would need to first verify all storage accesses and then do the storage
//! modifications.
//!
//! [`with_transaction`] provides a way to run a given closure in a transactional context.

use sp_io::storage::{commit_transaction, rollback_transaction, start_transaction};
use sp_runtime::{DispatchError, TransactionOutcome, TransactionalError};

/// The type that is being used to store the current number of active layers.
pub type Layer = u32;
/// The key that is holds the current number of active layers.
///
/// Encodes to `0x3a7472616e73616374696f6e5f6c6576656c3a`.
pub const TRANSACTION_LEVEL_KEY: &[u8] = b":transaction_level:";
/// The maximum number of nested layers.
pub const TRANSACTIONAL_LIMIT: Layer = 255;

/// Returns the current number of nested transactional layers.
fn get_transaction_level() -> Layer {
	crate::storage::unhashed::get_or_default::<Layer>(TRANSACTION_LEVEL_KEY)
}

/// Set the current number of nested transactional layers.
fn set_transaction_level(level: Layer) {
	crate::storage::unhashed::put::<Layer>(TRANSACTION_LEVEL_KEY, &level);
}

/// Kill the transactional layers storage.
fn kill_transaction_level() {
	crate::storage::unhashed::kill(TRANSACTION_LEVEL_KEY);
}

/// Increments the transaction level. Returns an error if levels go past the limit.
///
/// Returns a guard that when dropped decrements the transaction level automatically.
fn inc_transaction_level() -> Result<StorageLayerGuard, ()> {
	let existing_levels = get_transaction_level();
	if existing_levels >= TRANSACTIONAL_LIMIT {
		return Err(())
	}
	// Cannot overflow because of check above.
	set_transaction_level(existing_levels + 1);
	Ok(StorageLayerGuard)
}

fn dec_transaction_level() {
	let existing_levels = get_transaction_level();
	if existing_levels == 0 {
		log::warn!(
				"We are underflowing with calculating transactional levels. Not great, but let's not panic...",
			);
	} else if existing_levels == 1 {
		// Don't leave any trace of this storage item.
		kill_transaction_level();
	} else {
		// Cannot underflow because of checks above.
		set_transaction_level(existing_levels - 1);
	}
}

struct StorageLayerGuard;

impl Drop for StorageLayerGuard {
	fn drop(&mut self) {
		dec_transaction_level()
	}
}

/// Check if the current call is within a transactional layer.
pub fn is_transactional() -> bool {
	get_transaction_level() > 0
}

/// Execute the supplied function in a new storage transaction.
///
/// All changes to storage performed by the supplied function are discarded if the returned
/// outcome is `TransactionOutcome::Rollback`.
///
/// Transactions can be nested up to `TRANSACTIONAL_LIMIT` times; more than that will result in an
/// error.
///
/// Commits happen to the parent transaction.
pub fn with_transaction<T, E, F>(f: F) -> Result<T, E>
where
	E: From<DispatchError>,
	F: FnOnce() -> TransactionOutcome<Result<T, E>>,
{
	// This needs to happen before `start_transaction` below.
	// Otherwise we may rollback the increase, then decrease as the guard goes out of scope
	// and then end in some bad state.
	let _guard = inc_transaction_level().map_err(|()| TransactionalError::LimitReached.into())?;

	start_transaction();

	match f() {
		TransactionOutcome::Commit(res) => {
			commit_transaction();
			res
		},
		TransactionOutcome::Rollback(res) => {
			rollback_transaction();
			res
		},
	}
}

/// Same as [`with_transaction`] but casts any internal error to `()`.
///
/// This rids `E` of the `From<DispatchError>` bound that is required by `with_transaction`.
pub fn with_transaction_opaque_err<T, E, F>(f: F) -> Result<Result<T, E>, ()>
where
	F: FnOnce() -> TransactionOutcome<Result<T, E>>,
{
	with_transaction(move || -> TransactionOutcome<Result<Result<T, E>, DispatchError>> {
		match f() {
			TransactionOutcome::Commit(res) => TransactionOutcome::Commit(Ok(res)),
			TransactionOutcome::Rollback(res) => TransactionOutcome::Rollback(Ok(res)),
		}
	})
	.map_err(|_| ())
}

/// Same as [`with_transaction`] but without a limit check on nested transactional layers.
///
/// This is mostly for backwards compatibility before there was a transactional layer limit.
/// It is recommended to only use [`with_transaction`] to avoid users from generating too many
/// transactional layers.
pub fn with_transaction_unchecked<R, F>(f: F) -> R
where
	F: FnOnce() -> TransactionOutcome<R>,
{
	// This needs to happen before `start_transaction` below.
	// Otherwise we may rollback the increase, then decrease as the guard goes out of scope
	// and then end in some bad state.
	let maybe_guard = inc_transaction_level();

	if maybe_guard.is_err() {
		log::warn!(
			"The transactional layer limit has been reached, and new transactional layers are being
			spawned with `with_transaction_unchecked`. This could be caused by someone trying to
			attack your chain, and you should investigate usage of `with_transaction_unchecked` and
			potentially migrate to `with_transaction`, which enforces a transactional limit.",
		);
	}

	start_transaction();

	match f() {
		TransactionOutcome::Commit(res) => {
			commit_transaction();
			res
		},
		TransactionOutcome::Rollback(res) => {
			rollback_transaction();
			res
		},
	}
}

/// Execute the supplied function, adding a new storage layer.
///
/// This is the same as `with_transaction`, but assuming that any function returning an `Err` should
/// rollback, and any function returning `Ok` should commit. This provides a cleaner API to the
/// developer who wants this behavior.
pub fn with_storage_layer<T, E, F>(f: F) -> Result<T, E>
where
	E: From<DispatchError>,
	F: FnOnce() -> Result<T, E>,
{
	with_transaction(|| {
		let r = f();
		if r.is_ok() {
			TransactionOutcome::Commit(r)
		} else {
			TransactionOutcome::Rollback(r)
		}
	})
}

/// Execute the supplied function, ensuring we are at least in one storage layer.
///
/// If we are already in a storage layer, we just execute the provided closure.
/// If we are not, we execute the closure within a [`with_storage_layer`].
pub fn in_storage_layer<T, E, F>(f: F) -> Result<T, E>
where
	E: From<DispatchError>,
	F: FnOnce() -> Result<T, E>,
{
	if is_transactional() {
		f()
	} else {
		with_storage_layer(f)
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use crate::{assert_noop, assert_ok};
	use sp_io::TestExternalities;
	use sp_runtime::DispatchResult;

	#[test]
	fn is_transactional_should_return_false() {
		TestExternalities::default().execute_with(|| {
			assert!(!is_transactional());
		});
	}

	#[test]
	fn is_transactional_should_not_error_in_with_transaction() {
		TestExternalities::default().execute_with(|| {
			assert_ok!(with_transaction(|| -> TransactionOutcome<DispatchResult> {
				assert!(is_transactional());
				TransactionOutcome::Commit(Ok(()))
			}));

			assert_noop!(
				with_transaction(|| -> TransactionOutcome<DispatchResult> {
					assert!(is_transactional());
					TransactionOutcome::Rollback(Err("revert".into()))
				}),
				"revert"
			);
		});
	}

	fn recursive_transactional(num: u32) -> DispatchResult {
		if num == 0 {
			return Ok(())
		}

		with_transaction(|| -> TransactionOutcome<DispatchResult> {
			let res = recursive_transactional(num - 1);
			TransactionOutcome::Commit(res)
		})
	}

	#[test]
	fn transaction_limit_should_work() {
		TestExternalities::default().execute_with(|| {
			assert_eq!(get_transaction_level(), 0);

			assert_ok!(with_transaction(|| -> TransactionOutcome<DispatchResult> {
				assert_eq!(get_transaction_level(), 1);
				TransactionOutcome::Commit(Ok(()))
			}));

			assert_ok!(with_transaction(|| -> TransactionOutcome<DispatchResult> {
				assert_eq!(get_transaction_level(), 1);
				let res = with_transaction(|| -> TransactionOutcome<DispatchResult> {
					assert_eq!(get_transaction_level(), 2);
					TransactionOutcome::Commit(Ok(()))
				});
				TransactionOutcome::Commit(res)
			}));

			assert_ok!(recursive_transactional(255));
			assert_noop!(
				recursive_transactional(256),
				sp_runtime::TransactionalError::LimitReached
			);

			assert_eq!(get_transaction_level(), 0);
		});
	}

	#[test]
	fn in_storage_layer_works() {
		TestExternalities::default().execute_with(|| {
			assert_eq!(get_transaction_level(), 0);

			let res = in_storage_layer(|| -> DispatchResult {
				assert_eq!(get_transaction_level(), 1);
				in_storage_layer(|| -> DispatchResult {
					// We are still in the same layer :)
					assert_eq!(get_transaction_level(), 1);
					Ok(())
				})
			});

			assert_ok!(res);

			let res = in_storage_layer(|| -> DispatchResult {
				assert_eq!(get_transaction_level(), 1);
				in_storage_layer(|| -> DispatchResult {
					// We are still in the same layer :)
					assert_eq!(get_transaction_level(), 1);
					Err("epic fail".into())
				})
			});

			assert_noop!(res, "epic fail");
		});
	}
}