Implement initial Balance RPC service

go.mod

@@ -11,7 +11,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/dynamodb v1.59.0
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.102.2
 	github.com/code-payments/code-vm-indexer v1.2.0
-	github.com/code-payments/ocp-protobuf-api v1.14.1-0.20260818131607-042819236352
+	github.com/code-payments/ocp-protobuf-api v1.16.0
 	github.com/emirpasic/gods v1.12.0
 	github.com/envoyproxy/protoc-gen-validate v1.3.3
 	github.com/golang/protobuf v1.5.4

go.sum

@@ -80,8 +80,8 @@ github.com/cockroachdb/apd v1.1.0 h1:3LFP3629v+1aKXU5Q37mxmRxX/pIu1nijXydLShEq5I
 github.com/cockroachdb/apd v1.1.0/go.mod h1:8Sl8LxpKi29FqWXR16WEFZRNSz3SoPzUzeMeY4+DwBQ=
 github.com/code-payments/code-vm-indexer v1.2.0 h1:rSHpBMiT9BKgmKcXg/VIoi/h0t7jNxGx07Qz59m+6Q0=
 github.com/code-payments/code-vm-indexer v1.2.0/go.mod h1:vn91YN2qNqb+gGJeZe2+l+TNxVmEEiRHXXnIn2Y40h8=
-github.com/code-payments/ocp-protobuf-api v1.14.1-0.20260818131607-042819236352 h1:RsDDYM1VXBNcEYSZ+GvnMXCua9pMckYdIwdJnmiX+FI=
-github.com/code-payments/ocp-protobuf-api v1.14.1-0.20260818131607-042819236352/go.mod h1:tw6BooY5a8l6CtSZnKOruyKII0W04n89pcM4BizrgG8=
+github.com/code-payments/ocp-protobuf-api v1.16.0 h1:j8Z5TRsNd8b02mCjJSq8+e1CoRh0JCnRVDldTlb9+Ds=
+github.com/code-payments/ocp-protobuf-api v1.16.0/go.mod h1:tw6BooY5a8l6CtSZnKOruyKII0W04n89pcM4BizrgG8=
 github.com/containerd/continuity v0.0.0-20190827140505-75bee3e2ccb6 h1:NmTXa/uVnDyp0TY5MKi197+3HWcnYWfnHGyaFthlnGw=
 github.com/containerd/continuity v0.0.0-20190827140505-75bee3e2ccb6/go.mod h1:GL3xCUCBDV3CZiTSEKksMWbLE66hEyuu9qyDOOqM47Y=
 github.com/coreos/bbolt v1.3.2/go.mod h1:iRUV2dpdMOn7Bo10OQBFzIJO9kkE559Wcmn+qkEiiKk=

ocp/rpc/balance/server.go

