Merge pull request #216 from notional-labs/vuong/snapshotter

Vuong/snapshotter

modules/light-clients/08-wasm/keeper/keeper.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"crypto/sha256"
 	"encoding/hex"
+	"fmt"
 	"math"
 	"path/filepath"
 	"strings"
@@ -200,3 +201,32 @@ func (k Keeper) ExportGenesis(ctx sdk.Context) types.GenesisState {
 	}
 	return genesisState
 }
+
+// TODO: testing
+func (k Keeper) IterateCodeInfos(ctx sdk.Context, fn func(codeID string) (stop bool)) {
+	store := ctx.KVStore(k.storeKey)
+	prefixStore := prefix.NewStore(store, []byte(fmt.Sprintf("%s/", types.PrefixCodeIDKey)))
+
+	iter := prefixStore.Iterator(nil, nil)
+	defer iter.Close()
+
+	for ; iter.Valid(); iter.Next() {
+		if fn(string(iter.Value())) {
+			break
+		}
+	}
+}
+
+// TODO: testing
+func (k Keeper) GetWasmByte(ctx sdk.Context, codeID string) ([]byte, error) {
+	store := ctx.KVStore(k.storeKey)
+
+	byteCodeID, err := hex.DecodeString(codeID)
+	if err != nil {
+		return nil, status.Error(codes.InvalidArgument, "invalid code ID")
+	}
+
+	codeKey := types.CodeID(byteCodeID)
+	wasmBytes := store.Get(codeKey)
+	return wasmBytes, nil
+}

modules/light-clients/08-wasm/keeper/snapshotter.go

@@ -0,0 +1,134 @@
+package keeper
+
+import (
+	"io"
+
+	errorsmod "cosmossdk.io/errors"
+	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
+
+	snapshot "github.com/cosmos/cosmos-sdk/snapshots/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/ibc-go/v7/modules/light-clients/08-wasm/types"
+)
+
+var _ snapshot.ExtensionSnapshotter = &WasmSnapshotter{}
+
+// SnapshotFormat format 1 is just gzipped wasm byte code for each item payload. No protobuf envelope, no metadata.
+const SnapshotFormat = 1
+
+type WasmSnapshotter struct {
+	wasm *Keeper
+	cms  sdk.MultiStore
+}
+
+func NewWasmSnapshotter(cms sdk.MultiStore, wasm *Keeper) *WasmSnapshotter {
+	return &WasmSnapshotter{
+		wasm: wasm,
+		cms:  cms,
+	}
+}
+
+func (ws *WasmSnapshotter) SnapshotName() string {
+	return types.ModuleName
+}
+
+func (ws *WasmSnapshotter) SnapshotFormat() uint32 {
+	return SnapshotFormat
+}
+
+func (ws *WasmSnapshotter) SupportedFormats() []uint32 {
+	// If we support older formats, add them here and handle them in Restore
+	return []uint32{SnapshotFormat}
+}
+
+func (ws *WasmSnapshotter) SnapshotExtension(height uint64, payloadWriter snapshot.ExtensionPayloadWriter) error {
+	cacheMS, err := ws.cms.CacheMultiStoreWithVersion(int64(height))
+	if err != nil {
+		return err
+	}
+
+	ctx := sdk.NewContext(cacheMS, tmproto.Header{}, false, nil)
+	seenBefore := make(map[string]bool)
+	var rerr error
+
+	ws.wasm.IterateCodeInfos(ctx, func(codeID string) bool {
+		if seenBefore[codeID] {
+			return false
+		}
+		seenBefore[codeID] = true
+
+		// load code and abort on error
+		wasmBytes, err := ws.wasm.GetWasmByte(ctx, codeID)
+		if err != nil {
+			rerr = err
+			return true
+		}
+
+		compressedWasm, err := types.GzipIt(wasmBytes)
+		if err != nil {
+			rerr = err
+			return true
+		}
+
+		err = payloadWriter(compressedWasm)
+		if err != nil {
+			rerr = err
+			return true
+		}
+
+		return false
+	})
+
+	return rerr
+}
+
+func (ws *WasmSnapshotter) RestoreExtension(height uint64, format uint32, payloadReader snapshot.ExtensionPayloadReader) error {
+	if format == SnapshotFormat {
+		return ws.processAllItems(height, payloadReader, restoreV1, finalizeV1)
+	}
+	return snapshot.ErrUnknownFormat
+}
+
+func restoreV1(_ sdk.Context, k *Keeper, compressedCode []byte) error {
+	if !types.IsGzip(compressedCode) {
+		return types.ErrInvalid.Wrap("not a gzip")
+	}
+	wasmCode, err := types.Uncompress(compressedCode, uint64(types.MaxWasmSize))
+	if err != nil {
+		return errorsmod.Wrap(errorsmod.Wrap(err, "failed to store contract"), err.Error())
+	}
+
+	// FIXME: check which codeIDs the checksum matches??
+	_, err = k.wasmVM.StoreCode(wasmCode)
+	if err != nil {
+		return errorsmod.Wrap(errorsmod.Wrap(err, "failed to store contract"), err.Error())
+	}
+	return nil
+}
+
+func finalizeV1(ctx sdk.Context, k *Keeper) error {
+	return nil
+}
+
+func (ws *WasmSnapshotter) processAllItems(
+	height uint64,
+	payloadReader snapshot.ExtensionPayloadReader,
+	cb func(sdk.Context, *Keeper, []byte) error,
+	finalize func(sdk.Context, *Keeper) error,
+) error {
+	ctx := sdk.NewContext(ws.cms, tmproto.Header{Height: int64(height)}, false, nil)
+	for {
+		payload, err := payloadReader()
+		if err == io.EOF {
+			break
+		} else if err != nil {
+			return err
+		}
+
+		if err := cb(ctx, ws.wasm, payload); err != nil {
+			return errorsmod.Wrap(err, "processing snapshot item")
+		}
+	}
+
+	return finalize(ctx, ws.wasm)
+}

