hard limit on write_op and event sizes

jowenshaw · Oct 5, 2022 · 155f966 · 155f966
1 parent 4d9a6dc
commit 155f966
Show file tree

Hide file tree

Showing 8 changed files with 97 additions and 61 deletions.
diff --git a/aptos-move/aptos-aggregator/src/transaction.rs b/aptos-move/aptos-aggregator/src/transaction.rs
@@ -96,7 +96,7 @@ impl ChangeSetExt {
 
         Ok(Self {
             delta_change_set: delta_set,
-            change_set: ChangeSet::new(write_set.freeze()?, events),
+            change_set: ChangeSet::new(write_set.freeze()?, events)?,
         })
     }
 
@@ -140,7 +140,7 @@ impl ChangeSetExt {
 
         Ok(Self {
             delta_change_set: delta,
-            change_set: ChangeSet::new(write_set.freeze()?, events),
+            change_set: ChangeSet::new(write_set.freeze()?, events)?,
         })
     }
 

diff --git a/aptos-move/aptos-vm/src/move_vm_ext/session.rs b/aptos-move/aptos-vm/src/move_vm_ext/session.rs
@@ -249,7 +249,7 @@ impl SessionOutput {
             })
             .collect::<Result<Vec<_>, VMStatus>>()?;
 
-        let change_set = ChangeSet::new(write_set, events);
+        let change_set = ChangeSet::new(write_set, events)?;
         Ok(ChangeSetExt::new(delta_change_set, change_set))
     }
 

