Merge pull request #24 from flashbots/feature/builder-api-balance

Add GetBalance endpoint

miner/builder.go

@@ -172,6 +172,10 @@ func (b *Builder) AddBundles(bundles []*suavextypes.Bundle) ([]*suavextypes.Simu
 	return results, nil
 }
 
+func (b *Builder) GetBalance(addr common.Address) *big.Int {
+	return b.env.state.GetBalance(addr)
+}
+
 func (b *Builder) FillPending() error {
 	if err := b.wrk.commitPendingTxs(b.env); err != nil {
 		return err

miner/builder_test.go

@@ -270,6 +270,26 @@ func TestBuilder_Bid(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestBuilder_Balance(t *testing.T) {
+	t.Parallel()
+
+	config, backend := newMockBuilderConfig(t)
+
+	builder, err := NewBuilder(config, &BuilderArgs{})
+	require.NoError(t, err)
+
+	balance := builder.GetBalance(testBankAddress)
+	require.Equal(t, balance, testBankFunds)
+
+	// make a random txn that consumes gas
+	tx1 := backend.newRandomTx(true)
+	_, err = builder.AddTransaction(tx1)
+	require.NoError(t, err)
+
+	balance2 := builder.GetBalance(testBankAddress)
+	require.NotEqual(t, balance2, testBankFunds)
+}
+
 func newMockBuilderConfig(t *testing.T) (*BuilderConfig, *testWorkerBackend) {
 	var (
 		db     = rawdb.NewMemoryDatabase()

suave/builder/api/api.go

@@ -96,4 +96,5 @@ type API interface {
 	AddBundles(ctx context.Context, sessionId string, bundles []*Bundle) ([]*SimulateBundleResult, error)
 	BuildBlock(ctx context.Context, sessionId string) error
 	Bid(ctx context.Context, sessioId string, blsPubKey phase0.BLSPubKey) (*SubmitBlockRequest, error)
+	GetBalance(ctx context.Context, sessionId string, addr common.Address) (*big.Int, error)
 }

suave/builder/api/api_client.go

@@ -2,8 +2,10 @@ package api
 
 import (
 	"context"
+	"math/big"
 
 	"github.com/attestantio/go-eth2-client/spec/phase0"
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/rpc"
 )
@@ -63,3 +65,9 @@ func (a *APIClient) Bid(ctx context.Context, sessioId string, blsPubKey phase0.B
 	err := a.rpc.CallContext(ctx, &req, "suavex_bid", sessioId, blsPubKey)
 	return req, err
 }
+
+func (a *APIClient) GetBalance(ctx context.Context, sessionId string, addr common.Address) (*big.Int, error) {
+	var balance *big.Int
+	err := a.rpc.CallContext(ctx, &balance, "suavex_getBalance", sessionId, addr)
+	return balance, err
+}

suave/builder/api/api_server.go

@@ -2,8 +2,10 @@ package api
 
 import (
 	"context"
+	"math/big"
 
 	"github.com/attestantio/go-eth2-client/spec/phase0"
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
 )
 
@@ -17,6 +19,7 @@ type SessionManager interface {
 	AddBundles(sessionId string, bundles []*Bundle) ([]*SimulateBundleResult, error)
 	BuildBlock(sessionId string) error
 	Bid(sessionId string, blsPubKey phase0.BLSPubKey) (*SubmitBlockRequest, error)
+	GetBalance(sessionId string, addr common.Address) (*big.Int, error)
 }
 
 func NewServer(s SessionManager) *Server {
@@ -54,6 +57,10 @@ func (s *Server) Bid(ctx context.Context, sessionId string, blsPubKey phase0.BLS
 	return s.sessionMngr.Bid(sessionId, blsPubKey)
 }
 
+func (s *Server) GetBalance(ctx context.Context, sessionId string, addr common.Address) (*big.Int, error) {
+	return s.sessionMngr.GetBalance(sessionId, addr)
+}
+
 // TODO: Remove
 type MockServer struct {
 }
@@ -69,3 +76,7 @@ func (s *MockServer) AddTransaction(ctx context.Context, sessionId string, tx *t
 func (s *MockServer) BuildBlock(ctx context.Context) error {
 	return nil
 }
+
+func (s *MockServer) GetBalance(ctx context.Context, sessionId string, addr common.Address) (*big.Int, error) {
+	return big.NewInt(0), nil
+}

suave/builder/api/api_test.go

@@ -27,6 +27,10 @@ func TestAPI(t *testing.T) {
 	txn := types.NewTransaction(0, common.Address{}, big.NewInt(1), 1, big.NewInt(1), []byte{})
 	_, err = c.AddTransaction(context.Background(), "1", txn)
 	require.NoError(t, err)
+
+	_, err = c.GetBalance(context.Background(), "1", common.Address{})
+	require.NoError(t, err)
+
 	_, err = c.AddTransactions(context.Background(), "1", []*types.Transaction{txn})
 	require.NoError(t, err)
 
@@ -62,3 +66,7 @@ func (nullSessionManager) BuildBlock(sessionId string) error {
 func (nullSessionManager) Bid(sessioId string, blsPubKey phase0.BLSPubKey) (*SubmitBlockRequest, error) {
 	return nil, nil
 }
+
+func (nullSessionManager) GetBalance(sessionId string, addr common.Address) (*big.Int, error) {
+	return big.NewInt(0), nil
+}

suave/builder/session_manager.go

@@ -183,6 +183,14 @@ func (s *SessionManager) Bid(sessionId string, blsPubKey phase0.BLSPubKey) (*api
 	return builder.Bid(blsPubKey)
 }
 
+func (s *SessionManager) GetBalance(sessionId string, addr common.Address) (*big.Int, error) {
+	builder, err := s.getSession(sessionId)
+	if err != nil {
+		return nil, err
+	}
+	return builder.GetBalance(addr), nil
+}
+
 // CalcBaseFee calculates the basefee of the header.
 func CalcBaseFee(config *params.ChainConfig, parent *types.Header) *big.Int {
 	// If the current block is the first EIP-1559 block, return the InitialBaseFee.