snapshot.go

package keeper

import (
	"bytes"
	"encoding/binary"
	"io"
	"sort"

	"github.com/CosmWasm/wasmd/x/wasm/types"

	snapshottypes "github.com/cosmos/cosmos-sdk/snapshots/types"
	"github.com/cosmos/cosmos-sdk/store/prefix"
	storetypes "github.com/cosmos/cosmos-sdk/store/types"
	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
	protoio "github.com/gogo/protobuf/io"
)

var _ snapshottypes.ExtensionSnapshotter = (*Snapshotter)(nil)

type Snapshotter struct {
	cms    storetypes.MultiStore
	keeper *Keeper
}

func NewSnapshotter(cms storetypes.MultiStore, keeper *Keeper) *Snapshotter {
	return &Snapshotter{
		cms:    cms,
		keeper: keeper,
	}
}

// SnapshotName implements types.ExtensionSnapshotter
func (*Snapshotter) SnapshotName() string {
	return types.ModuleName
}

// SnapshotFormat implements types.ExtensionSnapshotter
func (*Snapshotter) SnapshotFormat() uint32 {
	return types.SnapshotFormat
}

// SupportedFormats implements types.ExtensionSnapshotter
func (*Snapshotter) SupportedFormats() []uint32 {
	return []uint32{types.SnapshotFormat}
}

// Snapshot implements types.Snapshotter
func (ws *Snapshotter) Snapshot(height uint64, protoWriter protoio.Writer) error {
	cacheMS, err := ws.cms.CacheMultiStoreWithVersion(int64(height))
	if err != nil {
		return err
	}

	prefixStore := prefix.NewStore(cacheMS.GetKVStore(ws.keeper.storeKey), types.CodeKeyPrefix)
	iter := prefixStore.Iterator(nil, nil)
	defer iter.Close()

	type codeItem struct {
		CodeID   uint64
		CodeHash []byte
	}
	uniqueHashes := make(map[string]bool)
	codeItems := []codeItem{}

	for ; iter.Valid(); iter.Next() {
		var c types.CodeInfo
		ws.keeper.cdc.MustUnmarshal(iter.Value(), &c)
		codeID := binary.BigEndian.Uint64(iter.Key())

		if _, ok := uniqueHashes[string(c.CodeHash)]; ok {
			continue
		}
		uniqueHashes[string(c.CodeHash)] = true

		codeItems = append(codeItems, codeItem{
			CodeID:   codeID,
			CodeHash: c.CodeHash,
		})
	}

	sort.SliceStable(codeItems, func(i, j int) bool {
		return codeItems[i].CodeID < codeItems[j].CodeID
	})

	for _, item := range codeItems {
		bytecode, err := ws.keeper.wasmVM.GetCode(item.CodeHash)
		if err != nil {
			return err
		}

		wasmItem := types.SnapshotWasmItem{
			CodeHash:     item.CodeHash,
			WASMByteCode: bytecode,
		}

		data, err := ws.keeper.cdc.Marshal(&wasmItem)
		if err != nil {
			return sdkerrors.Wrap(err, "cannot encode protobuf wasm message")
		}

		snapshottypes.WriteExtensionItem(protoWriter, data)
	}

	return nil
}

// Restore implements types.Snapshotter
func (ws Snapshotter) Restore(height uint64, format uint32, protoReader protoio.Reader) (snapshottypes.SnapshotItem, error) {
	if format != types.SnapshotFormat {
		return snapshottypes.SnapshotItem{}, snapshottypes.ErrUnknownFormat
	}

	var snapshotItem snapshottypes.SnapshotItem
	for {
		snapshotItem = snapshottypes.SnapshotItem{}
		err := protoReader.ReadMsg(&snapshotItem)
		if err == io.EOF {
			break
		} else if err != nil {
			return snapshottypes.SnapshotItem{}, sdkerrors.Wrap(err, "invalid protobuf message")
		}

		extension := snapshotItem.GetExtensionPayload()
		// break and return item to process other modules.
		if extension == nil {
			break
		}

		var wasmItem types.SnapshotWasmItem
		err = ws.keeper.cdc.Unmarshal(extension.Payload, &wasmItem)
		if err != nil {
			return snapshottypes.SnapshotItem{}, sdkerrors.Wrap(err, "invalid protobuf wasm message")
		}

		codeHash, err := ws.keeper.wasmVM.Create(wasmItem.WASMByteCode)
		if err != nil {
			return snapshottypes.SnapshotItem{}, sdkerrors.Wrap(types.ErrCreateFailed, err.Error())
		}

		if !bytes.Equal(wasmItem.CodeHash, codeHash) {
			return snapshottypes.SnapshotItem{}, sdkerrors.Wrap(types.ErrInvalid, "code hashes not same")
		}
	}
	// TODO: iniitialize pin codes

	return snapshotItem, nil
}

func (ws Snapshotter) getCodeInfo(store storetypes.KVStore, codeID uint64) *types.CodeInfo {
	var codeInfo types.CodeInfo
	codeInfoBz := store.Get(types.GetCodeKey(codeID))
	if codeInfoBz == nil {
		return nil
	}
	ws.keeper.cdc.MustUnmarshal(codeInfoBz, &codeInfo)
	return &codeInfo
}