diff --git a/config/src/config/execution_config.rs b/config/src/config/execution_config.rs
@@ -134,7 +134,7 @@ mod test {
     #[test]
     fn test_some_and_load_genesis() {
         let fake_genesis = Transaction::GenesisTransaction(WriteSetPayload::Direct(
-            ChangeSet::new(WriteSetMut::new(vec![]).freeze().unwrap(), vec![]),
+            ChangeSet::new(WriteSetMut::new(vec![]).freeze().unwrap(), vec![]).unwrap(),
         ));
         let (mut config, path) = generate_config();
         config.genesis = Some(fake_genesis.clone());

diff --git a/execution/executor/src/mock_vm/mod.rs b/execution/executor/src/mock_vm/mod.rs
@@ -351,10 +351,9 @@ fn encode_transaction(sender: AccountAddress, program: Script) -> Transaction {
 }
 
 pub fn encode_reconfiguration_transaction() -> Transaction {
-    Transaction::GenesisTransaction(WriteSetPayload::Direct(ChangeSet::new(
-        WriteSet::default(),
-        vec![],
-    )))
+    Transaction::GenesisTransaction(WriteSetPayload::Direct(
+        ChangeSet::new(WriteSet::default(), vec![]).unwrap(),
+    ))
 }
 
 fn decode_transaction(txn: &SignedTransaction) -> MockVMTransaction {

diff --git a/execution/executor/tests/db_bootstrapper_test.rs b/execution/executor/tests/db_bootstrapper_test.rs
@@ -215,52 +215,57 @@ fn test_new_genesis() {
 
     // New genesis transaction: set validator set, bump epoch and overwrite account1 balance.
     let configuration = get_configuration(&db);
-    let genesis_txn = Transaction::GenesisTransaction(WriteSetPayload::Direct(ChangeSet::new(
-        WriteSetMut::new(vec![
-            (
-                StateKey::AccessPath(access_path_for_config(ValidatorSet::CONFIG_ID)),
-                WriteOp::Modification(bcs::to_bytes(&ValidatorSet::new(vec![])).unwrap()),
-            ),
-            (
-                StateKey::AccessPath(AccessPath::new(
-                    CORE_CODE_ADDRESS,
-                    ConfigurationResource::resource_path(),
-                )),
-                WriteOp::Modification(bcs::to_bytes(&configuration.bump_epoch_for_test()).unwrap()),
-            ),
-            (
-                StateKey::AccessPath(AccessPath::new(
-                    account1,
-                    CoinStoreResource::resource_path(),
-                )),
-                WriteOp::Modification(
-                    bcs::to_bytes(&CoinStoreResource::new(
-                        1_000_000,
-                        false,
-                        EventHandle::random(0),
-                        EventHandle::random(0),
-                    ))
-                    .unwrap(),
+    let genesis_txn = Transaction::GenesisTransaction(WriteSetPayload::Direct(
+        ChangeSet::new(
+            WriteSetMut::new(vec![
+                (
+                    StateKey::AccessPath(access_path_for_config(ValidatorSet::CONFIG_ID)),
+                    WriteOp::Modification(bcs::to_bytes(&ValidatorSet::new(vec![])).unwrap()),
                 ),
-            ),
-        ])
-        .freeze()
+                (
+                    StateKey::AccessPath(AccessPath::new(
+                        CORE_CODE_ADDRESS,
+                        ConfigurationResource::resource_path(),
+                    )),
+                    WriteOp::Modification(
+                        bcs::to_bytes(&configuration.bump_epoch_for_test()).unwrap(),
+                    ),
+                ),
+                (
+                    StateKey::AccessPath(AccessPath::new(
+                        account1,
+                        CoinStoreResource::resource_path(),
+                    )),
+                    WriteOp::Modification(
+                        bcs::to_bytes(&CoinStoreResource::new(
+                            1_000_000,
+                            false,
+                            EventHandle::random(0),
+                            EventHandle::random(0),
+                        ))
+                        .unwrap(),
+                    ),
+                ),
+            ])
+            .freeze()
+            .unwrap(),
+            vec![
+                ContractEvent::new(
+                    *configuration.events().key(),
+                    0,
+                    TypeTag::Struct(Box::new(ConfigurationResource::struct_tag())),
+                    vec![],
+                ),
+                ContractEvent::new(
+                    new_block_event_key(),
+                    0,
+                    TypeTag::Struct(Box::new(NewBlockEvent::struct_tag())),
+                    vec![],
+                ),
+            ],
+        )
         .unwrap(),
-        vec![
-            ContractEvent::new(
-                *configuration.events().key(),
-                0,
-                TypeTag::Struct(Box::new(ConfigurationResource::struct_tag())),
-                vec![],
-            ),
-            ContractEvent::new(
-                new_block_event_key(),
-                0,
-                TypeTag::Struct(Box::new(NewBlockEvent::struct_tag())),
-                vec![],
-            ),
-        ],
-    )));
+    ));
 
     // Bootstrap DB into new genesis.
     let waypoint = generate_waypoint::<AptosVM>(&db, &genesis_txn).unwrap();

diff --git a/state-sync/inter-component/mempool-notifications/src/lib.rs b/state-sync/inter-component/mempool-notifications/src/lib.rs
@@ -382,12 +382,15 @@ mod tests {
     }
 
     fn create_genesis_transaction() -> Transaction {
-        Transaction::GenesisTransaction(WriteSetPayload::Direct(ChangeSet::new(
-            WriteSetMut::new(vec![])
-                .freeze()
-                .expect("freeze cannot fail"),
-            vec![],
-        )))
+        Transaction::GenesisTransaction(WriteSetPayload::Direct(
+            ChangeSet::new(
+                WriteSetMut::new(vec![])
+                    .freeze()
+                    .expect("freeze cannot fail"),
+                vec![],
+            )
+            .unwrap(),
+        ))
     }
 
     fn create_runtime() -> Runtime {

diff --git a/types/src/proptest_types.rs b/types/src/proptest_types.rs
@@ -104,7 +104,7 @@ impl Arbitrary for ChangeSet {
     type Parameters = ();
     fn arbitrary_with(_args: ()) -> Self::Strategy {
         (any::<WriteSet>(), vec(any::<ContractEvent>(), 0..10))
-            .prop_map(|(ws, events)| ChangeSet::new(ws, events))
+            .prop_map(|(ws, events)| ChangeSet::new(ws, events).unwrap())
             .boxed()
     }
 

diff --git a/types/src/transaction/change_set.rs b/types/src/transaction/change_set.rs
@@ -1,7 +1,9 @@
 // Copyright (c) Aptos
 // SPDX-License-Identifier: Apache-2.0
 
+use crate::write_set::WriteOp;
 use crate::{contract_event::ContractEvent, write_set::WriteSet};
+use move_core_types::vm_status::{StatusCode, VMStatus};
 use serde::{Deserialize, Serialize};
 
 #[derive(Clone, Debug, Hash, Eq, PartialEq, Serialize, Deserialize)]
@@ -11,8 +13,35 @@ pub struct ChangeSet {
 }
 
 impl ChangeSet {
-    pub fn new(write_set: WriteSet, events: Vec<ContractEvent>) -> Self {
-        Self { write_set, events }
+    pub fn new(write_set: WriteSet, events: Vec<ContractEvent>) -> Result<Self, VMStatus> {
+        static MAX_ITEM_SIZE_ALLOWED: usize = 1 << 20;
+        static MAX_EVENT_SIZE_ALLOWED: usize = 1 << 20;
+        static MAX_TOTAL_EVENT_SIZE_ALLOWED: usize = 10 << 20;
+
+        for (key, op) in &write_set {
+            match op {
+                WriteOp::Creation(data) | WriteOp::Modification(data) => {
+                    if data.len() + key.size() > MAX_ITEM_SIZE_ALLOWED {
+                        return Err(VMStatus::Error(StatusCode::STORAGE_WRITE_LIMIT_REACHED));
+                    }
+                }
+                WriteOp::Deletion => (),
+            }
+        }
+
+        let mut total_event_size = 0;
+        for event in &events {
+            let size = event.event_data().len();
+            if size > MAX_EVENT_SIZE_ALLOWED {
+                return Err(VMStatus::Error(StatusCode::STORAGE_WRITE_LIMIT_REACHED));
+            }
+            total_event_size += size;
+            if total_event_size > MAX_TOTAL_EVENT_SIZE_ALLOWED {
+                return Err(VMStatus::Error(StatusCode::STORAGE_WRITE_LIMIT_REACHED));
+            }
+        }
+
+        Ok(Self { write_set, events })
     }
 
     pub fn into_inner(self) -> (WriteSet, Vec<ContractEvent>) {