Add RemoveCode

CosmWasm · Jan 11, 2023 · 5cc3dce · 5cc3dce
1 parent 1b8f280
commit 5cc3dce
Show file tree

Hide file tree

Showing 7 changed files with 138 additions and 0 deletions.
diff --git a/internal/api/bindings.h b/internal/api/bindings.h
@@ -313,6 +313,10 @@ struct UnmanagedVector save_wasm(struct cache_t *cache,
                                  struct ByteSliceView wasm,
                                  struct UnmanagedVector *error_msg);
 
+void remove_wasm(struct cache_t *cache,
+                 struct ByteSliceView checksum,
+                 struct UnmanagedVector *error_msg);
+
 struct UnmanagedVector load_wasm(struct cache_t *cache,
                                  struct ByteSliceView checksum,
                                  struct UnmanagedVector *error_msg);

diff --git a/internal/api/lib.go b/internal/api/lib.go
@@ -67,6 +67,17 @@ func StoreCode(cache Cache, wasm []byte) ([]byte, error) {
 	return copyAndDestroyUnmanagedVector(checksum), nil
 }
 
+func RemoveCode(cache Cache, checksum []byte) error {
+	cs := makeView(checksum)
+	defer runtime.KeepAlive(checksum)
+	errmsg := newUnmanagedVector(nil)
+	_, err := C.remove_wasm(cache.ptr, cs, &errmsg)
+	if err != nil {
+		return errorWithMessage(err, errmsg)
+	}
+	return nil
+}
+
 func GetCode(cache Cache, checksum []byte) ([]byte, error) {
 	cs := makeView(checksum)
 	defer runtime.KeepAlive(checksum)

diff --git a/internal/api/lib_test.go b/internal/api/lib_test.go
@@ -72,6 +72,25 @@ func TestStoreCodeAndGetCode(t *testing.T) {
 	require.Equal(t, wasm, code)
 }
 
+func TestRemoveCode(t *testing.T) {
+	cache, cleanup := withCache(t)
+	defer cleanup()
+
+	wasm, err := ioutil.ReadFile("../../testdata/hackatom.wasm")
+	require.NoError(t, err)
+
+	checksum, err := StoreCode(cache, wasm)
+	require.NoError(t, err)
+
+	// First removal works
+	err = RemoveCode(cache, checksum)
+	require.NoError(t, err)
+
+	// Second removal fails
+	err = RemoveCode(cache, checksum)
+	require.ErrorContains(t, err, "Wasm file does not exist")
+}
+
 func TestStoreCodeFailsWithBadData(t *testing.T) {
 	cache, cleanup := withCache(t)
 	defer cleanup()

diff --git a/lib.go b/lib.go
@@ -60,6 +60,10 @@ func (vm *VM) StoreCode(code WasmCode) (Checksum, error) {
 	return api.StoreCode(vm.cache, code)
 }
 
+func (vm *VM) RemoveCode(checksum Checksum) error {
+	return api.RemoveCode(vm.cache, checksum)
+}
+
 // GetCode will load the original Wasm code for the given checksum.
 // This will only succeed if that checksum was previously returned from
 // a call to StoreCode.

diff --git a/lib_test.go b/lib_test.go
@@ -114,6 +114,22 @@ func TestStoreCodeAndGet(t *testing.T) {
 	require.Equal(t, WasmCode(wasm), code)
 }
 
+func TestRemoveCode(t *testing.T) {
+	vm := withVM(t)
+
+	wasm, err := ioutil.ReadFile(HACKATOM_TEST_CONTRACT)
+	require.NoError(t, err)
+
+	checksum, err := vm.StoreCode(wasm)
+	require.NoError(t, err)
+
+	err = vm.RemoveCode(checksum)
+	require.NoError(t, err)
+
+	err = vm.RemoveCode(checksum)
+	require.ErrorContains(t, err, "Wasm file does not exist")
+}
+
 func TestHappyPath(t *testing.T) {
 	vm := withVM(t)
 	checksum := createTestContract(t, vm, HACKATOM_TEST_CONTRACT)

diff --git a/libwasmvm/bindings.h b/libwasmvm/bindings.h
@@ -313,6 +313,10 @@ struct UnmanagedVector save_wasm(struct cache_t *cache,
                                  struct ByteSliceView wasm,
                                  struct UnmanagedVector *error_msg);
 
+void remove_wasm(struct cache_t *cache,
+                 struct ByteSliceView checksum,
+                 struct UnmanagedVector *error_msg);
+
 struct UnmanagedVector load_wasm(struct cache_t *cache,
                                  struct ByteSliceView checksum,
                                  struct UnmanagedVector *error_msg);

diff --git a/libwasmvm/src/cache.rs b/libwasmvm/src/cache.rs
@@ -104,6 +104,32 @@ fn do_save_wasm(
     Ok(checksum)
 }
 
+#[no_mangle]
+pub extern "C" fn remove_wasm(
+    cache: *mut cache_t,
+    checksum: ByteSliceView,
+    error_msg: Option<&mut UnmanagedVector>,
+) {
+    let r = match to_cache(cache) {
+        Some(c) => catch_unwind(AssertUnwindSafe(move || do_remove_wasm(c, checksum)))
+            .unwrap_or_else(|_| Err(Error::panic())),
+        None => Err(Error::unset_arg(CACHE_ARG)),
+    };
+    handle_c_error_default(r, error_msg)
+}
+
+fn do_remove_wasm(
+    cache: &mut Cache<GoApi, GoStorage, GoQuerier>,
+    checksum: ByteSliceView,
+) -> Result<(), Error> {
+    let checksum: Checksum = checksum
+        .read()
+        .ok_or_else(|| Error::unset_arg(CHECKSUM_ARG))?
+        .try_into()?;
+    cache.remove_wasm(&checksum)?;
+    Ok(())
+}
+
 #[no_mangle]
 pub extern "C" fn load_wasm(
     cache: *mut cache_t,
@@ -406,6 +432,60 @@ mod tests {
         release_cache(cache_ptr);
     }
 
+    #[test]
+    fn remove_wasm_works() {
+        let dir: String = TempDir::new().unwrap().path().to_str().unwrap().to_owned();
+        let capabilities = b"staking";
+
+        let mut error_msg = UnmanagedVector::default();
+        let cache_ptr = init_cache(
+            ByteSliceView::new(dir.as_bytes()),
+            ByteSliceView::new(capabilities),
+            512,
+            32,
+            Some(&mut error_msg),
+        );
+        assert!(error_msg.is_none());
+        let _ = error_msg.consume();
+
+        let mut error_msg = UnmanagedVector::default();
+        let checksum = save_wasm(
+            cache_ptr,
+            ByteSliceView::new(HACKATOM),
+            Some(&mut error_msg),
+        );
+        assert!(error_msg.is_none());
+        let _ = error_msg.consume();
+        let checksum = checksum.consume().unwrap_or_default();
+
+        // Removing once works
+        let mut error_msg = UnmanagedVector::default();
+        remove_wasm(
+            cache_ptr,
+            ByteSliceView::new(&checksum),
+            Some(&mut error_msg),
+        );
+        assert!(error_msg.is_none());
+        let _ = error_msg.consume();
+
+        // Removing again fails
+        let mut error_msg = UnmanagedVector::default();
+        remove_wasm(
+            cache_ptr,
+            ByteSliceView::new(&checksum),
+            Some(&mut error_msg),
+        );
+        let error_msg = error_msg
+            .consume()
+            .map(|e| String::from_utf8_lossy(&e).into_owned());
+        assert_eq!(
+            error_msg.unwrap(),
+            "Error calling the VM: Cache error: Wasm file does not exist"
+        );
+
+        release_cache(cache_ptr);
+    }
+
     #[test]
     fn load_wasm_works() {
         let dir: String = TempDir::new().unwrap().path().to_str().unwrap().to_owned();