feat(cosmos): support InstallBundle and SendChunk

Author: michaelfig

golang/cosmos/x/swingset/abci.go

@@ -35,6 +35,8 @@ type afterCommitBlockAction struct {
 func BeginBlock(ctx sdk.Context, req abci.RequestBeginBlock, keeper Keeper) error {
 	defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyBeginBlocker)
 
+	keeper.PruneExpiredBundleInstalls(ctx)
+
 	action := beginBlockAction{
 		ChainID: ctx.ChainID(),
 		Params:  keeper.GetParams(ctx),

golang/cosmos/x/swingset/keeper/grpc_query.go

@@ -61,3 +61,20 @@ func (k Querier) Mailbox(c context.Context, req *types.QueryMailboxRequest) (*ty
 		Value: value,
 	}, nil
 }
+
+func (k Querier) PendingInstall(c context.Context, req *types.QueryPendingInstallRequest) (*types.QueryPendingInstallResponse, error) {
+	if req == nil {
+		return nil, status.Error(codes.InvalidArgument, "empty request")
+	}
+	ctx := sdk.UnwrapSDKContext(c)
+
+	msg := k.GetPendingBundleInstall(ctx, req.PendingId)
+	if msg == nil {
+		return nil, status.Error(codes.NotFound, "pending install not found")
+	}
+
+	return &types.QueryPendingInstallResponse{
+		PendingId:    req.PendingId,
+		BundleChunks: msg.BundleChunks,
+	}, nil
+}

golang/cosmos/x/swingset/keeper/keeper.go

@@ -47,8 +47,11 @@ const (
 )
 
 const (
-	stateKey            = "state"
-	swingStoreKeyPrefix = "swingStore."
+	stateKey                      = "state"
+	pendingChunkDataKeyPrefix     = "pendingChunkData."
+	pendingBundleInstallKeyPrefix = "pendingBundleInstall."
+	pendingNodeKeyPrefix          = "pending."
+	swingStoreKeyPrefix           = "swingStore."
 )
 
 // Keeper maintains the link to data vstorage and exposes getter/setter methods for the various parts of the state machine
@@ -485,6 +488,158 @@ func (k Keeper) SetMailbox(ctx sdk.Context, peer string, mailbox string) {
 	k.vstorageKeeper.LegacySetStorageAndNotify(ctx, agoric.NewKVEntry(path, mailbox))
 }
 
+func (k Keeper) GetPendingChunkData(ctx sdk.Context, pendingId uint64, chunkIndex uint64) []byte {
+	store := ctx.KVStore(k.storeKey)
+	pendingChunkData := prefix.NewStore(store, []byte(pendingChunkDataKeyPrefix))
+
+	key := append(sdk.Uint64ToBigEndian(pendingId), sdk.Uint64ToBigEndian(chunkIndex)...)
+	if !pendingChunkData.Has(key) {
+		return nil
+	}
+	return pendingChunkData.Get(key)
+}
+
+func (k Keeper) SetPendingChunkData(ctx sdk.Context, pendingId uint64, chunkIndex uint64, data []byte) {
+	store := ctx.KVStore(k.storeKey)
+	pendingChunkData := prefix.NewStore(store, []byte(pendingChunkDataKeyPrefix))
+
+	key := append(sdk.Uint64ToBigEndian(pendingId), sdk.Uint64ToBigEndian(chunkIndex)...)
+	if len(data) == 0 {
+		pendingChunkData.Delete(key)
+		return
+	}
+	pendingChunkData.Set(key, data)
+}
+
+func (k Keeper) GetPendingBundleInstall(ctx sdk.Context, pendingId uint64) *types.MsgInstallBundle {
+	store := ctx.KVStore(k.storeKey)
+	pendingStore := prefix.NewStore(store, []byte(pendingBundleInstallKeyPrefix))
+	key := sdk.Uint64ToBigEndian(pendingId)
+	if !pendingStore.Has(key) {
+		return nil
+	}
+	bz := pendingStore.Get(key)
+	msg := &types.MsgInstallBundle{}
+	k.cdc.MustUnmarshal(bz, msg)
+	return msg
+}
+
+func (k Keeper) AddPendingBundleInstall(ctx sdk.Context, msg *types.MsgInstallBundle) uint64 {
+	state := k.GetState(ctx)
+	state.LastPendingId++
+	k.SetState(ctx, state)
+
+	pendingId := state.LastPendingId
+	k.SetPendingBundleInstall(ctx, pendingId, msg)
+
+	// Attach the pending install node to the linked list.
+	node := &types.PendingInstallNode{
+		PendingId:  pendingId,
+		StartTime:  ctx.BlockTime().Unix(),
+		StartBlock: ctx.BlockHeight(),
+	}
+	store := ctx.KVStore(k.storeKey)
+	key := sdk.Uint64ToBigEndian(pendingId)
+	startStore := prefix.NewStore(store, []byte(pendingNodeKeyPrefix))
+	bz := k.cdc.MustMarshal(node)
+	startStore.Set(key, bz)
+
+	return pendingId
+}
+
+// PruneExpiredBundleInstalls removes pending bundle installs that have passed
+// their deadline, as set by the keeper parameters.
+func (k Keeper) PruneExpiredBundleInstalls(ctx sdk.Context) {
+	params := k.GetParams(ctx)
+	currentSeconds := ctx.BlockTime().Unix()
+	currentBlocks := ctx.BlockHeight()
+	deadlineSeconds := params.PendingInstallationDeadlineSeconds
+	deadlineBlocks := params.PendingInstallationDeadlineBlocks
+
+	store := ctx.KVStore(k.storeKey)
+	startStore := prefix.NewStore(store, []byte(pendingNodeKeyPrefix))
+
+	state := k.GetState(ctx)
+	for state.FirstPendingId != 0 {
+		pendingId := state.FirstPendingId
+		key := sdk.Uint64ToBigEndian(pendingId)
+		bz := startStore.Get(key)
+		node := &types.PendingInstallNode{}
+		k.cdc.MustUnmarshal(bz, node)
+
+		if deadlineSeconds < 0 || currentSeconds-node.StartTime < deadlineSeconds {
+			if deadlineBlocks < 0 || currentBlocks-node.StartBlock < deadlineBlocks {
+				// Still alive.  Stop the search.
+				break
+			}
+		}
+
+		// This pending bundle install is dead.  Remove it.
+		k.SetPendingBundleInstall(ctx, pendingId, nil)
+
+		// Advance to the next node.
+		state.FirstPendingId = node.NextId
+	}
+
+	k.SetState(ctx, state)
+}
+
+func (k Keeper) makeListTools(ctx sdk.Context) *types.ListTools {
+	store := ctx.KVStore(k.storeKey)
+	listStore := prefix.NewStore(store, []byte(pendingNodeKeyPrefix))
+	return types.NewListTools(ctx, listStore, k.cdc)
+}
+
+func (k Keeper) SetPendingBundleInstall(ctx sdk.Context, pendingId uint64, newMsg *types.MsgInstallBundle) {
+	store := ctx.KVStore(k.storeKey)
+	pendingStore := prefix.NewStore(store, []byte(pendingBundleInstallKeyPrefix))
+
+	key := sdk.Uint64ToBigEndian(pendingId)
+	if newMsg != nil {
+		bz := k.cdc.MustMarshal(newMsg)
+		pendingStore.Set(key, bz)
+		return
+	}
+
+	var msg types.MsgInstallBundle
+	k.cdc.MustUnmarshal(pendingStore.Get(key), &msg)
+	if msg.BundleChunks != nil && len(msg.BundleChunks.Chunks) > 0 {
+		// Remove the chunks.
+		for i := range msg.BundleChunks.Chunks {
+			k.SetPendingChunkData(ctx, pendingId, uint64(i), nil)
+		}
+	}
+	pendingStore.Delete(key)
+
+	// Remove auxilliary data structure entries.
+	k.RemovePendingInstallNode(ctx, pendingId)
+}
+
+// RemovePendingInstallNode removes this pending install from the keeper's
+// ordered linked list structures, and deletes it from the store.
+func (k Keeper) RemovePendingInstallNode(ctx sdk.Context, pendingId uint64) types.PendingInstallNode {
+	lt := k.makeListTools(ctx)
+
+	victimKey := lt.Key(pendingId)
+	victimNode := lt.Fetch(victimKey)
+
+	// Remove the victim from the linked list, keeping the structure intact.
+	lt.Unlink(victimNode, func(firstp, lastp *uint64) {
+		state := k.GetState(ctx)
+		if firstp != nil {
+			state.FirstPendingId = *firstp
+		}
+		if lastp != nil {
+			state.LastPendingId = *lastp
+		}
+		k.SetState(ctx, state)
+	})
+
+	// Finally, delete the victim node's storage.
+	lt.Delete(victimKey)
+	return victimNode
+}
+
 func (k Keeper) GetSwingStore(ctx sdk.Context) sdk.KVStore {
 	store := ctx.KVStore(k.storeKey)
 	return prefix.NewStore(store, []byte(swingStoreKeyPrefix))

golang/cosmos/x/swingset/keeper/msg_server.go

@@ -1,7 +1,11 @@
 package keeper
 
 import (
+	"bytes"
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"fmt"
 
 	"github.com/Agoric/agoric-sdk/golang/cosmos/vm"
 	"github.com/Agoric/agoric-sdk/golang/cosmos/x/swingset/types"
@@ -190,6 +194,26 @@ type installBundleAction struct {
 
 func (keeper msgServer) InstallBundle(goCtx context.Context, msg *types.MsgInstallBundle) (*types.MsgInstallBundleResponse, error) {
 	ctx := sdk.UnwrapSDKContext(goCtx)
+	if msg.BundleChunks == nil || len(msg.BundleChunks.Chunks) == 0 {
+		return keeper.InstallFinishedBundle(goCtx, msg)
+	}
+
+	// Mark all the chunks as in-flight.
+	bc := *msg.BundleChunks
+	bc.Chunks = make([]*types.ChunkInfo, len(bc.Chunks))
+	for i, chunk := range bc.Chunks {
+		ci := *chunk
+		ci.State = types.ChunkState_IN_FLIGHT
+		bc.Chunks[i] = &ci
+	}
+	msg.BundleChunks = &bc
+
+	pendingId := keeper.AddPendingBundleInstall(ctx, msg)
+	return &types.MsgInstallBundleResponse{PendingId: pendingId}, nil
+}
+
+func (keeper msgServer) InstallFinishedBundle(goCtx context.Context, msg *types.MsgInstallBundle) (*types.MsgInstallBundleResponse, error) {
+	ctx := sdk.UnwrapSDKContext(goCtx)
 
 	err := msg.Uncompress()
 	if err != nil {
@@ -207,3 +231,102 @@ func (keeper msgServer) InstallBundle(goCtx context.Context, msg *types.MsgInsta
 
 	return &types.MsgInstallBundleResponse{}, nil
 }
+
+func (keeper msgServer) SendChunk(goCtx context.Context, msg *types.MsgSendChunk) (*types.MsgSendChunkResponse, error) {
+	ctx := sdk.UnwrapSDKContext(goCtx)
+
+	inst := keeper.GetPendingBundleInstall(ctx, msg.PendingId)
+	if inst == nil {
+		return nil, fmt.Errorf("no upload in progress for pending Id %d", msg.PendingId)
+	}
+
+	bc := inst.BundleChunks
+
+	if msg.ChunkIndex < 0 || msg.ChunkIndex >= uint64(len(bc.Chunks)) {
+		return nil, fmt.Errorf("chunk index %d out of range for pending Id %d", msg.ChunkIndex, msg.PendingId)
+	}
+
+	if bc.Chunks[msg.ChunkIndex].State != types.ChunkState_IN_FLIGHT {
+		return nil, fmt.Errorf("chunk %d is not in flight for pending Id %d", msg.ChunkIndex, msg.PendingId)
+	}
+
+	// Verify the chunk data.
+	ci := bc.Chunks[msg.ChunkIndex]
+	if ci.ChunkSize != uint64(len(msg.ChunkData)) {
+		return nil, fmt.Errorf("chunk %d size mismatch for pending Id %d", msg.ChunkIndex, msg.PendingId)
+	}
+
+	sha256Hash, err := hex.DecodeString(ci.Hash)
+	if err != nil {
+		return nil, fmt.Errorf("chunk %d cannot decode hash %s: %w", msg.ChunkIndex, ci.Hash, err)
+	}
+
+	hasher := sha256.New()
+	sum := hasher.Sum(msg.ChunkData)
+	if !bytes.Equal(sum, sha256Hash) {
+		return nil, fmt.Errorf("chunk %d hash mismatch; expected %x, got %x", msg.ChunkIndex, sha256Hash, sum)
+	}
+
+	// Data is valid, so store it.
+	keeper.SetPendingChunkData(ctx, msg.PendingId, msg.ChunkIndex, msg.ChunkData)
+
+	// Mark the chunk as received, and store the pending installation.
+	ci.State = types.ChunkState_RECEIVED
+	keeper.SetPendingBundleInstall(ctx, msg.PendingId, inst)
+
+	err = keeper.MaybeFinalizeBundle(ctx, msg.PendingId)
+
+	return &types.MsgSendChunkResponse{
+		PendingId: msg.PendingId,
+		Chunk:     ci,
+	}, err
+}
+
+func (keeper msgServer) MaybeFinalizeBundle(ctx sdk.Context, pendingId uint64) error {
+	msg := keeper.GetPendingBundleInstall(ctx, pendingId)
+	if msg == nil {
+		return nil
+	}
+
+	// If any chunks are not received, then bail (without error).
+	bc := msg.BundleChunks
+	totalChunkSize := uint64(0)
+	for _, chunk := range bc.Chunks {
+		if chunk.State != types.ChunkState_RECEIVED {
+			return nil
+		}
+		totalChunkSize += chunk.ChunkSize
+	}
+
+	chunkData := make([]byte, 0, totalChunkSize)
+	for i := range bc.Chunks {
+		bz := keeper.GetPendingChunkData(ctx, pendingId, uint64(i))
+		chunkData = append(chunkData, bz...)
+	}
+
+	// Verify the hash of the concatenated chunks.
+	hasher := sha256.New()
+	sum := hasher.Sum(chunkData)
+	sha256Hash, err := hex.DecodeString(bc.BundleHash)
+	if err != nil {
+		return fmt.Errorf("cannot decode hash %s: %w", bc.BundleHash, err)
+	}
+	if !bytes.Equal(sum, sha256Hash) {
+		return fmt.Errorf("bundle hash mismatch; expected %x, got %x", sha256Hash, sum)
+	}
+
+	// Is it compressed or not?
+	if msg.UncompressedSize > 0 {
+		msg.CompressedBundle = chunkData
+	} else {
+		msg.Bundle = string(chunkData)
+	}
+
+	// Clean up the pending installation state.
+	msg.BundleChunks = nil
+	keeper.SetPendingBundleInstall(ctx, pendingId, nil)
+
+	// Install the bundle now that all the chunks are processed.
+	_, err = keeper.InstallFinishedBundle(ctx, msg)
+	return err
+}