Add Getter to BalanceHandler (#1685)

ava-labs · Oct 25, 2024 · 0459bd4 · 0459bd4
1 parent 960b790
commit 0459bd4
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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)
+}