From 0459bd402bfb8ae96fcaba3569fb73c4cad465de Mon Sep 17 00:00:00 2001
From: rodrigo <77309055+RodrigoVillar@users.noreply.github.com>
Date: Fri, 25 Oct 2024 12:23:04 -0400
Subject: [PATCH] Add Getter to `BalanceHandler` (#1685)

---
 api/dependencies.go                          |  1 +
 api/jsonrpc/client.go                        | 14 +++++++++
 api/jsonrpc/server.go                        | 30 ++++++++++++++++++++
 chain/dependencies.go                        |  4 +++
 examples/morpheusvm/storage/state_manager.go |  4 +++
 x/contracts/vm/storage/state_manager.go      |  4 +++
 6 files changed, 57 insertions(+)

diff --git a/api/dependencies.go b/api/dependencies.go
index ee9856ea4d..95175d7cf3 100644
--- a/api/dependencies.go
+++ b/api/dependencies.go
@@ -42,4 +42,5 @@ type VM interface {
 	GetVerifyAuth() bool
 	ReadState(ctx context.Context, keys [][]byte) ([][]byte, []error)
 	ImmutableState(ctx context.Context) (state.Immutable, error)
+	BalanceHandler() chain.BalanceHandler
 }
diff --git a/api/jsonrpc/client.go b/api/jsonrpc/client.go
index f851570704..d2253c2982 100644
--- a/api/jsonrpc/client.go
+++ b/api/jsonrpc/client.go
@@ -257,3 +257,17 @@ func (cli *JSONRPCClient) SimulateActions(ctx context.Context, actions chain.Act
 
 	return resp.ActionResults, nil
 }
+
+func (cli *JSONRPCClient) GetBalance(ctx context.Context, addr codec.Address) (uint64, error) {
+	args := &GetBalanceArgs{
+		Address: addr,
+	}
+	resp := new(GetBalanceReply)
+	err := cli.requester.SendRequest(
+		ctx,
+		"getBalance",
+		args,
+		resp,
+	)
+	return resp.Balance, err
+}
diff --git a/api/jsonrpc/server.go b/api/jsonrpc/server.go
index a2bdca155a..88b7bf94c7 100644
--- a/api/jsonrpc/server.go
+++ b/api/jsonrpc/server.go
@@ -318,3 +318,33 @@ func (j *JSONRPCServer) SimulateActions(
 	}
 	return nil
 }
+
+type GetBalanceArgs struct {
+	Address codec.Address `json:"address"`
+}
+
+type GetBalanceReply struct {
+	Balance uint64 `json:"balance"`
+}
+
+func (j *JSONRPCServer) GetBalance(
+	req *http.Request,
+	args *GetBalanceArgs,
+	reply *GetBalanceReply,
+) error {
+	ctx, span := j.vm.Tracer().Start(req.Context(), "JSONRPCServer.GetBalance")
+	defer span.End()
+
+	im, err := j.vm.ImmutableState(ctx)
+	if err != nil {
+		return err
+	}
+
+	balance, err := j.vm.BalanceHandler().GetBalance(ctx, args.Address, im)
+	if err != nil {
+		return err
+	}
+
+	reply.Balance = balance
+	return nil
+}
diff --git a/chain/dependencies.go b/chain/dependencies.go
index a69809d160..22839866e1 100644
--- a/chain/dependencies.go
+++ b/chain/dependencies.go
@@ -172,6 +172,10 @@ type BalanceHandler interface {
 
 	// AddBalance adds [amount] to [addr].
 	AddBalance(ctx context.Context, addr codec.Address, mu state.Mutable, amount uint64, createAccount bool) error
+
+	// GetBalance returns the balance of [addr].
+	// If [addr] does not exist, this should return 0 and no error.
+	GetBalance(ctx context.Context, addr codec.Address, im state.Immutable) (uint64, error)
 }
 
 type Object interface {
diff --git a/examples/morpheusvm/storage/state_manager.go b/examples/morpheusvm/storage/state_manager.go
index fc997245d0..c5d9cbc8cd 100644
--- a/examples/morpheusvm/storage/state_manager.go
+++ b/examples/morpheusvm/storage/state_manager.go
@@ -57,3 +57,7 @@ func (*BalanceHandler) AddBalance(
 	_, err := AddBalance(ctx, mu, addr, amount, createAccount)
 	return err
 }
+
+func (*BalanceHandler) GetBalance(ctx context.Context, addr codec.Address, im state.Immutable) (uint64, error) {
+	return GetBalance(ctx, im, addr)
+}
diff --git a/x/contracts/vm/storage/state_manager.go b/x/contracts/vm/storage/state_manager.go
index fc997245d0..c5d9cbc8cd 100644
--- a/x/contracts/vm/storage/state_manager.go
+++ b/x/contracts/vm/storage/state_manager.go
@@ -57,3 +57,7 @@ func (*BalanceHandler) AddBalance(
 	_, err := AddBalance(ctx, mu, addr, amount, createAccount)
 	return err
 }
+
+func (*BalanceHandler) GetBalance(ctx context.Context, addr codec.Address, im state.Immutable) (uint64, error) {
+	return GetBalance(ctx, im, addr)
+}