From c598523e9d644a3c61beab09bffebd19f389d287 Mon Sep 17 00:00:00 2001
From: Piotr Sarna <sarna@chiselstrike.com>
Date: Fri, 19 May 2023 12:23:42 +0200
Subject: [PATCH] udf: fix memory leak for libsql_wasm_engine_new

The leak was quite obvious - there was no destructor
for the per-connection Wasm engine.
---
 .gitignore                            |  1 +
 ext/udf/wasm_bindings.h               |  1 +
 ext/udf/wasmedge_bindings.c           |  3 +++
 src/main.c                            |  5 +++++
 src/rust/wasmtime-bindings/Cargo.toml |  2 +-
 src/rust/wasmtime-bindings/src/lib.rs | 10 ++++++----
 6 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/.gitignore b/.gitignore
index b704d88a9e..6092471009 100644
--- a/.gitignore
+++ b/.gitignore
@@ -59,3 +59,4 @@ test/rust_suite/Cargo.lock
 testfixture
 libsql
 src/rust/libsql-shell/target/
+src/rust/wasmtime-bindings/Cargo.lock
diff --git a/ext/udf/wasm_bindings.h b/ext/udf/wasm_bindings.h
index fcef3f4a74..4c1e88b3e8 100644
--- a/ext/udf/wasm_bindings.h
+++ b/ext/udf/wasm_bindings.h
@@ -48,6 +48,7 @@ void libsql_free_wasm_module(void *module);
 ** Creates a new wasm engine
 */
 libsql_wasm_engine_t *libsql_wasm_engine_new();
+void libsql_wasm_engine_free(libsql_wasm_engine_t *);
 
 #endif //LIBSQL_WASM_BINDINGS_H
 #endif //LIBSQL_ENABLE_WASM_RUNTIME
diff --git a/ext/udf/wasmedge_bindings.c b/ext/udf/wasmedge_bindings.c
index 17e312aa45..dccf39a1f4 100644
--- a/ext/udf/wasmedge_bindings.c
+++ b/ext/udf/wasmedge_bindings.c
@@ -175,6 +175,9 @@ libsql_wasm_engine_t *libsql_wasm_engine_new() {
   return NULL;
 }
 
+void libsql_wasm_engine_free(libsql_wasm_engine_t *eng) {
+}
+
 libsql_wasm_module_t *libsql_compile_wasm_module(libsql_wasm_engine_t* engine, const char *pSrcBody, int nBody,
     void *(*alloc_err_buf)(unsigned long long), char **err_msg_buf) {
 
diff --git a/src/main.c b/src/main.c
index 10018c8729..60439874a1 100644
--- a/src/main.c
+++ b/src/main.c
@@ -1266,6 +1266,11 @@ static int sqlite3Close(sqlite3 *db, int forceZombie){
   /* Convert the connection into a zombie and then close it.
   */
   db->eOpenState = SQLITE_STATE_ZOMBIE;
+#ifdef LIBSQL_ENABLE_WASM_RUNTIME
+  if (db->wasm.engine) {
+    libsql_wasm_engine_free(db->wasm.engine);
+  }
+#endif
   sqlite3LeaveMutexAndCloseZombie(db);
   return SQLITE_OK;
 }
diff --git a/src/rust/wasmtime-bindings/Cargo.toml b/src/rust/wasmtime-bindings/Cargo.toml
index cd26519466..f4cb273659 100644
--- a/src/rust/wasmtime-bindings/Cargo.toml
+++ b/src/rust/wasmtime-bindings/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "libsql-wasmtime-bindings"
-version = "0.2.0"
+version = "0.2.1"
 edition = "2021"
 license = "Apache-2.0"
 description = "Wasmtime bindings for libSQL user-defined functions"
diff --git a/src/rust/wasmtime-bindings/src/lib.rs b/src/rust/wasmtime-bindings/src/lib.rs
index 4108c20261..db704305ad 100644
--- a/src/rust/wasmtime-bindings/src/lib.rs
+++ b/src/rust/wasmtime-bindings/src/lib.rs
@@ -85,10 +85,12 @@ pub fn libsql_wasm_engine_new() -> *const c_void {
         Err(_) => return std::ptr::null() as *const c_void,
     };
 
-    let engine = Box::new(engine);
-    let engine_ptr = &*engine as *const Engine as *const c_void;
-    std::mem::forget(engine);
-    engine_ptr
+    Box::into_raw(Box::new(engine)) as *const c_void
+}
+
+#[no_mangle]
+pub fn libsql_wasm_engine_free(engine: *mut c_void) {
+    unsafe { Box::from_raw(engine as *mut Engine) };
 }
 
 #[repr(C)]