Add validateInitCode to userOp checks (#67)

Author: hazim-j

internal/testutils/constants.go

@@ -0,0 +1,8 @@
+package testutils
+
+import "math/big"
+
+var (
+	OneETH                 = big.NewInt(1000000000000000000)
+	DefaultUnstakeDelaySec = uint32(86400)
+)

internal/testutils/opmock.go

@@ -1,6 +1,9 @@
 package testutils
 
-import "github.com/stackup-wallet/stackup-bundler/pkg/userop"
+import (
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/stackup-wallet/stackup-bundler/pkg/userop"
+)
 
 var (
 	MockUserOpData = map[string]any{
@@ -16,7 +19,7 @@ var (
 		"preVerificationGas":   "0xc650",
 		"signature":            "0xa925dcc5e5131636e244d4405334c25f034ebdd85c0cb12e8cdb13c15249c2d466d0bade18e2cafd3513497f7f968dcbb63e519acd9b76dcae7acd61f11aa8421b",
 	}
-	MockByteCode = "0x6080604052"
+	MockByteCode = common.Hex2Bytes("6080604052")
 )
 
 // Returns a valid initial userOperation for an EIP-4337 account.

pkg/modules/checks/initcode.go

@@ -0,0 +1,28 @@
+package checks
+
+import (
+	"errors"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/stackup-wallet/stackup-bundler/pkg/userop"
+)
+
+// ValidateInitCode checks if initCode is not empty and gets the factory address. If factory address is valid
+// it calls a generic function that can retrieve the stake from the EntryPoint.
+func ValidateInitCode(op *userop.UserOperation, gs GetStakeFunc) error {
+	if len(op.InitCode) == 0 {
+		return nil
+	}
+
+	f := op.GetFactory()
+	if f == common.HexToAddress("0x") {
+		return errors.New("initCode: does not contain a valid address")
+	}
+
+	_, err := gs(f)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

pkg/modules/checks/initcode_test.go

@@ -0,0 +1,67 @@
+package checks
+
+import (
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/stackup-wallet/stackup-bundler/internal/testutils"
+	"github.com/stackup-wallet/stackup-bundler/pkg/entrypoint"
+)
+
+// TestInitCodeDNE calls checks.ValidateInitCode where initCode does not exist. Expect nil.
+func TestInitCodeDNE(t *testing.T) {
+	op := testutils.MockValidInitUserOp()
+	op.InitCode = []byte{}
+	err := ValidateInitCode(op, func(f common.Address) (*entrypoint.IStakeManagerDepositInfo, error) {
+		return nil, nil
+	})
+
+	if err != nil {
+		t.Fatalf(`got err %v, want nil`, err)
+	}
+}
+
+// TestInitCodeContainsAddress calls checks.ValidateInitCode where initCode exist without a valid address.
+// Expect error.
+func TestInitCodeContainsAddress(t *testing.T) {
+	op := testutils.MockValidInitUserOp()
+	op.InitCode = []byte("1234")
+	err := ValidateInitCode(op, func(f common.Address) (*entrypoint.IStakeManagerDepositInfo, error) {
+		return nil, nil
+	})
+
+	if err == nil {
+		t.Fatalf("got nil, want err")
+	}
+}
+
+// TestGetStakeFuncReceivesFactory calls checks.ValidateInitCode where initCode exist and calls getStakeFunc
+// with the correct factory address.
+func TestGetStakeFuncReceivesFactory(t *testing.T) {
+	op := testutils.MockValidInitUserOp()
+	isCalled := false
+	_ = ValidateInitCode(op, func(f common.Address) (*entrypoint.IStakeManagerDepositInfo, error) {
+		if f != op.GetFactory() {
+			t.Fatalf("got %s, want %s", f.String(), op.GetFactory())
+		}
+
+		isCalled = true
+		return nil, nil
+	})
+
+	if !isCalled {
+		t.Fatalf("getStakeFunc was not called")
+	}
+}
+
+// TestInitCodeExists calls checks.ValidateInitCode where valid initCode does exist. Expect nil.
+func TestInitCodeExists(t *testing.T) {
+	op := testutils.MockValidInitUserOp()
+	err := ValidateInitCode(op, func(f common.Address) (*entrypoint.IStakeManagerDepositInfo, error) {
+		return nil, nil
+	})
+
+	if err != nil {
+		t.Fatalf(`got err %v, want nil`, err)
+	}
+}

pkg/modules/checks/sender.go

@@ -1,16 +1,15 @@
 package checks
 
 import (
-	"context"
 	"errors"
 
-	"github.com/ethereum/go-ethereum/ethclient"
 	"github.com/stackup-wallet/stackup-bundler/pkg/userop"
 )
 
-// Checks that the sender is an existing contract, or the initCode is not empty (but not both)
-func checkSender(eth *ethclient.Client, op *userop.UserOperation) error {
-	bytecode, err := eth.CodeAt(context.Background(), op.Sender, nil)
+// ValidateSender accepts a userOp and a generic function that can retrieve the bytecode of the sender.
+// Either the sender is deployed (non-zero length bytecode) or the initCode is not empty (but not both).
+func ValidateSender(op *userop.UserOperation, gc GetCodeFunc) error {
+	bytecode, err := gc(op.Sender)
 	if err != nil {
 		return err
 	}

pkg/modules/checks/sender_test.go

@@ -3,68 +3,60 @@ package checks
 import (
 	"testing"
 
-	"github.com/ethereum/go-ethereum/ethclient"
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/stackup-wallet/stackup-bundler/internal/testutils"
 )
 
-// TestSenderExistAndInitCodeDNE calls checks.CheckSender where sender contract exist and initCode does not.
-// Expect nil.
+// TestSenderExistAndInitCodeDNE calls checks.ValidateSender where sender contract exist and initCode does
+// not. Expect nil.
 func TestSenderExistAndInitCodeDNE(t *testing.T) {
-	server := testutils.EthMock(testutils.MethodMocks{
-		"eth_getCode": testutils.MockByteCode,
-	})
-	defer server.Close()
-
-	eth, _ := ethclient.Dial(server.URL)
 	op := testutils.MockValidInitUserOp()
 	op.InitCode = []byte{}
-	if err := checkSender(eth, op); err != nil {
+	err := ValidateSender(op, func(s common.Address) ([]byte, error) {
+		return testutils.MockByteCode, nil
+	})
+
+	if err != nil {
 		t.Fatalf(`got err %v, want nil`, err)
 	}
 }
 
-// TestSenderAndInitCodeExist calls checks.CheckSender where sender contract and initCode exist. Expect
+// TestSenderAndInitCodeExist calls checks.ValidateSender where sender contract and initCode exist. Expect
 // error.
 func TestSenderAndInitCodeExist(t *testing.T) {
-	server := testutils.EthMock(testutils.MethodMocks{
-		"eth_getCode": testutils.MockByteCode,
+	op := testutils.MockValidInitUserOp()
+	err := ValidateSender(op, func(s common.Address) ([]byte, error) {
+		return testutils.MockByteCode, nil
 	})
-	defer server.Close()
 
-	eth, _ := ethclient.Dial(server.URL)
-	op := testutils.MockValidInitUserOp()
-	if err := checkSender(eth, op); err == nil {
+	if err == nil {
 		t.Fatalf(`got nil, want err`)
 	}
 }
 
-// TestSenderDNEAndInitCodeExist calls checks.CheckSender where sender contract does not exist and initCode
-// does. Expect nil.
+// TestSenderDNEAndInitCodeExist calls checks.ValidateSender where sender contract does not exist and
+// initCode does. Expect nil.
 func TestSenderDNEAndInitCodeExist(t *testing.T) {
-	server := testutils.EthMock(testutils.MethodMocks{
-		"eth_getCode": "0x",
+	op := testutils.MockValidInitUserOp()
+	err := ValidateSender(op, func(s common.Address) ([]byte, error) {
+		return []byte{}, nil
 	})
-	defer server.Close()
 
-	eth, _ := ethclient.Dial(server.URL)
-	op := testutils.MockValidInitUserOp()
-	if err := checkSender(eth, op); err != nil {
+	if err != nil {
 		t.Fatalf(`got err %v, want nil`, err)
 	}
 }
 
-// TestSenderAndInitCodeDNE calls checks.CheckSender where sender contract and initCode does not exist.
+// TestSenderAndInitCodeDNE calls checks.ValidateSender where sender contract and initCode does not exist.
 // Expect error.
 func TestSenderAndInitCodeDNE(t *testing.T) {
-	server := testutils.EthMock(testutils.MethodMocks{
-		"eth_getCode": "0x",
-	})
-	defer server.Close()
-
-	eth, _ := ethclient.Dial(server.URL)
 	op := testutils.MockValidInitUserOp()
 	op.InitCode = []byte{}
-	if err := checkSender(eth, op); err == nil {
+	err := ValidateSender(op, func(s common.Address) ([]byte, error) {
+		return []byte{}, nil
+	})
+
+	if err == nil {
 		t.Fatalf(`got nil, want err`)
 	}
 }

pkg/modules/checks/standalone.go

@@ -39,8 +39,15 @@ func (s *Standalone) ValidateOpValues() modules.UserOpHandlerFunc {
 			return err
 		}
 
+		getCode := getCodeWithEthClient(s.eth)
+		getStake, err := getStakeWithEthClient(ctx, s.eth)
+		if err != nil {
+			return err
+		}
+
 		g := new(errgroup.Group)
-		g.Go(func() error { return checkSender(s.eth, ctx.UserOp) })
+		g.Go(func() error { return ValidateSender(ctx.UserOp, getCode) })
+		g.Go(func() error { return ValidateInitCode(ctx.UserOp, getStake) })
 		g.Go(func() error { return checkVerificationGas(s.maxVerificationGas, ctx.UserOp) })
 		g.Go(func() error { return checkPaymasterAndData(s.eth, ep, ctx.UserOp) })
 		g.Go(func() error { return checkCallGasLimit(ctx.UserOp) })

pkg/modules/checks/utils.go

@@ -0,0 +1,42 @@
+package checks
+
+import (
+	"context"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/ethclient"
+	"github.com/stackup-wallet/stackup-bundler/pkg/entrypoint"
+	"github.com/stackup-wallet/stackup-bundler/pkg/modules"
+)
+
+// GetCodeFunc provides a general interface for retrieving the bytecode for a given address.
+type GetCodeFunc = func(addr common.Address) ([]byte, error)
+
+// GetStakeFunc provides a general interface for retrieving the EntryPoint stake for a given address.
+type GetStakeFunc = func(entity common.Address) (*entrypoint.IStakeManagerDepositInfo, error)
+
+// getCodeWithEthClient returns a GetCodeFunc that uses an eth client to call eth_getCode.
+func getCodeWithEthClient(eth *ethclient.Client) GetCodeFunc {
+	return func(addr common.Address) ([]byte, error) {
+		return eth.CodeAt(context.Background(), addr, nil)
+	}
+}
+
+// getStakeWithEthClient returns a GetStakeFunc that uses an EntryPoint binding to get stake info and adds it
+// to the current context.
+func getStakeWithEthClient(ctx *modules.UserOpHandlerCtx, eth *ethclient.Client) (GetStakeFunc, error) {
+	ep, err := entrypoint.NewEntrypoint(ctx.EntryPoint, eth)
+	if err != nil {
+		return nil, err
+	}
+
+	return func(addr common.Address) (*entrypoint.IStakeManagerDepositInfo, error) {
+		dep, err := ep.GetDepositInfo(nil, addr)
+		if err != nil {
+			return nil, err
+		}
+
+		ctx.AddDepositInfo(addr, &dep)
+		return &dep, nil
+	}, nil
+}

pkg/modules/context.go

@@ -4,9 +4,21 @@ import (
 	"math/big"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/stackup-wallet/stackup-bundler/pkg/entrypoint"
 	"github.com/stackup-wallet/stackup-bundler/pkg/userop"
 )
 
+// BatchHandlerCtx is the object passed to BatchHandler functions during the Bundler's Run process. It
+// also contains a Data field for adding arbitrary key-value pairs to the context. These values will be
+// logged by the Bundler at the end of each run.
+type BatchHandlerCtx struct {
+	Batch          []*userop.UserOperation
+	PendingRemoval []*userop.UserOperation
+	EntryPoint     common.Address
+	ChainID        *big.Int
+	Data           map[string]any
+}
+
 // NewBatchHandlerContext creates a new BatchHandlerCtx using a copy of the given batch.
 func NewBatchHandlerContext(batch []*userop.UserOperation, entryPoint common.Address, chainID *big.Int) *BatchHandlerCtx {
 	var copy []*userop.UserOperation
@@ -21,11 +33,52 @@ func NewBatchHandlerContext(batch []*userop.UserOperation, entryPoint common.Add
 	}
 }
 
+// MarkOpIndexForRemoval will remove the op by index from the batch and add it to the pending removal array.
+// This should be used for ops that are not to be included on-chain and dropped from the mempool.
+func (c *BatchHandlerCtx) MarkOpIndexForRemoval(index int) {
+	batch := []*userop.UserOperation{}
+	var op *userop.UserOperation
+	for i, curr := range c.Batch {
+		if i == index {
+			op = curr
+		} else {
+			batch = append(batch, curr)
+		}
+	}
+	if op == nil {
+		return
+	}
+
+	c.Batch = batch
+	c.PendingRemoval = append(c.PendingRemoval, op)
+}
+
+// UserOpHandlerCtx is the object passed to UserOpHandler functions during the Client's SendUserOperation
+// process.
+type UserOpHandlerCtx struct {
+	UserOp     *userop.UserOperation
+	EntryPoint common.Address
+	ChainID    *big.Int
+	deposits   map[common.Address]*entrypoint.IStakeManagerDepositInfo
+}
+
 // NewUserOpHandlerContext creates a new UserOpHandlerCtx using a given op.
 func NewUserOpHandlerContext(op *userop.UserOperation, entryPoint common.Address, chainID *big.Int) *UserOpHandlerCtx {
 	return &UserOpHandlerCtx{
 		UserOp:     op,
 		EntryPoint: entryPoint,
 		ChainID:    chainID,
+		deposits:   make(map[common.Address]*entrypoint.IStakeManagerDepositInfo),
 	}
 }
+
+// AddDepositInfo adds any entity's EntryPoint stake info to the current context.
+func (c *UserOpHandlerCtx) AddDepositInfo(entity common.Address, dep *entrypoint.IStakeManagerDepositInfo) {
+	c.deposits[entity] = dep
+}
+
+// GetDepositInfo retrieves any entity's EntryPoint stake info from the current context if it was previously
+// added. Otherwise returns nil
+func (c *UserOpHandlerCtx) GetDepositInfo(entity common.Address) *entrypoint.IStakeManagerDepositInfo {
+	return c.deposits[entity]
+}