Add WasmSnapshotter

This helps state sync to track the WASM files too, as they're not in IAVL

app/app.go

@@ -682,6 +682,11 @@ func NewWasmApp(
 	app.scopedIBCKeeper = scopedIBCKeeper
 	app.scopedTransferKeeper = scopedTransferKeeper
 	app.scopedWasmKeeper = scopedWasmKeeper
+
+	app.SnapshotManager().RegisterExtensions(
+		wasm.NewWasmSnapshotter(filepath.Join(wasmDir, "wasm", "state", "wasm")),
+	)
+
 	return app
 }
 

x/wasm/wasm_snapshotter.go

@@ -0,0 +1,150 @@
+package wasm
+
+import (
+	"io"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"regexp"
+	"sort"
+	"strings"
+
+	snapshot "github.com/cosmos/cosmos-sdk/snapshots/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+	protoio "github.com/gogo/protobuf/io"
+)
+
+/*
+API to implement:
+
+// Snapshotter is something that can create and restore snapshots, consisting of streamed binary
+// chunks - all of which must be read from the channel and closed. If an unsupported format is
+// given, it must return ErrUnknownFormat (possibly wrapped with fmt.Errorf).
+type Snapshotter interface {
+	// Snapshot writes snapshot items into the protobuf writer.
+	Snapshot(height uint64, protoWriter protoio.Writer) error
+
+	// Restore restores a state snapshot from the protobuf items read from the reader.
+	// If the ready channel is non-nil, it returns a ready signal (by being closed) once the
+	// restorer is ready to accept chunks.
+	Restore(height uint64, format uint32, protoReader protoio.Reader) (SnapshotItem, error)
+}
+
+// ExtensionSnapshotter is an extension Snapshotter that is appended to the snapshot stream.
+// ExtensionSnapshotter has an unique name and manages it's own internal formats.
+type ExtensionSnapshotter interface {
+	Snapshotter
+
+	// SnapshotName returns the name of snapshotter, it should be unique in the manager.
+	SnapshotName() string
+
+	// SnapshotFormat returns the default format the extension snapshotter use to encode the
+	// payloads when taking a snapshot.
+	// It's defined within the extension, different from the global format for the whole state-sync snapshot.
+	SnapshotFormat() uint32
+
+	// SupportedFormats returns a list of formats it can restore from.
+	SupportedFormats() []uint32
+}
+*/
+
+type WasmSnapshotter struct {
+	wasmDirectory string
+}
+
+func NewWasmSnapshotter(wasmDirectory string) *WasmSnapshotter {
+	return &WasmSnapshotter{
+		wasmDirectory,
+	}
+}
+
+func (ws *WasmSnapshotter) SnapshotName() string {
+	return "WASM Files Snapshot"
+}
+
+func (ws *WasmSnapshotter) SnapshotFormat() uint32 {
+	return 1
+}
+
+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
+	}
+
+	// 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)
+		if err != nil {
+			return err
+		}
+
+		// snapshotItem is 64 bytes of the file name, then the actual WASM bytes
+		snapshotItem := append([]byte(wasmFile.Name()), wasmBytes...)
+
+		snapshot.WriteExtensionItem(protoWriter, snapshotItem)
+	}
+
+	return nil
+}
+
+func (ws *WasmSnapshotter) Restore(
+	height uint64, format uint32, protoReader protoio.Reader,
+) (snapshot.SnapshotItem, error) {
+	if format != 1 {
+		return snapshot.SnapshotItem{}, snapshot.ErrUnknownFormat
+	}
+
+	// Create .compute directory if it doesn't exist already
+	err := os.MkdirAll(ws.wasmDirectory, os.ModePerm)
+	if err != nil {
+		return snapshot.SnapshotItem{}, sdkerrors.Wrapf(err, "failed to create directory '%s'", ws.wasmDirectory)
+	}
+
+	for {
+		item := &snapshot.SnapshotItem{}
+		err = protoReader.ReadMsg(item)
+		if err == io.EOF {
+			break
+		} else if err != nil {
+			return snapshot.SnapshotItem{}, sdkerrors.Wrap(err, "invalid protobuf message")
+		}
+
+		payload := item.GetExtensionPayload()
+		if payload == nil {
+			return snapshot.SnapshotItem{}, sdkerrors.Wrap(err, "invalid protobuf message")
+		}
+
+		// snapshotItem is 64 bytes of the file name, then the actual WASM bytes
+		if len(payload.Payload) < 64 {
+			return snapshot.SnapshotItem{}, sdkerrors.Wrapf(err, "wasm snapshot must be at least 64 bytes, got %v bytes", len(payload.Payload))
+		}
+
+		wasmFileName := string(payload.Payload[0:64])
+		wasmBytes := payload.Payload[64:]
+
+		wasmFilePath := filepath.Join(ws.wasmDirectory, wasmFileName)
+
+		err = ioutil.WriteFile(wasmFilePath, wasmBytes, 0664 /* -rw-rw-r-- */)
+		if err != nil {
+			return snapshot.SnapshotItem{}, sdkerrors.Wrapf(err, "failed to write wasm file '%v' to disk", wasmFilePath)
+		}
+	}
+
+	return snapshot.SnapshotItem{}, nil
+}