@@ -0,0 +1,144 @@
+package balance
+
+import (
+	"context"
+
+	"go.uber.org/zap"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	balancepb "github.com/code-payments/ocp-protobuf-api/generated/go/balance/v1"
+
+	"github.com/code-payments/ocp-server/grpc/client"
+	"github.com/code-payments/ocp-server/ocp/balance"
+	"github.com/code-payments/ocp-server/ocp/common"
+	currency_util "github.com/code-payments/ocp-server/ocp/currency"
+	ocp_data "github.com/code-payments/ocp-server/ocp/data"
+	"github.com/code-payments/ocp-server/ocp/rpc"
+	"github.com/code-payments/ocp-server/solana/currencycreator"
+)
+
+type server struct {
+	log              *zap.Logger
+	data             ocp_data.Provider
+	mintDataProvider *currency_util.MintDataProvider
+
+	balancepb.UnimplementedBalanceServer
+}
+
+func NewBalanceServer(log *zap.Logger, data ocp_data.Provider, mintDataProvider *currency_util.MintDataProvider) balancepb.BalanceServer {
+	return &server{
+		log:              log,
+		data:             data,
+		mintDataProvider: mintDataProvider,
+	}
+}
+
+func (s *server) GetBalance(ctx context.Context, req *balancepb.GetBalanceRequest) (*balancepb.GetBalanceResponse, error) {
+	log := s.log.With(zap.String("method", "GetBalance"))
+	log = client.InjectLoggingMetadata(ctx, log, rpc.UserAgentName)
+
+	owner, err := common.NewAccountFromProto(req.Owner)
+	if err != nil {
+		log.With(zap.Error(err)).Warn("invalid owner account")
+		return nil, status.Error(codes.Internal, "")
+	}
+	log = log.With(zap.String("owner_account", owner.PublicKey().ToBase58()))
+
+	ownerMetadata, err := common.GetOwnerMetadata(ctx, s.data, owner)
+	if err == common.ErrOwnerNotFound {
+		return &balancepb.GetBalanceResponse{
+			Result: balancepb.GetBalanceResponse_NOT_FOUND,
+		}, nil
+	} else if err != nil {
+		log.With(zap.Error(err)).Warn("failure getting owner metadata")
+		return nil, status.Error(codes.Internal, "")
+	}
+
+	if ownerMetadata.Type != common.OwnerTypeUser12Words {
+		return &balancepb.GetBalanceResponse{
+			Result: balancepb.GetBalanceResponse_NOT_FOUND,
+		}, nil
+	}
+
+	coreMintValue, err := s.calculateCoreMintValue(ctx, owner)
+	if err != nil {
+		log.With(zap.Error(err)).Warn("failure calculating core mint value")
+		return nil, status.Error(codes.Internal, "")
+	}
+
+	return &balancepb.GetBalanceResponse{
+		Result:        balancepb.GetBalanceResponse_OK,
+		CoreMintValue: coreMintValue,
+	}, nil
+}
+
+func (s *server) calculateCoreMintValue(ctx context.Context, owner *common.Account) (uint64, error) {
+	recordsByMintAndType, err := common.GetLatestCodeTimelockAccountRecordsForOwner(ctx, s.data, owner)
+	if err != nil {
+		return 0, err
+	}
+
+	// Accounts that have left the L2 system don't have a cached balance that can
+	// be trusted, so they're excluded from the calculation.
+	mintByTokenAccount := make(map[string]string)
+	var managedByCodeRecords []*common.AccountRecords
+	for mint, recordsByType := range recordsByMintAndType {
+		for _, recordsList := range recordsByType {
+			for _, records := range recordsList {
+				if !records.IsManagedByCode(ctx) {
+					continue
+				}
+
+				mintByTokenAccount[records.General.TokenAccount] = mint
+				managedByCodeRecords = append(managedByCodeRecords, records)
+			}
+		}
+	}
+
+	if len(managedByCodeRecords) == 0 {
+		return 0, nil
+	}
+
+	balanceByTokenAccount, err := balance.BatchCalculateFromCacheWithAccountRecords(ctx, s.data, managedByCodeRecords...)
+	if err != nil {
+		return 0, err
+	}
+
+	quarksByMint := make(map[string]uint64)
+	for tokenAccount, quarks := range balanceByTokenAccount {
+		quarksByMint[mintByTokenAccount[tokenAccount]] += quarks
+	}
+
+	var coreMintValue uint64
+	for mint, quarks := range quarksByMint {
+		if quarks == 0 {
+			continue
+		}
+
+		if mint == common.CoreMintAccount.PublicKey().ToBase58() {
+			coreMintValue += quarks
+			continue
+		}
+
+		mintAccount, err := common.NewAccountFromPublicKeyString(mint)
+		if err != nil {
+			return 0, err
+		}
+
+		reserveState, err := s.mintDataProvider.GetLiveReserveState(ctx, mintAccount)
+		if err != nil {
+			return 0, err
+		}
+
+		sellValue, _ := currencycreator.EstimateSell(&currencycreator.EstimateSellArgs{
+			CurrentSupplyInQuarks: reserveState.SupplyFromBonding,
+			SellAmountInQuarks:    quarks,
+			ValueMintDecimals:     uint8(common.CoreMintDecimals),
+			SellFeeBps:            0,
+		})
+		coreMintValue += sellValue
+	}
+
+	return coreMintValue, nil
+}

ocp/rpc/balance/server_test.go

