fixup! feat(x/swingset): auto-provision smart wallet

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

@@ -80,7 +80,7 @@ type walletAction struct {
 func (keeper msgServer) WalletAction(goCtx context.Context, msg *types.MsgWalletAction) (*types.MsgWalletActionResponse, error) {
 	ctx := sdk.UnwrapSDKContext(goCtx)
 
-	err := keeper.provisionIfNeeded(ctx)
+	err := keeper.provisionIfNeeded(ctx, msg.Owner)
 	if err != nil {
 		return nil, err
 	}
@@ -113,7 +113,7 @@ type walletSpendAction struct {
 func (keeper msgServer) WalletSpendAction(goCtx context.Context, msg *types.MsgWalletSpendAction) (*types.MsgWalletSpendActionResponse, error) {
 	ctx := sdk.UnwrapSDKContext(goCtx)
 
-	err := keeper.provisionIfNeeded(ctx)
+	err := keeper.provisionIfNeeded(ctx, msg.Owner)
 	if err != nil {
 		return nil, err
 	}
@@ -141,22 +141,37 @@ type provisionAction struct {
 	AutoProvision bool   `json:"autoProvision"`
 }
 
-// provisionIfNeeded generates a provision action if needed for the transaction
-func (keeper msgServer) provisionIfNeeded(ctx sdk.Context) error {
-	msg, txShouldProvision := ctx.Context().Value(types.TxShouldProvision).(types.MsgProvision)
-	if !txShouldProvision {
+// provisionIfNeeded generates a provision action if no smart wallet is already
+// provisioned for the account. This assumes that all messages for
+// non-provisioned smart wallets allowed by the admission AnteHandler should
+// auto-provision the smart wallet.
+func (keeper msgServer) provisionIfNeeded(ctx sdk.Context, owner sdk.AccAddress) error {
+	// TODO: If/when we implement marking of pending smart wallet provisions,
+	// here we'll need to generate a provision action until the smart wallet has
+	// been fully provisioned by the controller. This is because a provision is
+	// not guaranteed to succeed (e.g. lack of provision pool funds)
+	provisioned, err := keeper.HasSmartWallet(ctx, owner)
+	if err != nil {
+		return err
+	} else if provisioned {
 		return nil
 	}
 
+	msg := &types.MsgProvision{
+		Address:    owner,
+		Submitter:  owner,
+		PowerFlags: []string{types.PowerFlagSmartWallet},
+	}
+
 	action := &provisionAction{
-		MsgProvision:  &msg,
+		MsgProvision:  msg,
 		Type:          "PLEASE_PROVISION",
 		BlockHeight:   ctx.BlockHeight(),
 		BlockTime:     ctx.BlockTime().Unix(),
 		AutoProvision: true,
 	}
 
-	err := keeper.routeAction(ctx, &msg, action)
+	err = keeper.routeAction(ctx, msg, action)
 	// fmt.Fprintln(os.Stderr, "Returned from SwingSet", out, err)
 	if err != nil {
 		return err

golang/cosmos/x/swingset/types/msgs.go

@@ -3,7 +3,6 @@ package types
 import (
 	"bytes"
 	"compress/gzip"
-	"context"
 	"encoding/json"
 	"io"
 	"strings"
@@ -15,8 +14,6 @@ import (
 
 const RouterKey = ModuleName // this was defined in your key.go file
 
-const TxShouldProvision = sdk.ContextKey("tx-should-provision")
-
 var (
 	_ sdk.Msg = &MsgDeliverInbound{}
 	_ sdk.Msg = &MsgProvision{}
@@ -52,49 +49,42 @@ func chargeAdmission(ctx sdk.Context, keeper SwingSetKeeper, addr sdk.AccAddress
 }
 
 // checkSmartWalletProvisioned verifies if a smart wallet message can be
-// delivered for the owner's address.
-func checkSmartWalletProvisioned(ctx sdk.Context, keeper SwingSetKeeper, addr sdk.AccAddress) (sdk.Context, error) {
-	// While multiple swingset messages are not currently allowed in the same
-	// transaction, if that is ever relaxed, we would end up generating multiple
-	// provision actions, and the provisioning code is not fully idempotent.
-	_, txShouldProvision := ctx.Context().Value(TxShouldProvision).(MsgProvision)
-	if txShouldProvision {
-		return ctx, sdkerrors.Wrapf(sdkerrors.ErrInvalidRequest, "cannot auto-provision for multiple smart wallet messages")
-	}
-
+// delivered for the owner's address. A messaged is allowed if a smart wallet
+// is already provisioned for the address, or if the provisioning fee is
+// charged successfully.
+// All messages for non-provisioned smart wallets allowed here will result in
+// an auto-provision action generated by the msg server.
+func checkSmartWalletProvisioned(ctx sdk.Context, keeper SwingSetKeeper, addr sdk.AccAddress) error {
 	// This checks if a smart wallet has already been provisioned by the controller
 	// However a provision (either explicit or automatic) may be pending execution
 	// and this check will not catch that case, resulting in the owner being charged
 	// for provisioning again.
+	// Furthermore, while multiple swingset messages are not currently allowed in
+	// the same transaction, if that is ever relaxed, we would end up generating
+	// multiple provision actions, and the provisioning code is not fully
+	// idempotent. However this is only a degenerate case of the multiple txs case,
+	// and as such is charged again similarly.
 	// TODO: mark that a smart wallet provision is pending. However in that case,
 	// auto-provisioning should still be performed (but without fees being charged),
 	// until the controller actually provisions the smart wallet (the operation may
 	// transiently fail, requiring retries until success)
 	isProvisioned, err := keeper.HasSmartWallet(ctx, addr)
 	if err != nil {
-		return ctx, err
-	} else if isProvisioned {
-		return ctx, nil
+		return err
 	}
 
-	beansPerUnit := keeper.GetBeansPerUnit(ctx)
-	beans := beansPerUnit[BeansPerSmartWalletProvision]
-
-	// This is a separate charge from the smart wallet action which triggered the check
-	err = keeper.ChargeBeans(ctx, addr, beans)
-	if err != nil {
-		return ctx, err
-	}
+	if !isProvisioned {
+		beansPerUnit := keeper.GetBeansPerUnit(ctx)
+		beans := beansPerUnit[BeansPerSmartWalletProvision]
 
-	// Track on the context that the smart wallet requires auto-provisioning
-	provisionMsg := MsgProvision{
-		Address:    addr,
-		Submitter:  addr,
-		PowerFlags: []string{PowerFlagSmartWallet},
+		// This is a separate charge from the smart wallet action which triggered the check
+		err = keeper.ChargeBeans(ctx, addr, beans)
+		if err != nil {
+			return err
+		}
 	}
-	newCtx := ctx.WithContext(context.WithValue(ctx.Context(), TxShouldProvision, provisionMsg))
 
-	return newCtx, nil
+	return nil
 }
 
 func NewMsgDeliverInbound(msgs *Messages, submitter sdk.AccAddress) *MsgDeliverInbound {
@@ -186,7 +176,7 @@ func (msg MsgWalletAction) CheckAdmissibility(ctx sdk.Context, data interface{})
 		return ctx, sdkerrors.Wrapf(sdkerrors.ErrInvalidRequest, "data must be a SwingSetKeeper, not a %T", data)
 	}
 
-	ctx, err := checkSmartWalletProvisioned(ctx, keeper, msg.Owner)
+	err := checkSmartWalletProvisioned(ctx, keeper, msg.Owner)
 	if err != nil {
 		return ctx, err
 	}
@@ -258,7 +248,7 @@ func (msg MsgWalletSpendAction) CheckAdmissibility(ctx sdk.Context, data interfa
 		return ctx, sdkerrors.Wrapf(sdkerrors.ErrInvalidRequest, "data must be a SwingSetKeeper, not a %T", data)
 	}
 
-	ctx, err := checkSmartWalletProvisioned(ctx, keeper, msg.Owner)
+	err := checkSmartWalletProvisioned(ctx, keeper, msg.Owner)
 	if err != nil {
 		return ctx, err
 	}
@@ -330,18 +320,9 @@ func (msg MsgProvision) ValidateBasic() error {
 
 // CheckAdmissibility implements the vm.ControllerAdmissionMsg interface.
 func (msg MsgProvision) CheckAdmissibility(ctx sdk.Context, data interface{}) (sdk.Context, error) {
-	// While multiple swingset messages are not currently allowed in the same
-	// transaction, if that is ever relaxed, it would be invalid for a provision
-	// message to appear after a smart wallet message
-	// TODO: the more thorough check would be to disallow a provision message for
-	// a smart wallet if the smart wallet is already provisioned or pending
-	// provisioning. However we/ currently do not track whether is smart wallet
-	// is pending provisioning, except in the narrow case of an auto-provisioned
-	// smart wallet in the same transaction
-	_, txShouldProvision := ctx.Context().Value(TxShouldProvision).(MsgProvision)
-	if txShouldProvision {
-		return ctx, sdkerrors.Wrap(sdkerrors.ErrInvalidRequest, "Unexpected provision message after auto-provisioning message")
-	}
+	// TODO: consider disallowing a provision message for a smart wallet if the
+	// smart wallet is already provisioned or pending provisioning. However we
+	// currently do not track whether a smart wallet is pending provisioning.
 
 	// For explicitly provisioning, swingset will take care of charging,
 	// so we skip admission fees.