From 5a70b3e939d307ed9da033c5933fc4933c536003 Mon Sep 17 00:00:00 2001
From: codchen <codchen03@gmail.com>
Date: Mon, 6 Feb 2023 15:40:49 +0800
Subject: [PATCH] Add mutex around WasmVM write calls

---
 x/wasm/keeper/keeper.go     |   2 +-
 x/wasm/keeper/vm_wrapper.go | 122 ++++++++++++++++++++++++++++++++++++
 2 files changed, 123 insertions(+), 1 deletion(-)
 create mode 100644 x/wasm/keeper/vm_wrapper.go

diff --git a/x/wasm/keeper/keeper.go b/x/wasm/keeper/keeper.go
index 3fdcbec..20fcf78 100644
--- a/x/wasm/keeper/keeper.go
+++ b/x/wasm/keeper/keeper.go
@@ -118,7 +118,7 @@ func NewKeeper(
 	keeper := &Keeper{
 		storeKey:          storeKey,
 		cdc:               cdc,
-		wasmVM:            wasmer,
+		wasmVM:            NewVMWrapper(wasmer),
 		accountKeeper:     accountKeeper,
 		bank:              NewBankCoinTransferrer(bankKeeper),
 		portKeeper:        portKeeper,
diff --git a/x/wasm/keeper/vm_wrapper.go b/x/wasm/keeper/vm_wrapper.go
new file mode 100644
index 0000000..f26d69b
--- /dev/null
+++ b/x/wasm/keeper/vm_wrapper.go
@@ -0,0 +1,122 @@
+package keeper
+
+import (
+	"sync"
+
+	"github.com/CosmWasm/wasmd/x/wasm/types"
+	wasmvm "github.com/CosmWasm/wasmvm"
+	wasmvmtypes "github.com/CosmWasm/wasmvm/types"
+)
+
+type VMWrapper struct {
+	types.WasmerEngine
+
+	mu *sync.Mutex
+}
+
+func NewVMWrapper(inner types.WasmerEngine) types.WasmerEngine {
+	return &VMWrapper{
+		inner,
+		&sync.Mutex{},
+	}
+}
+
+func (w *VMWrapper) Create(code wasmvm.WasmCode) (wasmvm.Checksum, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.WasmerEngine.Create(code)
+}
+
+func (w *VMWrapper) Instantiate(
+	checksum wasmvm.Checksum,
+	env wasmvmtypes.Env,
+	info wasmvmtypes.MessageInfo,
+	initMsg []byte,
+	store wasmvm.KVStore,
+	goapi wasmvm.GoAPI,
+	querier wasmvm.Querier,
+	gasMeter wasmvm.GasMeter,
+	gasLimit uint64,
+	deserCost wasmvmtypes.UFraction,
+) (*wasmvmtypes.Response, uint64, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.WasmerEngine.Instantiate(checksum, env, info, initMsg, store, goapi, querier, gasMeter, gasLimit, deserCost)
+}
+
+func (w *VMWrapper) Execute(
+	code wasmvm.Checksum,
+	env wasmvmtypes.Env,
+	info wasmvmtypes.MessageInfo,
+	executeMsg []byte,
+	store wasmvm.KVStore,
+	goapi wasmvm.GoAPI,
+	querier wasmvm.Querier,
+	gasMeter wasmvm.GasMeter,
+	gasLimit uint64,
+	deserCost wasmvmtypes.UFraction,
+) (*wasmvmtypes.Response, uint64, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.WasmerEngine.Execute(code, env, info, executeMsg, store, goapi, querier, gasMeter, gasLimit, deserCost)
+}
+
+func (w *VMWrapper) Migrate(
+	checksum wasmvm.Checksum,
+	env wasmvmtypes.Env,
+	migrateMsg []byte,
+	store wasmvm.KVStore,
+	goapi wasmvm.GoAPI,
+	querier wasmvm.Querier,
+	gasMeter wasmvm.GasMeter,
+	gasLimit uint64,
+	deserCost wasmvmtypes.UFraction,
+) (*wasmvmtypes.Response, uint64, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.WasmerEngine.Migrate(checksum, env, migrateMsg, store, goapi, querier, gasMeter, gasLimit, deserCost)
+}
+
+func (w *VMWrapper) Sudo(
+	checksum wasmvm.Checksum,
+	env wasmvmtypes.Env,
+	sudoMsg []byte,
+	store wasmvm.KVStore,
+	goapi wasmvm.GoAPI,
+	querier wasmvm.Querier,
+	gasMeter wasmvm.GasMeter,
+	gasLimit uint64,
+	deserCost wasmvmtypes.UFraction,
+) (*wasmvmtypes.Response, uint64, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.WasmerEngine.Sudo(checksum, env, sudoMsg, store, goapi, querier, gasMeter, gasLimit, deserCost)
+}
+
+func (w *VMWrapper) Reply(
+	checksum wasmvm.Checksum,
+	env wasmvmtypes.Env,
+	reply wasmvmtypes.Reply,
+	store wasmvm.KVStore,
+	goapi wasmvm.GoAPI,
+	querier wasmvm.Querier,
+	gasMeter wasmvm.GasMeter,
+	gasLimit uint64,
+	deserCost wasmvmtypes.UFraction,
+) (*wasmvmtypes.Response, uint64, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.WasmerEngine.Reply(checksum, env, reply, store, goapi, querier, gasMeter, gasLimit, deserCost)
+}
+
+func (w *VMWrapper) Unpin(checksum wasmvm.Checksum) error {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.WasmerEngine.Unpin(checksum)
+}
+
+func (w *VMWrapper) Pin(checksum wasmvm.Checksum) error {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.WasmerEngine.Pin(checksum)
+}