@@ -0,0 +1,217 @@
+package balance
+
+import (
+	"context"
+	"fmt"
+	"math/rand"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap/zaptest"
+	"google.golang.org/grpc"
+
+	balancepb "github.com/code-payments/ocp-protobuf-api/generated/go/balance/v1"
+	commonpb "github.com/code-payments/ocp-protobuf-api/generated/go/common/v1"
+
+	"github.com/code-payments/ocp-server/ocp/common"
+	currency_util "github.com/code-payments/ocp-server/ocp/currency"
+	ocp_data "github.com/code-payments/ocp-server/ocp/data"
+	"github.com/code-payments/ocp-server/ocp/data/account"
+	exchange_memory "github.com/code-payments/ocp-server/ocp/data/currency/exchange/memory"
+	"github.com/code-payments/ocp-server/ocp/data/currency/holder"
+	holder_memory "github.com/code-payments/ocp-server/ocp/data/currency/holder/memory"
+	"github.com/code-payments/ocp-server/ocp/data/currency/reserve"
+	reserve_memory "github.com/code-payments/ocp-server/ocp/data/currency/reserve/memory"
+	"github.com/code-payments/ocp-server/ocp/data/deposit"
+	"github.com/code-payments/ocp-server/ocp/data/transaction"
+	"github.com/code-payments/ocp-server/solana/currencycreator"
+	timelock_token_v1 "github.com/code-payments/ocp-server/solana/timelock/v1"
+	"github.com/code-payments/ocp-server/testutil"
+)
+
+type testEnv struct {
+	ctx          context.Context
+	client       balancepb.BalanceClient
+	data         ocp_data.Provider
+	reserveStore reserve.Store
+	holderStore  holder.Store
+}
+
+func setup(t *testing.T) (env testEnv, cleanup func()) {
+	log := zaptest.NewLogger(t)
+
+	conn, serv, err := testutil.NewServer(log)
+	require.NoError(t, err)
+
+	env.ctx = context.Background()
+	env.client = balancepb.NewBalanceClient(conn)
+	env.data = ocp_data.NewTestDataProvider()
+	env.reserveStore = reserve_memory.New()
+	env.holderStore = holder_memory.New()
+	testutil.SetupRandomSubsidizer(t, env.data)
+
+	exchangeRateStore := exchange_memory.New()
+	mintDataProvider := currency_util.NewMintDataProvider(log, env.data, exchangeRateStore, env.reserveStore, env.holderStore, 0, time.Second, time.Second)
+	s := NewBalanceServer(log, env.data, mintDataProvider)
+
+	serv.RegisterService(func(server *grpc.Server) {
+		balancepb.RegisterBalanceServer(server, s)
+	})
+
+	require.NoError(t, mintDataProvider.Start(env.ctx))
+
+	serverCleanup, err := serv.Serve()
+	require.NoError(t, err)
+
+	cleanup = func() {
+		mintDataProvider.Stop()
+		serverCleanup()
+	}
+	return env, cleanup
+}
+
+func TestGetBalance_HappyPath(t *testing.T) {
+	env, cleanup := setup(t)
+	defer cleanup()
+
+	coreVmConfig := testutil.NewRandomVmConfig(t, true)
+	launchpadMint := testutil.SetupLaunchpadCurrency(t, env.data, env.reserveStore, env.holderStore)
+	launchpadVmConfig, err := common.GetVmConfigForMint(env.ctx, env.data, launchpadMint)
+	require.NoError(t, err)
+
+	ownerAccount := testutil.NewRandomAccount(t)
+
+	req := &balancepb.GetBalanceRequest{
+		Owner: ownerAccount.ToProto(),
+	}
+
+	resp, err := env.client.GetBalance(env.ctx, req)
+	require.NoError(t, err)
+	assert.Equal(t, balancepb.GetBalanceResponse_NOT_FOUND, resp.Result)
+	assert.EqualValues(t, 0, resp.CoreMintValue)
+
+	primaryCoreMintAccountRecords := setupAccountRecords(t, env, ownerAccount, ownerAccount, coreVmConfig, 0, commonpb.AccountType_PRIMARY)
+	primaryLaunchpadMintAccountRecords := setupAccountRecords(t, env, ownerAccount, ownerAccount, launchpadVmConfig, 0, commonpb.AccountType_PRIMARY)
+
+	resp, err = env.client.GetBalance(env.ctx, req)
+	require.NoError(t, err)
+	assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result)
+	assert.EqualValues(t, 0, resp.CoreMintValue)
+
+	setupCachedBalance(t, env, primaryCoreMintAccountRecords, common.ToCoreMintQuarks(42))
+	setupCachedBalance(t, env, primaryLaunchpadMintAccountRecords, currencycreator.ToQuarks(100))
+
+	// The launchpad currency's value is what the entire position would currently
+	// sell for on the bonding curve.
+	expectedLaunchpadMintValue, _ := currencycreator.EstimateSell(&currencycreator.EstimateSellArgs{
+		CurrentSupplyInQuarks: currencycreator.ToQuarks(1_000),
+		SellAmountInQuarks:    currencycreator.ToQuarks(100),
+		ValueMintDecimals:     uint8(common.CoreMintDecimals),
+		SellFeeBps:            0,
+	})
+	require.NotZero(t, expectedLaunchpadMintValue)
+
+	resp, err = env.client.GetBalance(env.ctx, req)
+	require.NoError(t, err)
+	assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result)
+	assert.EqualValues(t, common.ToCoreMintQuarks(42)+expectedLaunchpadMintValue, resp.CoreMintValue)
+}
+
+func TestGetBalance_UnmanagedAccountsExcluded(t *testing.T) {
+	env, cleanup := setup(t)
+	defer cleanup()
+
+	coreVmConfig := testutil.NewRandomVmConfig(t, true)
+
+	ownerAccount := testutil.NewRandomAccount(t)
+
+	primaryCoreMintAccountRecords := setupAccountRecords(t, env, ownerAccount, ownerAccount, coreVmConfig, 0, commonpb.AccountType_PRIMARY)
+
+	setupCachedBalance(t, env, primaryCoreMintAccountRecords, common.ToCoreMintQuarks(42))
+
+	// The pool account has left the L2 system, so there isn't a cached balance that
+	// can be trusted for it.
+	primaryCoreMintAccountRecords.Timelock.VaultState = timelock_token_v1.StateUnlocked
+	primaryCoreMintAccountRecords.Timelock.Block += 1
+	require.NoError(t, env.data.SaveTimelock(env.ctx, primaryCoreMintAccountRecords.Timelock))
+
+	resp, err := env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{
+		Owner: ownerAccount.ToProto(),
+	})
+	require.NoError(t, err)
+	assert.Equal(t, balancepb.GetBalanceResponse_OK, resp.Result)
+	assert.EqualValues(t, common.ToCoreMintQuarks(0), resp.CoreMintValue)
+}
+
+func TestGetBalance_GiftCardOwnerAccount(t *testing.T) {
+	env, cleanup := setup(t)
+	defer cleanup()
+
+	coreVmConfig := testutil.NewRandomVmConfig(t, true)
+
+	giftCardOwnerAccount := testutil.NewRandomAccount(t)
+
+	giftCardAccountRecords := setupAccountRecords(t, env, giftCardOwnerAccount, giftCardOwnerAccount, coreVmConfig, 0, commonpb.AccountType_REMOTE_SEND_GIFT_CARD)
+	setupCachedBalance(t, env, giftCardAccountRecords, common.ToCoreMintQuarks(42))
+
+	resp, err := env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{
+		Owner: giftCardOwnerAccount.ToProto(),
+	})
+	require.NoError(t, err)
+	assert.Equal(t, balancepb.GetBalanceResponse_NOT_FOUND, resp.Result)
+	assert.EqualValues(t, 0, resp.CoreMintValue)
+}
+
+func TestGetBalance_UnknownOwnerAccount(t *testing.T) {
+	env, cleanup := setup(t)
+	defer cleanup()
+
+	resp, err := env.client.GetBalance(env.ctx, &balancepb.GetBalanceRequest{
+		Owner: testutil.NewRandomAccount(t).ToProto(),
+	})
+	require.NoError(t, err)
+	assert.Equal(t, balancepb.GetBalanceResponse_NOT_FOUND, resp.Result)
+	assert.EqualValues(t, 0, resp.CoreMintValue)
+}
+
+func setupAccountRecords(t *testing.T, env testEnv, ownerAccount, authorityAccount *common.Account, vmConfig *common.VmConfig, index uint64, accountType commonpb.AccountType) *common.AccountRecords {
+	timelockAccounts, err := authorityAccount.GetTimelockAccounts(vmConfig)
+	require.NoError(t, err)
+
+	timelockRecord := timelockAccounts.ToDBRecord()
+	timelockRecord.VaultState = timelock_token_v1.StateLocked
+	timelockRecord.Block += 1
+
+	accountInfoRecord := &account.Record{
+		OwnerAccount:     ownerAccount.PublicKey().ToBase58(),
+		AuthorityAccount: authorityAccount.PublicKey().ToBase58(),
+		TokenAccount:     timelockAccounts.Vault.PublicKey().ToBase58(),
+		MintAccount:      vmConfig.Mint.PublicKey().ToBase58(),
+
+		AccountType: accountType,
+
+		Index: index,
+	}
+
+	require.NoError(t, env.data.CreateAccountInfo(env.ctx, accountInfoRecord))
+	require.NoError(t, env.data.SaveTimelock(env.ctx, timelockRecord))
+
+	return &common.AccountRecords{
+		General:  accountInfoRecord,
+		Timelock: timelockRecord,
+	}
+}
+
+func setupCachedBalance(t *testing.T, env testEnv, accountRecords *common.AccountRecords, balance uint64) {
+	depositRecord := &deposit.Record{
+		Signature:   fmt.Sprintf("txn%d", rand.Uint64()),
+		Destination: accountRecords.General.TokenAccount,
+		Amount:      balance,
+
+		ConfirmationState: transaction.ConfirmationFinalized,
+		Slot:              12345,
+	}
+	require.NoError(t, env.data.SaveExternalDeposit(env.ctx, depositRecord))
+}