Make snapshot from keeper

app/app.go

@@ -684,7 +684,7 @@ func NewWasmApp(
 	app.scopedWasmKeeper = scopedWasmKeeper
 
 	app.SnapshotManager().RegisterExtensions(
-		wasm.NewWasmSnapshotter(filepath.Join(wasmDir, "wasm", "state", "wasm")),
+		wasm.NewWasmSnapshotter(app.CommitMultiStore(), app.wasmKeeper),
 	)
 
 	return app

x/wasm/wasm_snapshotter.go

@@ -5,13 +5,16 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
-	"regexp"
-	"sort"
-	"strings"
 
 	snapshot "github.com/cosmos/cosmos-sdk/snapshots/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	protoio "github.com/gogo/protobuf/io"
+	"github.com/tendermint/tendermint/libs/log"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
+
+	"github.com/CosmWasm/wasmd/x/wasm/keeper"
+	"github.com/CosmWasm/wasmd/x/wasm/types"
 )
 
 /*
@@ -49,17 +52,21 @@ type ExtensionSnapshotter interface {
 */
 
 type WasmSnapshotter struct {
+	wasm keeper.Keeper
+	cms  sdk.CommitMultiStore
+	// obsolete placeholder to compile
 	wasmDirectory string
 }
 
-func NewWasmSnapshotter(wasmDirectory string) *WasmSnapshotter {
+func NewWasmSnapshotter(cms sdk.CommitMultiStore, wasm keeper.Keeper) *WasmSnapshotter {
 	return &WasmSnapshotter{
-		wasmDirectory,
+		wasm: wasm,
+		cms:  cms,
 	}
 }
 
 func (ws *WasmSnapshotter) SnapshotName() string {
-	return "WASM Files Snapshot"
+	return "WASM Code Snapshot"
 }
 
 func (ws *WasmSnapshotter) SnapshotFormat() uint32 {
@@ -70,35 +77,24 @@ func (ws *WasmSnapshotter) SupportedFormats() []uint32 {
 	return []uint32{1}
 }
 
-var wasmFileNameRegex = regexp.MustCompile(`^[a-f0-9]{64}$`)
-
 func (ws *WasmSnapshotter) Snapshot(height uint64, protoWriter protoio.Writer) error {
-	wasmFiles, err := ioutil.ReadDir(ws.wasmDirectory)
-	if err != nil {
-		return err
-	}
+	var rerr error
+	ctx := sdk.NewContext(ws.cms, tmproto.Header{}, false, log.NewNopLogger())
 
-	// In case snapshotting needs to be deterministic
-	sort.SliceStable(wasmFiles, func(i, j int) bool {
-		return strings.Compare(wasmFiles[i].Name(), wasmFiles[j].Name()) < 0
-	})
-
-	for _, wasmFile := range wasmFiles {
-		if !wasmFileNameRegex.MatchString(wasmFile.Name()) {
-			continue
-		}
-
-		wasmFilePath := filepath.Join(ws.wasmDirectory, wasmFile.Name())
-		wasmBytes, err := ioutil.ReadFile(wasmFilePath)
+	ws.wasm.IterateCodeInfos(ctx, func(id uint64, info types.CodeInfo) bool {
+		// load code and abort on error
+		wasmBytes, err := ws.wasm.GetByteCode(ctx, id)
 		if err != nil {
-			return err
+			rerr = err
+			return true
 		}
 
-		// snapshotItem is 64 bytes of the file name, then the actual WASM bytes
-		snapshotItem := append([]byte(wasmFile.Name()), wasmBytes...)
+		// TODO: compress wasm bytes
+		// TODO: embed in a protobuf message
+		snapshot.WriteExtensionItem(protoWriter, wasmBytes)
 
-		snapshot.WriteExtensionItem(protoWriter, snapshotItem)
-	}
+		return false
+	})
 
 	return nil
 }