Skip to content

Commit 7ccca1f

Browse files
0xVolosnikovantoniolocascio
0xVolosnikov
and
antoniolocascio
authored
feat: Improve validations in history map rollback (#44)
## What ❔ 1. Adds internal error instead of panic on attempt to rollback to invalid snapshot 2. Forbids rollbacks to non-existent snapshots 3. Clarifies `materialize_element` logic ## Is this a breaking change? - [ ] Yes - [X] No ## Checklist <!-- Check your PR fulfills the following items. --> <!-- For draft PRs check the boxes as you complete them. --> - [X] PR title corresponds to the body of PR (we generate changelog entries from PRs). - [ ] Tests for the changes have been added / updated. - [ ] Documentation comments have been added / updated. - [X] Code has been formatted. --------- Co-authored-by: Antonio Locascio <antonio.locascio1@gmail.com>
1 parent 72a770b commit 7ccca1f

File tree

10 files changed

+122
-47
lines changed

10 files changed

+122
-47
lines changed

‎basic_system/src/system_implementation/flat_storage_model/account_cache.rs‎

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@ where
8484
}
8585
}
8686

87+
/// Read element and initialize it if needed
8788
fn materialize_element<const PROOF_ENV: bool>(
8889
&mut self,
8990
ee_type: ExecutionEnvironmentType,
@@ -111,10 +112,13 @@ where
111112
let native = R::Native::from_computational(WARM_ACCOUNT_CACHE_ACCESS_NATIVE_COST);
112113
resources.charge(&R::from_ergs_and_native(ergs, native))?;
113114

114-
let mut cold_read_charged = false;
115+
let mut initialized_element = false;
115116

116117
self.cache
117118
.get_or_insert(address.into(), || {
119+
// Element doesn't exist in cache yet, initialize it
120+
initialized_element = true;
121+
118122
// - first get a hash of properties from storage
119123
match ee_type {
120124
ExecutionEnvironmentType::NoEE => {
@@ -141,8 +145,6 @@ where
141145
_ => return Err(InternalError("Unsupported EE").into()),
142146
}
143147

144-
cold_read_charged = true;
145-
146148
// to avoid divergence we read as-if infinite ergs
147149
let hash = resources.with_infinite_ergs(|inf_resources| {
148150
storage.read_special_account_property::<AccountAggregateDataHash>(
@@ -179,16 +181,19 @@ where
179181
}
180182
};
181183

184+
// Note: we initialize it as cold, should be warmed up separately
185+
// Since in case of revert it should become cold again and initial record can't be rolled back
182186
Ok(CacheRecord::new(acc_data.0, acc_data.1))
183187
})
184-
// We're adding a read snapshot for case when we're rollbacking the initial read.
185188
.and_then(|mut x| {
189+
// Warm up element according to EVM rules if needed
186190
let is_warm = x
187191
.current()
188192
.metadata()
189193
.considered_warm(self.current_tx_number);
190194
if is_warm == false {
191-
if cold_read_charged == false {
195+
if initialized_element == false {
196+
// Element exists in cache, but wasn't touched in current tx yet
192197
match ee_type {
193198
ExecutionEnvironmentType::NoEE => {
194199
// Access list accesses are always done in NoEE.
@@ -383,9 +388,15 @@ where
383388
self.cache.snapshot()
384389
}
385390

386-
pub fn finish_frame(&mut self, rollback_handle: Option<&CacheSnapshotId>) {
391+
#[must_use]
392+
pub fn finish_frame(
393+
&mut self,
394+
rollback_handle: Option<&CacheSnapshotId>,
395+
) -> Result<(), InternalError> {
387396
if let Some(x) = rollback_handle {
388-
self.cache.rollback(*x);
397+
self.cache.rollback(*x)
398+
} else {
399+
Ok(())
389400
}
390401
}
391402

‎basic_system/src/system_implementation/flat_storage_model/mod.rs‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -477,12 +477,17 @@ where
477477
}
478478
}
479479

480-
fn finish_frame(&mut self, rollback_handle: Option<&Self::StateSnapshot>) {
480+
fn finish_frame(
481+
&mut self,
482+
rollback_handle: Option<&Self::StateSnapshot>,
483+
) -> Result<(), InternalError> {
481484
self.storage_cache
482-
.finish_frame(rollback_handle.map(|x| &x.storage));
485+
.finish_frame(rollback_handle.map(|x| &x.storage))?;
483486
self.preimages_cache
484-
.finish_frame(rollback_handle.map(|x| &x.preimages));
487+
.finish_frame(rollback_handle.map(|x| &x.preimages))?;
485488
self.account_data_cache
486-
.finish_frame(rollback_handle.map(|x| &x.account_data));
489+
.finish_frame(rollback_handle.map(|x| &x.account_data))?;
490+
491+
Ok(())
487492
}
488493
}

‎basic_system/src/system_implementation/flat_storage_model/preimage_cache.rs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,10 @@ impl<R: Resources, A: Allocator + Clone> SnapshottableIo
278278
self.publication_storage.start_frame()
279279
}
280280

281-
fn finish_frame(&mut self, rollback_handle: Option<&Self::StateSnapshot>) {
282-
self.publication_storage.finish_frame(rollback_handle);
281+
fn finish_frame(
282+
&mut self,
283+
rollback_handle: Option<&Self::StateSnapshot>,
284+
) -> Result<(), InternalError> {
285+
self.publication_storage.finish_frame(rollback_handle)
283286
}
284287
}

‎basic_system/src/system_implementation/flat_storage_model/storage_cache.rs‎

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use storage_models::common_structs::{AccountAggregateDataHash, StorageCacheModel
1111
use zk_ee::common_structs::cache_record::{Appearance, CacheRecord};
1212
use zk_ee::common_traits::key_like_with_bounds::{KeyLikeWithBounds, TyEq};
1313
use zk_ee::execution_environment_type::ExecutionEnvironmentType;
14+
use zk_ee::system::errors::InternalError;
1415
use zk_ee::{
1516
common_structs::{WarmStorageKey, WarmStorageValue},
1617
kv_markers::{StorageAddress, UsizeDeserializable},
@@ -140,12 +141,19 @@ where
140141
}
141142

142143
#[track_caller]
143-
pub fn finish_frame_impl(&mut self, rollback_handle: Option<&CacheSnapshotId>) {
144+
#[must_use]
145+
pub fn finish_frame_impl(
146+
&mut self,
147+
rollback_handle: Option<&CacheSnapshotId>,
148+
) -> Result<(), InternalError> {
144149
if let Some(x) = rollback_handle {
145-
self.cache.rollback(*x);
150+
self.cache.rollback(*x)
151+
} else {
152+
Ok(())
146153
}
147154
}
148155

156+
/// Read element and initialize it if needed
149157
fn materialize_element<'a>(
150158
cache: &'a mut HistoryMap<K, CacheRecord<V, StorageElementMetadata>, A>,
151159
resources_policy: &mut P,
@@ -159,10 +167,13 @@ where
159167
) -> Result<(AddressItem<'a, K, V, A>, IsWarmRead), SystemError> {
160168
resources_policy.charge_warm_storage_read(ee_type, resources)?;
161169

162-
let mut cold_read_charged = false;
170+
let mut initialized_element = false;
163171

164172
cache
165-
.get_or_insert( key, || {
173+
.get_or_insert(key, || {
174+
// Element doesn't exist in cache yet, initialize it
175+
initialized_element = true;
176+
166177
let mut dst =
167178
core::mem::MaybeUninit::<InitialStorageSlotData<EthereumIOTypesConfig>>::uninit(
168179
);
@@ -179,21 +190,24 @@ where
179190
// correctly.
180191
let data_from_oracle = unsafe { dst.assume_init() } ;
181192

182-
resources_policy.charge_cold_storage_read_extra(ee_type, resources, data_from_oracle.is_new_storage_slot,is_access_list)?;
183-
cold_read_charged = true;
193+
resources_policy.charge_cold_storage_read_extra(ee_type, resources, data_from_oracle.is_new_storage_slot, is_access_list)?;
184194

185195
let appearance = match data_from_oracle.is_new_storage_slot {
186196
true => Appearance::Unset,
187197
false => Appearance::Retrieved,
188198
};
199+
200+
// Note: we initialize it as cold, should be warmed up separately
201+
// Since in case of revert it should become cold again and initial record can't be rolled back
189202
Ok(CacheRecord::new(data_from_oracle.initial_value.into(), appearance))
190203
})
191-
// We're adding a read snapshot for case when we're rollbacking the initial read.
192204
.and_then(|mut x| {
205+
// Warm up element according to EVM rules if needed
193206
let is_warm_read = x.current().metadata().considered_warm(current_tx_number);
194207
if is_warm_read == false {
195-
if cold_read_charged == false {
196-
resources_policy.charge_cold_storage_read_extra(ee_type, resources,false,is_access_list)?;
208+
if initialized_element == false {
209+
// Element exists in cache, but wasn't touched in current tx yet
210+
resources_policy.charge_cold_storage_read_extra(ee_type, resources,false, is_access_list)?;
197211
}
198212

199213
x.update(|cache_record| {
@@ -524,8 +538,11 @@ where
524538
self.0.start_frame()
525539
}
526540

527-
fn finish_frame(&mut self, rollback_handle: Option<&Self::StateSnapshot>) {
528-
self.0.finish_frame_impl(rollback_handle);
541+
fn finish_frame(
542+
&mut self,
543+
rollback_handle: Option<&Self::StateSnapshot>,
544+
) -> Result<(), InternalError> {
545+
self.0.finish_frame_impl(rollback_handle)
529546
}
530547
}
531548

‎basic_system/src/system_implementation/system/io_subsystem.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -378,9 +378,9 @@ where
378378
&mut self,
379379
rollback_handle: Option<&FullIOStateSnapshot>,
380380
) -> Result<(), InternalError> {
381-
self.storage.finish_frame(rollback_handle.map(|x| &x.io));
381+
self.storage.finish_frame(rollback_handle.map(|x| &x.io))?;
382382
self.transient_storage
383-
.finish_frame(rollback_handle.map(|x| &x.transient));
383+
.finish_frame(rollback_handle.map(|x| &x.transient))?;
384384
self.logs_storage
385385
.finish_frame(rollback_handle.map(|x| x.messages));
386386
self.events_storage

‎storage_models/src/common_structs/generic_transient_storage.rs‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use zk_ee::common_structs::history_map::HistoryMapItemRefMut;
44
use core::alloc::Allocator;
55
use core::marker::PhantomData;
66
use zk_ee::common_traits::key_like_with_bounds::KeyLikeWithBounds;
7-
use zk_ee::system::errors::SystemError;
7+
use zk_ee::system::errors::{InternalError, SystemError};
88
use zk_ee::{
99
common_structs::history_map::{CacheSnapshotId, HistoryMap},
1010
memory::stack_trait::{StackCtor, StackCtorConst},
@@ -49,8 +49,8 @@ where
4949

5050
pub fn begin_new_tx(&mut self) {
5151
// Just discard old history
52-
// Note: it will reset snapshots counter
53-
// Could be also done in separate finish_tx (to not reset for first tx), but it doesn't matter
52+
// Note: it will reset snapshots counter, old snapshots handlers can't be used anymore
53+
// Note: We will reset it redundantly for first tx
5454
self.cache = HistoryMap::new(self.alloc.clone());
5555
self.current_tx_number += 1;
5656
}
@@ -60,6 +60,7 @@ where
6060
self.cache.snapshot()
6161
}
6262

63+
/// Read element and initialize it if needed
6364
fn materialize_element<'a>(
6465
cache: &'a mut HistoryMap<K, V, A>,
6566
key: &'a K,
@@ -93,9 +94,14 @@ where
9394
}
9495

9596
#[track_caller]
96-
pub fn finish_frame(&mut self, rollback_handle: Option<&CacheSnapshotId>) {
97+
pub fn finish_frame(
98+
&mut self,
99+
rollback_handle: Option<&CacheSnapshotId>,
100+
) -> Result<(), InternalError> {
97101
if let Some(x) = rollback_handle {
98-
self.cache.rollback(*x);
102+
self.cache.rollback(*x)
103+
} else {
104+
Ok(())
99105
}
100106
}
101107
}

‎storage_models/src/common_structs/traits/snapshottable_io.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,14 @@
1+
use zk_ee::system::errors::InternalError;
2+
13
/// TODO
24
pub trait SnapshottableIo {
35
type StateSnapshot;
46

57
fn begin_new_tx(&mut self);
68

79
fn start_frame(&mut self) -> Self::StateSnapshot;
8-
fn finish_frame(&mut self, rollback_handle: Option<&Self::StateSnapshot>);
10+
fn finish_frame(
11+
&mut self,
12+
rollback_handle: Option<&Self::StateSnapshot>,
13+
) -> Result<(), InternalError>;
914
}

‎zk_ee/src/common_structs/history_map.rs‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -113,11 +113,17 @@ where
113113
snapshot_id
114114
}
115115

116+
#[must_use]
116117
/// Rollbacks the data to the state at the provided `snapshot_id`.
117-
pub fn rollback(&mut self, snapshot_id: CacheSnapshotId) {
118+
pub fn rollback(&mut self, snapshot_id: CacheSnapshotId) -> Result<(), InternalError> {
118119
if snapshot_id < self.state.frozen_snapshot_id {
119-
// TODO: replace with internal error
120-
panic!("Rolling below frozen snapshot is illegal and will cause UB.")
120+
return Err(InternalError("History map: rollback below frozen snapshot"));
121+
}
122+
123+
if snapshot_id >= self.state.next_snapshot_id {
124+
return Err(InternalError(
125+
"History map: rollback to non-existent snapshot",
126+
));
121127
}
122128

123129
// Go over all elements changed since last `commit` and roll them back
@@ -145,6 +151,8 @@ where
145151
}
146152
}
147153
}
154+
155+
Ok(())
148156
}
149157

150158
/// Commits (freezes) changes up to this point and frees memory taken by snapshots that can't be
@@ -723,7 +731,7 @@ mod tests {
723731

724732
map.snapshot();
725733

726-
map.rollback(ss);
734+
map.rollback(ss).expect("Correct snapshot");
727735

728736
map.for_total_diff_operands::<_, ()>(|l, r, k| {
729737
assert_eq!(1, *l);
@@ -764,7 +772,7 @@ mod tests {
764772
// Just for fun.
765773
map.snapshot();
766774

767-
map.rollback(ss);
775+
map.rollback(ss).expect("Correct snapshot");
768776

769777
let mut v = map.get_or_insert::<()>(&1, || Ok(5)).unwrap();
770778

0 commit comments

Comments
 (0)