modules/light-clients/08-wasm/keeper/utils.go

@@ -62,3 +62,20 @@ func (l *LimitedReader) Read(p []byte) (n int, err error) {
 	}
 	return l.r.Read(p)
 }
+
+// GzipIt compresses the input ([]byte)
+func GzipIt(input []byte) ([]byte, error) {
+	// Create gzip writer.
+	var b bytes.Buffer
+	w := gzip.NewWriter(&b)
+	_, err := w.Write(input)
+	if err != nil {
+		return nil, err
+	}
+	err = w.Close() // You must close this first to flush the bytes to the buffer.
+	if err != nil {
+		return nil, err
+	}
+
+	return b.Bytes(), nil
+}

modules/light-clients/08-wasm/types/upgrade.go

@@ -5,6 +5,7 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	clienttypes "github.com/cosmos/ibc-go/v7/modules/core/02-client/types"
+
 	"github.com/cosmos/ibc-go/v7/modules/core/exported"
 )
 

modules/light-clients/08-wasm/types/utils.go

@@ -0,0 +1,71 @@
+package types
+
+import (
+	"bytes"
+	"compress/gzip"
+	"io"
+)
+
+// Copied gzip feature from wasmd
+// https://github.com/CosmWasm/wasmd/blob/v0.31.0/x/wasm/ioutils/utils.go
+
+// Note: []byte can never be const as they are inherently mutable
+
+// magic bytes to identify gzip.
+// See https://www.ietf.org/rfc/rfc1952.txt
+// and https://github.com/golang/go/blob/master/src/net/http/sniff.go#L186
+var gzipIdent = []byte("\x1F\x8B\x08")
+
+// IsGzip returns checks if the file contents are gzip compressed
+func IsGzip(input []byte) bool {
+	return len(input) >= 3 && bytes.Equal(gzipIdent, input[0:3])
+}
+
+// Uncompress expects a valid gzip source to unpack or fails. See IsGzip
+func Uncompress(gzipSrc []byte, limit uint64) ([]byte, error) {
+	if uint64(len(gzipSrc)) > limit {
+		return nil, ErrWasmCodeTooLarge
+	}
+	zr, err := gzip.NewReader(bytes.NewReader(gzipSrc))
+	if err != nil {
+		return nil, err
+	}
+	zr.Multistream(false)
+	defer zr.Close()
+	return io.ReadAll(limitReader(zr, int64(limit)))
+}
+
+// limitReader returns a Reader that reads from r
+// but stops with types.ErrLimit after n bytes.
+// The underlying implementation is a *io.LimitedReader.
+func limitReader(r io.Reader, n int64) io.Reader {
+	return &limitedReader{r: &io.LimitedReader{R: r, N: n}}
+}
+
+type limitedReader struct {
+	r *io.LimitedReader
+}
+
+func (l *limitedReader) Read(p []byte) (n int, err error) {
+	if l.r.N <= 0 {
+		return 0, ErrWasmCodeTooLarge
+	}
+	return l.r.Read(p)
+}
+
+// GzipIt compresses the input ([]byte)
+func GzipIt(input []byte) ([]byte, error) {
+	// Create gzip writer.
+	var b bytes.Buffer
+	w := gzip.NewWriter(&b)
+	_, err := w.Write(input)
+	if err != nil {
+		return nil, err
+	}
+	err = w.Close() // You must close this first to flush the bytes to the buffer.
+	if err != nil {
+		return nil, err
+	}
+
+	return b.Bytes(), nil
+}

modules/light-clients/08-wasm/types/validation.go

@@ -1,6 +1,6 @@
 package types
 
-var MaxWasmSize = 3 * 1024 * 1024
+const MaxWasmSize = 3 * 1024 * 1024
 
 func ValidateWasmCode(code []byte) (bool, error) {
 	if len(code) == 0 {