From 40829e12bbb42bb95bef92a68c645b899d0714d5 Mon Sep 17 00:00:00 2001 From: jeffyanta Date: Wed, 26 Aug 2026 09:31:40 -0400 Subject: [PATCH] Implement initial Balance RPC service --- go.mod | 2 +- go.sum | 4 +- ocp/rpc/balance/server.go | 144 ++++++++++++++++++++++ ocp/rpc/balance/server_test.go | 217 +++++++++++++++++++++++++++++++++ 4 files changed, 364 insertions(+), 3 deletions(-) create mode 100644 ocp/rpc/balance/server.go create mode 100644 ocp/rpc/balance/server_test.go diff --git a/go.mod b/go.mod index 8efce77..0f39fb4 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 7a1cb1d..a7fe8e0 100644 --- a/go.sum +++ b/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= diff --git a/ocp/rpc/balance/server.go b/ocp/rpc/balance/server.go new file mode 100644 index 0000000..d83413c --- /dev/null +++ b/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(¤cycreator.EstimateSellArgs{ + CurrentSupplyInQuarks: reserveState.SupplyFromBonding, + SellAmountInQuarks: quarks, + ValueMintDecimals: uint8(common.CoreMintDecimals), + SellFeeBps: 0, + }) + coreMintValue += sellValue + } + + return coreMintValue, nil +} diff --git a/ocp/rpc/balance/server_test.go b/ocp/rpc/balance/server_test.go new file mode 100644 index 0000000..27d90eb --- /dev/null +++ b/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(¤cycreator.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)) +}