Add GetBalance endpoint

allnil · Mar 14, 2024 · 786d4b9 · 786d4b9
1 parent d20d30c
commit 786d4b9
Show file tree

Hide file tree

Showing 7 changed files with 59 additions and 0 deletions.
diff --git a/miner/builder.go b/miner/builder.go
@@ -99,6 +99,10 @@ func (b *Builder) AddTransaction(txn *types.Transaction) (*suavextypes.SimulateT
 	return receiptToSimResult(&types.Receipt{Logs: logs}), 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

diff --git a/miner/builder_test.go b/miner/builder_test.go
@@ -124,6 +124,26 @@ func TestBuilder_Bid(t *testing.T) {
 	fmt.Println("-- req --", req)
 }
 
+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()

diff --git a/suave/builder/api/api.go b/suave/builder/api/api.go
@@ -79,4 +79,5 @@ type API interface {
 	AddTransaction(ctx context.Context, sessionId string, tx *types.Transaction) (*SimulateTransactionResult, 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)
 }
diff --git a/suave/builder/api/api_client.go b/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"
 )
@@ -51,3 +53,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
+}
diff --git a/suave/builder/api/api_server.go b/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"
 )
 
@@ -15,6 +17,7 @@ type SessionManager interface {
 	AddTransaction(sessionId string, tx *types.Transaction) (*SimulateTransactionResult, 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 {
@@ -44,6 +47,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 {
 }
@@ -59,3 +66,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
+}
diff --git a/suave/builder/api/api_test.go b/suave/builder/api/api_test.go
@@ -27,6 +27,9 @@ 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)
 }
 
 type nullSessionManager struct{}
@@ -50,3 +53,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
+}
diff --git a/suave/builder/session_manager.go b/suave/builder/session_manager.go
@@ -167,6 +167,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.