Просмотр исходного кода

runtime: Add db.contains_key functionality.

parazyd 3 лет назад
Родитель
Сommit
9c0a0e1c10
4 измененных файлов с 111 добавлено и 0 удалено
  1. 74 0
      src/runtime/import/db.rs
  2. 6 0
      src/runtime/vm_runtime.rs
  3. 25 0
      src/sdk/src/db.rs
  4. 6 0
      src/sdk/src/error.rs

+ 74 - 0
src/runtime/import/db.rs

@@ -47,6 +47,10 @@ impl DbHandle {
         Ok(None)
     }
 
+    pub fn contains_key(&self, key: &[u8]) -> Result<bool> {
+        self.tree.contains_key(key)
+    }
+
     pub fn apply_batch(&self, batch: sled::Batch) -> Result<()> {
         Ok(self.tree.apply_batch(batch)?)
     }
@@ -344,3 +348,73 @@ pub(crate) fn db_get(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i6
         _ => -1,
     }
 }
+
+/// Everyone can call this. Will check if a given db contains given key.
+pub(crate) fn db_contains_key(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i32 {
+    let env = ctx.data();
+    match env.contract_section {
+        ContractSection::Deploy |
+        ContractSection::Exec |
+        ContractSection::Update |
+        ContractSection::Metadata => {
+            let memory_view = env.memory_view(&ctx);
+
+            let Ok(mem_slice) = ptr.slice(&memory_view, len) else {
+                error!(target: "wasm_runtime::db_contains_key", "Failed to make slice from ptr");
+                return -2
+            };
+
+            let mut buf = vec![0_u8; len as usize];
+            if let Err(e) = mem_slice.read_slice(&mut buf) {
+                error!(target: "wasm_runtime:db_contains_key", "Failed to read from memory slice: {}", e);
+                return -2
+            };
+
+            let mut buf_reader = Cursor::new(buf);
+
+            // FIXME: There's a type DbHandle=u32, but this should maybe be renamed
+            let db_handle: u32 = match Decodable::decode(&mut buf_reader) {
+                Ok(v) => v,
+                Err(e) => {
+                    error!(target: "wasm_runtime::db_contains_key", "Failed to decode DbHandle: {}", e);
+                    return -2
+                }
+            };
+            let db_handle = db_handle as usize;
+
+            let key: Vec<u8> = match Decodable::decode(&mut buf_reader) {
+                Ok(v) => v,
+                Err(e) => {
+                    error!(target: "wasm_runtime::db_contains_key", "Failed to decode key vec: {}", e);
+                    return -2
+                }
+            };
+
+            // TODO: Ensure we've read the entire buffer above.
+
+            let db_handles = env.db_handles.borrow();
+
+            if db_handles.len() <= db_handle {
+                error!(target: "wasm_runtime::db_contains_key", "Requested DbHandle that is out of bounds");
+                return -2
+            }
+
+            let handle_idx = db_handle;
+            let db_handle = &db_handles[handle_idx];
+
+            match db_handle.contains_key(&key) {
+                Ok(v) => {
+                    if v {
+                        return 1
+                    } else {
+                        return 0
+                    }
+                }
+                Err(e) => {
+                    error!(target: "wasm_runtime::db_contains_key", "sled.tree.contains_key failed: {}", e);
+                    return -2
+                }
+            }
+        }
+    }
+}

+ 6 - 0
src/runtime/vm_runtime.rs

@@ -199,6 +199,12 @@ impl Runtime {
                     import::db::db_get,
                 ),
 
+                "db_contains_key_" => Function::new_typed_with_env(
+                    &mut store,
+                    &ctx,
+                    import::db::db_contains_key,
+                ),
+
                 "db_set_" => Function::new_typed_with_env(
                     &mut store,
                     &ctx,

+ 25 - 0
src/sdk/src/db.rs

@@ -104,6 +104,30 @@ pub fn db_get(db_handle: DbHandle, key: &[u8]) -> GenericResult<Option<Vec<u8>>>
     Ok(Some(buf))
 }
 
+/// Everyone can call this. Checks if a key is contained in the key-value store.
+///
+/// ```
+/// if db_contains_key(db_handle, key) {
+///     println!("true");
+/// }
+/// ```
+pub fn db_contains_key(db_handle: DbHandle, key: &[u8]) -> GenericResult<bool> {
+    let mut len = 0;
+    let mut buf = vec![];
+    len += db_handle.encode(&mut buf)?;
+    len += key.to_vec().encode(&mut buf)?;
+
+    let ret = unsafe { db_contains_key_(buf.as_ptr(), len as u32) };
+
+    match ret {
+        0 => return Ok(false),
+        1 => return Ok(true),
+        -1 => return Err(ContractError::CallerAccessDenied),
+        -2 => return Err(ContractError::DbContainsKeyFailed),
+        _ => unimplemented!(),
+    }
+}
+
 /// Only update() can call this. Set a value within the transaction.
 ///
 /// ```
@@ -131,5 +155,6 @@ extern "C" {
     fn db_init_(ptr: *const u8, len: u32) -> i32;
     fn db_lookup_(ptr: *const u8, len: u32) -> i32;
     fn db_get_(ptr: *const u8, len: u32) -> i64;
+    fn db_contains_key_(ptr: *const u8, len: u32) -> i32;
     fn db_set_(ptr: *const u8, len: u32) -> i32;
 }

+ 6 - 0
src/sdk/src/error.rs

@@ -65,6 +65,9 @@ pub enum ContractError {
 
     #[error("Db get failed")]
     DbGetFailed,
+
+    #[error("Db contains_key failed")]
+    DbContainsKeyFailed,
 }
 
 /// Builtin return values occupy the upper 32 bits
@@ -87,6 +90,7 @@ pub const DB_NOT_FOUND: i64 = to_builtin!(10);
 pub const DB_SET_FAILED: i64 = to_builtin!(11);
 pub const DB_LOOKUP_FAILED: i64 = to_builtin!(12);
 pub const DB_GET_FAILED: i64 = to_builtin!(13);
+pub const DB_CONTAINS_KEY_FAILED: i64 = to_builtin!(14);
 
 impl From<ContractError> for i64 {
     fn from(err: ContractError) -> Self {
@@ -103,6 +107,7 @@ impl From<ContractError> for i64 {
             ContractError::DbSetFailed => DB_SET_FAILED,
             ContractError::DbLookupFailed => DB_LOOKUP_FAILED,
             ContractError::DbGetFailed => DB_GET_FAILED,
+            ContractError::DbContainsKeyFailed => DB_CONTAINS_KEY_FAILED,
             ContractError::Custom(error) => {
                 if error == 0 {
                     CUSTOM_ZERO
@@ -130,6 +135,7 @@ impl From<i64> for ContractError {
             DB_SET_FAILED => Self::DbSetFailed,
             DB_LOOKUP_FAILED => Self::DbLookupFailed,
             DB_GET_FAILED => Self::DbGetFailed,
+            DB_CONTAINS_KEY_FAILED => Self::DbContainsKeyFailed,
             _ => Self::Custom(error as u32),
         }
     }