From 5546d198fba1903ddcc3d0cfade35f7ff3071236 Mon Sep 17 00:00:00 2001
From: Assaf Morami <assaf.morami@gmail.com>
Date: Mon, 31 Oct 2022 12:03:48 +0200
Subject: [PATCH] WASM snapshotter: write directly to disk

---
 app/app.go                                    |  3 +-
 x/compute/internal/keeper/wasm_snapshotter.go | 30 +++++++++++--------
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/app/app.go b/app/app.go
index 98bab2f09..d1fce2e56 100644
--- a/app/app.go
+++ b/app/app.go
@@ -5,6 +5,7 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"path"
 	"path/filepath"
 
 	"github.com/scrtlabs/SecretNetwork/app/keepers"
@@ -379,7 +380,7 @@ func NewSecretNetworkApp(
 
 	if manager := app.BaseApp.SnapshotManager(); manager != nil {
 		err := manager.RegisterExtensions(
-			compute.NewWasmSnapshotter(app.BaseApp.CommitMultiStore(), app.AppKeepers.ComputeKeeper),
+			compute.NewWasmSnapshotter(app.BaseApp.CommitMultiStore(), app.AppKeepers.ComputeKeeper, path.Join(homePath, ".compute", "wasm", "wasm")),
 		)
 		if err != nil {
 			panic(fmt.Errorf("failed to register snapshot extension: %s", err))
diff --git a/x/compute/internal/keeper/wasm_snapshotter.go b/x/compute/internal/keeper/wasm_snapshotter.go
index 61e98c2b4..6dfd21c59 100644
--- a/x/compute/internal/keeper/wasm_snapshotter.go
+++ b/x/compute/internal/keeper/wasm_snapshotter.go
@@ -1,9 +1,11 @@
 package keeper
 
 import (
+	"crypto/sha256"
 	"encoding/hex"
-	"fmt"
 	"io"
+	"os"
+	"path/filepath"
 
 	snapshottypes "github.com/cosmos/cosmos-sdk/snapshots/types"
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
@@ -52,14 +54,16 @@ type ExtensionSnapshotter interface {
 var _ snapshottypes.ExtensionSnapshotter = (*WasmSnapshotter)(nil)
 
 type WasmSnapshotter struct {
-	cms    storetypes.MultiStore
-	keeper *Keeper
+	cms           storetypes.MultiStore
+	keeper        *Keeper
+	wasmDirectory string
 }
 
-func NewWasmSnapshotter(cms storetypes.MultiStore, keeper *Keeper) *WasmSnapshotter {
+func NewWasmSnapshotter(cms storetypes.MultiStore, keeper *Keeper, wasmDirectory string) *WasmSnapshotter {
 	return &WasmSnapshotter{
-		cms:    cms,
-		keeper: keeper,
+		cms:           cms,
+		keeper:        keeper,
+		wasmDirectory: wasmDirectory,
 	}
 }
 
@@ -144,15 +148,17 @@ func (ws *WasmSnapshotter) Restore(
 			return item, nil
 		}
 
-		wasmCode := payload.Payload
+		wasmBytes := payload.Payload
 
-		fmt.Println("Writing wasm file to disk")
+		hash := sha256.Sum256(wasmBytes)
 
-		// Store the WASM bytes using the existing API
-		// FIXME: check which codeIDs the checksum matches??
-		_, err = ws.keeper.wasmer.Create(wasmCode)
+		wasmFileName := hex.EncodeToString(hash[:])
+
+		wasmFilePath := filepath.Join(ws.wasmDirectory, wasmFileName)
+
+		err = os.WriteFile(wasmFilePath, wasmBytes, 0o600 /* -rw-rw-r-- */)
 		if err != nil {
-			return snapshottypes.SnapshotItem{}, sdkerrors.Wrap(types.ErrCreateFailed, err.Error())
+			return snapshottypes.SnapshotItem{}, sdkerrors.Wrapf(err, "failed to write wasm file '%v' to disk", wasmFilePath)
 		}
 	}
 }