Fail when decoding from storage and not all bytes consumed

crates/engine/src/ext.rs

@@ -241,32 +241,26 @@ impl Engine {
     }
 
     /// Returns the decoded contract storage at the key if any.
-    pub fn get_storage(&mut self, key: &[u8], output: &mut &mut [u8]) -> Result {
+    pub fn get_storage(&mut self, key: &[u8]) -> core::result::Result<&[u8], Error> {
         let callee = self.get_callee();
         let account_id = AccountId::from_bytes(&callee[..]);
 
         self.debug_info.inc_reads(account_id);
         match self.database.get_from_contract_storage(&callee, key) {
-            Some(val) => {
-                set_output(output, val);
-                Ok(())
-            }
+            Some(val) => Ok(val),
             None => Err(Error::KeyNotFound),
         }
     }
 
     /// Removes the storage entries at the given key,
     /// returning previously stored value at the key if any.
-    pub fn take_storage(&mut self, key: &[u8], output: &mut &mut [u8]) -> Result {
+    pub fn take_storage(&mut self, key: &[u8]) -> core::result::Result<Vec<u8>, Error> {
         let callee = self.get_callee();
         let account_id = AccountId::from_bytes(&callee[..]);
 
         self.debug_info.inc_writes(account_id);
         match self.database.remove_contract_storage(&callee, key) {
-            Some(val) => {
-                set_output(output, &val);
-                Ok(())
-            }
+            Some(val) => Ok(val),
             None => Err(Error::KeyNotFound),
         }
     }

crates/engine/src/test_api.rs

@@ -334,17 +334,16 @@ mod tests {
         // given
         let mut engine = Engine::new();
         let key: &[u8; 32] = &[0x42; 32];
-        let mut buf = [0_u8; 32];
 
         // when
         engine.set_callee(vec![1; 32]);
         engine.set_storage(key, &[0x05_u8; 5]);
         engine.set_storage(key, &[0x05_u8; 6]);
-        engine.get_storage(key, &mut &mut buf[..]).unwrap();
+        engine.get_storage(key).unwrap();
 
         engine.set_callee(vec![2; 32]);
         engine.set_storage(key, &[0x07_u8; 7]);
-        engine.get_storage(key, &mut &mut buf[..]).unwrap();
+        engine.get_storage(key).unwrap();
 
         // then
         assert_eq!(engine.count_writes(), 3);

crates/engine/src/tests.rs

@@ -40,17 +40,16 @@ fn store_load_clear() {
     let mut engine = Engine::new();
     engine.set_callee(vec![1; 32]);
     let key: &[u8; 32] = &[0x42; 32];
-    let output = &mut &mut get_buffer()[..];
-    let res = engine.get_storage(key, output);
+    let res = engine.get_storage(key);
     assert_eq!(res, Err(Error::KeyNotFound));
 
     engine.set_storage(key, &[0x05_u8; 5]);
-    let res = engine.get_storage(key, output);
-    assert_eq!(res, Ok(()),);
-    assert_eq!(output[..5], [0x05; 5]);
+    let res = engine.get_storage(key);
+    assert!(res.is_ok());
+    assert_eq!(res.unwrap()[..5], [0x05; 5]);
 
     engine.clear_storage(key);
-    let res = engine.get_storage(key, output);
+    let res = engine.get_storage(key);
     assert_eq!(res, Err(Error::KeyNotFound));
 }
 
@@ -180,26 +179,6 @@ fn value_transferred() {
     assert_eq!(output, value);
 }
 
-#[test]
-#[should_panic(
-    expected = "the output buffer is too small! the decoded storage is of size 16 bytes, but the output buffer has only room for 8."
-)]
-fn must_panic_when_buffer_too_small() {
-    // given
-    let mut engine = Engine::new();
-    engine.set_callee(vec![1; 32]);
-    let key: &[u8; 32] = &[0x42; 32];
-    engine.set_storage(key, &[0x05_u8; 16]);
-
-    // when
-    let mut small_buffer = [0; 8];
-    let output = &mut &mut small_buffer[..];
-    let _ = engine.get_storage(key, output);
-
-    // then
-    unreachable!("`get_storage` must already have panicked");
-}
-
 #[test]
 fn ecdsa_recovery_test_from_contracts_pallet() {
     // given

crates/env/src/engine/off_chain/impls.rs

@@ -46,7 +46,10 @@ use ink_engine::{
     ext,
     ext::Engine,
 };
-use ink_storage_traits::Storable;
+use ink_storage_traits::{
+    decode_all,
+    Storable,
+};
 use schnorrkel::{
     PublicKey,
     Signature,
@@ -202,32 +205,29 @@ impl EnvBackend for EnvInstance {
         K: scale::Encode,
         R: Storable,
     {
-        let mut output: [u8; 9600] = [0; 9600];
-        match self.engine.get_storage(&key.encode(), &mut &mut output[..]) {
-            Ok(_) => (),
-            Err(ext::Error::KeyNotFound) => return Ok(None),
+        match self.engine.get_storage(&key.encode()) {
+            Ok(res) => {
+                let decoded = decode_all(&mut &res[..])?;
+                Ok(Some(decoded))
+            }
+            Err(ext::Error::KeyNotFound) => Ok(None),
             Err(_) => panic!("encountered unexpected error"),
         }
-        let decoded = Storable::decode(&mut &output[..])?;
-        Ok(Some(decoded))
     }
 
     fn take_contract_storage<K, R>(&mut self, key: &K) -> Result<Option<R>>
     where
         K: scale::Encode,
         R: Storable,
     {
-        let mut output: [u8; 9600] = [0; 9600];
-        match self
-            .engine
-            .take_storage(&key.encode(), &mut &mut output[..])
-        {
-            Ok(_) => (),
-            Err(ext::Error::KeyNotFound) => return Ok(None),
+        match self.engine.take_storage(&key.encode()) {
+            Ok(output) => {
+                let decoded = decode_all(&mut &output[..])?;
+                Ok(Some(decoded))
+            }
+            Err(ext::Error::KeyNotFound) => Ok(None),
             Err(_) => panic!("encountered unexpected error"),
         }
-        let decoded = Storable::decode(&mut &output[..])?;
-        Ok(Some(decoded))
     }
 
     fn contains_contract_storage<K>(&mut self, key: &K) -> Option<u32>

crates/env/src/engine/on_chain/impls.rs

@@ -48,7 +48,10 @@ use crate::{
     ReturnFlags,
     TypedEnvBackend,
 };
-use ink_storage_traits::Storable;
+use ink_storage_traits::{
+    decode_all,
+    Storable,
+};
 
 impl CryptoHash for Blake2x128 {
     fn hash(input: &[u8], output: &mut <Self as HashOutput>::Type) {
@@ -236,7 +239,7 @@ impl EnvBackend for EnvInstance {
             Err(ExtError::KeyNotFound) => return Ok(None),
             Err(_) => panic!("encountered unexpected error"),
         }
-        let decoded = Storable::decode(&mut &output[..])?;
+        let decoded = decode_all(&mut &output[..])?;
         Ok(Some(decoded))
     }
 
@@ -253,7 +256,7 @@ impl EnvBackend for EnvInstance {
             Err(ExtError::KeyNotFound) => return Ok(None),
             Err(_) => panic!("encountered unexpected error"),
         }
-        let decoded = Storable::decode(&mut &output[..])?;
+        let decoded = decode_all(&mut &output[..])?;
         Ok(Some(decoded))
     }
 

crates/storage/traits/src/lib.rs

@@ -47,6 +47,7 @@ pub use self::{
         ResolverKey,
     },
     storage::{
+        decode_all,
         AutoStorableHint,
         Packed,
         Storable,

crates/storage/traits/src/storage.rs

@@ -44,6 +44,19 @@ where
     }
 }
 
+/// Decode and consume all of the given input data.
+///
+/// If not all data is consumed, an error is returned.
+pub fn decode_all<T: Storable>(input: &mut &[u8]) -> Result<T, scale::Error> {
+    let res = <T as Storable>::decode(input)?;
+
+    if input.is_empty() {
+        Ok(res)
+    } else {
+        Err("Input buffer has still data left after decoding!".into())
+    }
+}
+
 pub(crate) mod private {
     /// Seals the implementation of `Packed`.
     pub trait Sealed {}

integration-tests/contract-storage/.gitignore

@@ -0,0 +1,9 @@
+# Ignore build artifacts from the local tests sub-crate.
+/target/
+
+# Ignore backup files creates by cargo fmt.
+**/*.rs.bk
+
+# Remove Cargo.lock when creating an executable, leave it for libraries
+# More information here http://doc.crates.io/guide.html#cargotoml-vs-cargolock
+Cargo.lock

integration-tests/contract-storage/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "contract-storage"
+version = "4.2.0"
+authors = ["Parity Technologies <admin@parity.io>"]
+edition = "2021"
+publish = false
+
+[dependencies]
+ink = { path = "../../crates/ink", default-features = false }
+
+[dev-dependencies]
+ink_e2e = { path = "../../crates/e2e" }
+
+[lib]
+path = "lib.rs"
+
+[features]
+default = ["std"]
+std = [
+    "ink/std",
+]
+ink-as-dependency = []
+e2e-tests = []

integration-tests/contract-storage/e2e_tests.rs

@@ -0,0 +1,122 @@
+use super::contract_storage::*;
+use ink_e2e::ContractsBackend;
+
+type E2EResult<T> = std::result::Result<T, Box<dyn std::error::Error>>;
+
+#[ink_e2e::test]
+async fn get_contract_storage_consumes_entire_buffer<Client: E2EBackend>(
+    mut client: Client,
+) -> E2EResult<()> {
+    // given
+    let mut constructor = ContractStorageRef::new();
+    let contract = client
+        .instantiate("contract-storage", &ink_e2e::alice(), &mut constructor)
+        .submit()
+        .await
+        .expect("instantiate failed");
+    let call = contract.call::<ContractStorage>();
+
+    // when
+    let result = client
+        .call(
+            &ink_e2e::alice(),
+            &call.set_and_get_storage_all_data_consumed(),
+        )
+        .submit()
+        .await
+        .expect("Calling `insert_balance` failed")
+        .return_value();
+
+    assert!(result.is_ok());
+
+    Ok(())
+}
+
+#[ink_e2e::test]
+async fn get_contract_storage_fails_when_extra_data<Client: E2EBackend>(
+    mut client: Client,
+) -> E2EResult<()> {
+    // given
+    let mut constructor = ContractStorageRef::new();
+    let contract = client
+        .instantiate("contract-storage", &ink_e2e::alice(), &mut constructor)
+        .submit()
+        .await
+        .expect("instantiate failed");
+    let call = contract.call::<ContractStorage>();
+
+    // when
+    let result = client
+        .call(
+            &ink_e2e::alice(),
+            &call.set_and_get_storage_partial_data_consumed(),
+        )
+        .submit()
+        .await;
+
+    assert!(
+        result.is_err(),
+        "Expected the contract to revert when only partially consuming the buffer"
+    );
+
+    Ok(())
+}
+
+#[ink_e2e::test]
+async fn take_contract_storage_consumes_entire_buffer<Client: E2EBackend>(
+    mut client: Client,
+) -> E2EResult<()> {
+    // given
+    let mut constructor = ContractStorageRef::new();
+    let contract = client
+        .instantiate("contract-storage", &ink_e2e::alice(), &mut constructor)
+        .submit()
+        .await
+        .expect("instantiate failed");
+    let call = contract.call::<ContractStorage>();
+
+    // when
+    let result = client
+        .call(
+            &ink_e2e::alice(),
+            &call.set_and_take_storage_all_data_consumed(),
+        )
+        .submit()
+        .await
+        .expect("Calling `insert_balance` failed")
+        .return_value();
+
+    assert!(result.is_ok());
+
+    Ok(())
+}
+
+#[ink_e2e::test]
+async fn take_contract_storage_fails_when_extra_data<Client: E2EBackend>(
+    mut client: Client,
+) -> E2EResult<()> {
+    // given
+    let mut constructor = ContractStorageRef::new();
+    let contract = client
+        .instantiate("contract-storage", &ink_e2e::alice(), &mut constructor)
+        .submit()
+        .await
+        .expect("instantiate failed");
+    let call = contract.call::<ContractStorage>();
+
+    // when
+    let result = client
+        .call(
+            &ink_e2e::alice(),
+            &call.set_and_take_storage_partial_data_consumed(),
+        )
+        .submit()
+        .await;
+
+    assert!(
+        result.is_err(),
+        "Expected the contract to revert when only partially consuming the buffer"
+    );
+
+    Ok(())
+}