Merge pull request #6028 from AndriiDiachuk/expose-GetAccountBalance-…

…GetAccountKeys-script-execution-module

[Access] Expose GetAccountBalance and GetAccountKeys in script execution module

engine/access/rpc/backend/script_executor.go

@@ -102,6 +102,30 @@ func (s *ScriptExecutor) GetAccountAtBlockHeight(ctx context.Context, address fl
 	return s.scriptExecutor.GetAccountAtBlockHeight(ctx, address, height)
 }
 
+// GetAccountBalance returns
+// Expected errors:
+// - Script execution related errors
+// - storage.ErrHeightNotIndexed if the data for the block height is not available
+func (s *ScriptExecutor) GetAccountBalance(ctx context.Context, address flow.Address, height uint64) (uint64, error) {
+	if err := s.checkDataAvailable(height); err != nil {
+		return 0, err
+	}
+
+	return s.scriptExecutor.GetAccountBalance(ctx, address, height)
+}
+
+// GetAccountKeys returns
+// Expected errors:
+// - Script execution related errors
+// - storage.ErrHeightNotIndexed if the data for the block height is not available
+func (s *ScriptExecutor) GetAccountKeys(ctx context.Context, address flow.Address, height uint64) ([]flow.AccountPublicKey, error) {
+	if err := s.checkDataAvailable(height); err != nil {
+		return nil, err
+	}
+
+	return s.scriptExecutor.GetAccountKeys(ctx, address, height)
+}
+
 func (s *ScriptExecutor) checkDataAvailable(height uint64) error {
 	if !s.initialized.Load() {
 		return fmt.Errorf("%w: script executor not initialized", storage.ErrHeightNotIndexed)

engine/execution/computation/query/executor.go

@@ -50,6 +50,26 @@ type Executor interface {
 		*flow.Account,
 		error,
 	)
+
+	GetAccountBalance(
+		ctx context.Context,
+		addr flow.Address,
+		header *flow.Header,
+		snapshot snapshot.StorageSnapshot,
+	) (
+		uint64,
+		error,
+	)
+
+	GetAccountKeys(
+		ctx context.Context,
+		addr flow.Address,
+		header *flow.Header,
+		snapshot snapshot.StorageSnapshot,
+	) (
+		[]flow.AccountPublicKey,
+		error,
+	)
 }
 
 type QueryConfig struct {
@@ -248,3 +268,50 @@ func (e *QueryExecutor) GetAccount(
 
 	return account, nil
 }
+
+func (e *QueryExecutor) GetAccountBalance(ctx context.Context, address flow.Address, blockHeader *flow.Header, snapshot snapshot.StorageSnapshot) (uint64, error) {
+
+	// TODO(ramtin): utilize ctx
+	blockCtx := fvm.NewContextFromParent(
+		e.vmCtx,
+		fvm.WithBlockHeader(blockHeader),
+		fvm.WithDerivedBlockData(
+			e.derivedChainData.NewDerivedBlockDataForScript(blockHeader.ID())))
+
+	accountBalance, err := fvm.GetAccountBalance(
+		blockCtx,
+		address,
+		snapshot)
+
+	if err != nil {
+		return 0, fmt.Errorf(
+			"failed to get account balance (%s) at block (%s): %w",
+			address.String(),
+			blockHeader.ID(),
+			err)
+	}
+
+	return accountBalance, nil
+}
+
+func (e *QueryExecutor) GetAccountKeys(ctx context.Context, address flow.Address, blockHeader *flow.Header, snapshot snapshot.StorageSnapshot) ([]flow.AccountPublicKey, error) {
+	// TODO(ramtin): utilize ctx
+	blockCtx := fvm.NewContextFromParent(
+		e.vmCtx,
+		fvm.WithBlockHeader(blockHeader),
+		fvm.WithDerivedBlockData(
+			e.derivedChainData.NewDerivedBlockDataForScript(blockHeader.ID())))
+
+	accountKeys, err := fvm.GetAccountKeys(blockCtx,
+		address,
+		snapshot)
+	if err != nil {
+		return nil, fmt.Errorf(
+			"failed to get account keys (%s) at block (%s): %w",
+			address.String(),
+			blockHeader.ID(),
+			err)
+	}
+
+	return accountKeys, nil
+}

fvm/environment/account_info.go

@@ -15,12 +15,13 @@ import (
 // AccountInfo exposes various account balance and storage statistics.
 type AccountInfo interface {
 	// Cadence's runtime APIs.
-	GetStorageUsed(runtimeaddress common.Address) (uint64, error)
+	GetStorageUsed(runtimeAddress common.Address) (uint64, error)
 	GetStorageCapacity(runtimeAddress common.Address) (uint64, error)
 	GetAccountBalance(runtimeAddress common.Address) (uint64, error)
 	GetAccountAvailableBalance(runtimeAddress common.Address) (uint64, error)
 
 	GetAccount(address flow.Address) (*flow.Account, error)
+	GetAccountKeys(address flow.Address) ([]flow.AccountPublicKey, error)
 }
 
 type ParseRestrictedAccountInfo struct {
@@ -103,6 +104,19 @@ func (info ParseRestrictedAccountInfo) GetAccount(
 		address)
 }
 
+func (info ParseRestrictedAccountInfo) GetAccountKeys(
+	address flow.Address,
+) (
+	[]flow.AccountPublicKey,
+	error,
+) {
+	return parseRestrict1Arg1Ret(
+		info.txnState,
+		trace.FVMEnvGetAccountKeys,
+		info.impl.GetAccountKeys,
+		address)
+}
+
 type accountInfo struct {
 	tracer tracing.TracerSpan
 	meter  Meter
@@ -152,7 +166,7 @@ func (info *accountInfo) GetStorageUsed(
 }
 
 // StorageMBUFixToBytesUInt converts the return type of storage capacity which
-// is a UFix64 with the unit of megabytes to UInt with the unit of bytes
+// is a UFix64 with the unit of megabytes to UInt with the unit of bytes.
 func StorageMBUFixToBytesUInt(result cadence.Value) uint64 {
 	// Divide the unsigned int by (1e8 (the scale of Fix64) / 1e6 (for mega))
 	// to get bytes (rounded down)
@@ -254,3 +268,20 @@ func (info *accountInfo) GetAccount(
 
 	return account, nil
 }
+
+func (info *accountInfo) GetAccountKeys(
+	address flow.Address,
+) (
+	[]flow.AccountPublicKey,
+	error,
+) {
+	defer info.tracer.StartChildSpan(trace.FVMEnvGetAccountKeys).End()
+
+	accountKeys, err := info.accounts.GetPublicKeys(address)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return accountKeys, nil
+}

fvm/environment/account_key_updater_test.go

@@ -165,6 +165,9 @@ func (f FakeAccounts) GetPublicKey(address flow.Address, keyIndex uint64) (flow.
 func (f FakeAccounts) SetPublicKey(_ flow.Address, _ uint64, _ flow.AccountPublicKey) ([]byte, error) {
 	return nil, nil
 }
+func (f FakeAccounts) GetPublicKeys(address flow.Address) ([]flow.AccountPublicKey, error) {
+	return make([]flow.AccountPublicKey, f.keyCount), nil
+}
 func (f FakeAccounts) GetContractNames(_ flow.Address) ([]string, error)      { return nil, nil }
 func (f FakeAccounts) GetContract(_ string, _ flow.Address) ([]byte, error)   { return nil, nil }
 func (f FakeAccounts) ContractExists(_ string, _ flow.Address) (bool, error)  { return false, nil }

fvm/environment/accounts.go

@@ -27,6 +27,7 @@ type Accounts interface {
 	AppendPublicKey(address flow.Address, key flow.AccountPublicKey) error
 	GetPublicKey(address flow.Address, keyIndex uint64) (flow.AccountPublicKey, error)
 	SetPublicKey(address flow.Address, keyIndex uint64, publicKey flow.AccountPublicKey) ([]byte, error)
+	GetPublicKeys(address flow.Address) ([]flow.AccountPublicKey, error)
 	GetContractNames(address flow.Address) ([]string, error)
 	GetContract(contractName string, address flow.Address) ([]byte, error)
 	ContractExists(contractName string, address flow.Address) (bool, error)
@@ -256,9 +257,7 @@ func (a *StatefulAccounts) GetPublicKeys(
 ) {
 	count, err := a.GetPublicKeyCount(address)
 	if err != nil {
-		return nil, fmt.Errorf(
-			"failed to get public key count of account: %w",
-			err)
+		return nil, err
 	}
 	publicKeys = make([]flow.AccountPublicKey, count)
 

fvm/environment/env.go

@@ -70,6 +70,7 @@ type Environment interface {
 
 	// AccountInfo
 	GetAccount(address flow.Address) (*flow.Account, error)
+	GetAccountKeys(address flow.Address) ([]flow.AccountPublicKey, error)
 
 	// RandomSourceHistory is the current block's derived random source.
 	// This source is only used by the core-contract that tracks the random source