Support POOL accounts with a new PublicDistribution intent

Author: jeffyanta

go.mod

@@ -6,7 +6,7 @@ require (
 	filippo.io/edwards25519 v1.1.0
 	github.com/aws/aws-sdk-go-v2 v0.17.0
 	github.com/bits-and-blooms/bloom/v3 v3.1.0
-	github.com/code-payments/code-protobuf-api v1.19.1-0.20250610140050-4cadbcc86f16
+	github.com/code-payments/code-protobuf-api v1.19.1-0.20250618155621-c66659ab4ff5
 	github.com/code-payments/code-vm-indexer v0.1.11-0.20241028132209-23031e814fba
 	github.com/emirpasic/gods v1.12.0
 	github.com/envoyproxy/protoc-gen-validate v1.2.1

go.sum

@@ -80,8 +80,8 @@ github.com/cncf/udpa/go v0.0.0-20200629203442-efcf912fb354/go.mod h1:WmhPx2Nbnht
 github.com/cncf/udpa/go v0.0.0-20201120205902-5459f2c99403/go.mod h1:WmhPx2Nbnhtbo57+VJT5O0JRkEi1Wbu0z5j0R8u5Hbk=
 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-protobuf-api v1.19.1-0.20250610140050-4cadbcc86f16 h1:drAMKRdbyObW8E4H6xc1pKIDxoFYgpaTdMlEnIKBIJ0=
-github.com/code-payments/code-protobuf-api v1.19.1-0.20250610140050-4cadbcc86f16/go.mod h1:ee6TzKbgMS42ZJgaFEMG3c4R3dGOiffHSu6MrY7WQvs=
+github.com/code-payments/code-protobuf-api v1.19.1-0.20250618155621-c66659ab4ff5 h1:fsZBRUCPGSZQngq/1rsJcEwhiYrZokr0wketRRTgGmI=
+github.com/code-payments/code-protobuf-api v1.19.1-0.20250618155621-c66659ab4ff5/go.mod h1:ee6TzKbgMS42ZJgaFEMG3c4R3dGOiffHSu6MrY7WQvs=
 github.com/code-payments/code-vm-indexer v0.1.11-0.20241028132209-23031e814fba h1:Bkp+gmeb6Y2PWXfkSCTMBGWkb2P1BujRDSjWeI+0j5I=
 github.com/code-payments/code-vm-indexer v0.1.11-0.20241028132209-23031e814fba/go.mod h1:jSiifpiBpyBQ8q2R0MGEbkSgWC6sbdRTyDBntmW+j1E=
 github.com/containerd/continuity v0.0.0-20190827140505-75bee3e2ccb6 h1:NmTXa/uVnDyp0TY5MKi197+3HWcnYWfnHGyaFthlnGw=

pkg/code/aml/guard.go

@@ -63,6 +63,9 @@ func (g *Guard) AllowMoneyMovement(ctx context.Context, intentRecord *intent.Rec
 	case intent.ReceivePaymentsPublicly:
 		// Public receives are always allowed
 		return true, nil
+	case intent.PublicDistribution:
+		// Public distributions are always allowed
+		return true, nil
 	default:
 		err := errors.New("intent record must be a send or receive payment")
 		tracer.OnError(err)

pkg/code/antispam/guard.go

@@ -3,6 +3,8 @@ package antispam
 import (
 	"context"
 
+	transactionpb "github.com/code-payments/code-protobuf-api/generated/go/transaction/v2"
+
 	"github.com/code-payments/code-server/pkg/code/common"
 	"github.com/code-payments/code-server/pkg/metrics"
 )
@@ -15,11 +17,11 @@ func NewGuard(integration Integration) *Guard {
 	return &Guard{integration: integration}
 }
 
-func (g *Guard) AllowOpenAccounts(ctx context.Context, owner *common.Account) (bool, error) {
+func (g *Guard) AllowOpenAccounts(ctx context.Context, owner *common.Account, accountSet transactionpb.OpenAccountsMetadata_AccountSet) (bool, error) {
 	tracer := metrics.TraceMethodCall(ctx, metricsStructName, "AllowOpenAccounts")
 	defer tracer.End()
 
-	allow, reason, err := g.integration.AllowOpenAccounts(ctx, owner)
+	allow, reason, err := g.integration.AllowOpenAccounts(ctx, owner, accountSet)
 	if err != nil {
 		return false, err
 	}
@@ -70,3 +72,17 @@ func (g *Guard) AllowReceivePayments(ctx context.Context, owner *common.Account,
 	}
 	return allow, nil
 }
+
+func (g *Guard) AllowDistribution(ctx context.Context, owner *common.Account, isPublic bool) (bool, error) {
+	tracer := metrics.TraceMethodCall(ctx, metricsStructName, "AllowDistribution")
+	defer tracer.End()
+
+	allow, reason, err := g.integration.AllowDistribution(ctx, owner, isPublic)
+	if err != nil {
+		return false, err
+	}
+	if !allow {
+		recordDenialEvent(ctx, actionDistribution, reason)
+	}
+	return allow, nil
+}

pkg/code/antispam/integration.go

@@ -3,19 +3,23 @@ package antispam
 import (
 	"context"
 
+	transactionpb "github.com/code-payments/code-protobuf-api/generated/go/transaction/v2"
+
 	"github.com/code-payments/code-server/pkg/code/common"
 )
 
 // Integration is an antispam guard integration that apps can implement to check
 // whether operations of interest are allowed to be performed.
 type Integration interface {
-	AllowOpenAccounts(ctx context.Context, owner *common.Account) (bool, string, error)
+	AllowOpenAccounts(ctx context.Context, owner *common.Account, accountSet transactionpb.OpenAccountsMetadata_AccountSet) (bool, string, error)
 
 	AllowWelcomeBonus(ctx context.Context, owner *common.Account) (bool, string, error)
 
 	AllowSendPayment(ctx context.Context, owner, destination *common.Account, isPublic bool) (bool, string, error)
 
 	AllowReceivePayments(ctx context.Context, owner *common.Account, isPublic bool) (bool, string, error)
+
+	AllowDistribution(ctx context.Context, owner *common.Account, isPublic bool) (bool, string, error)
 }
 
 type allowEverythingIntegration struct {
@@ -26,7 +30,7 @@ func NewAllowEverything() Integration {
 	return &allowEverythingIntegration{}
 }
 
-func (i *allowEverythingIntegration) AllowOpenAccounts(ctx context.Context, owner *common.Account) (bool, string, error) {
+func (i *allowEverythingIntegration) AllowOpenAccounts(ctx context.Context, owner *common.Account, accountSet transactionpb.OpenAccountsMetadata_AccountSet) (bool, string, error) {
 	return true, "", nil
 }
 
@@ -41,3 +45,7 @@ func (i *allowEverythingIntegration) AllowSendPayment(ctx context.Context, owner
 func (i *allowEverythingIntegration) AllowReceivePayments(ctx context.Context, owner *common.Account, isPublic bool) (bool, string, error) {
 	return true, "", nil
 }
+
+func (i *allowEverythingIntegration) AllowDistribution(ctx context.Context, owner *common.Account, isPublic bool) (bool, string, error) {
+	return true, "", nil
+}

pkg/code/antispam/metrics.go

@@ -14,6 +14,7 @@ const (
 	actionOpenAccounts    = "OpenAccounts"
 	actionSendPayment     = "SendPayment"
 	actionReceivePayments = "ReceivePayments"
+	actionDistribution    = "Distribution"
 
 	actionWelcomeBonus = "WelcomeBonus"
 )

pkg/code/async/account/gift_card.go

@@ -211,7 +211,7 @@ func InitiateProcessToAutoReturnGiftCard(ctx context.Context, data code_data.Pro
 			return err
 		}
 
-		return balanceLock.OnCommit(ctx, data)
+		return balanceLock.OnNewBalanceVersion(ctx, data)
 	})
 }
 

pkg/code/async/sequencer/intent_handler.go

@@ -125,6 +125,37 @@ func (h *ReceivePaymentsPubliclyIntentHandler) OnActionUpdated(ctx context.Conte
 	return nil
 }
 
+type PublicDistributionIntentHandler struct {
+	data code_data.Provider
+}
+
+func NewPublicDistributionIntentHandler(data code_data.Provider) IntentHandler {
+	return &PublicDistributionIntentHandler{
+		data: data,
+	}
+}
+
+func (h *PublicDistributionIntentHandler) OnActionUpdated(ctx context.Context, intentId string) error {
+	actionRecords, err := h.data.GetAllActionsByIntent(ctx, intentId)
+	if err != nil {
+		return err
+	}
+
+	for _, actionRecord := range actionRecords {
+		// Intent is failed if at least one transfer or withdraw action fails
+		if actionRecord.State == action.StateFailed {
+			return markIntentFailed(ctx, h.data, intentId)
+		}
+
+		if actionRecord.State != action.StateConfirmed {
+			return nil
+		}
+	}
+
+	// Intent is confirmed when all transfer and withdraw actions are confirmed
+	return markIntentConfirmed(ctx, h.data, intentId)
+}
+
 func validateIntentState(record *intent.Record, states ...intent.State) error {
 	for _, validState := range states {
 		if record.State == validState {
@@ -176,5 +207,6 @@ func getIntentHandlers(data code_data.Provider) map[intent.Type]IntentHandler {
 	handlersByType[intent.OpenAccounts] = NewOpenAccountsIntentHandler(data)
 	handlersByType[intent.SendPublicPayment] = NewSendPublicPaymentIntentHandler(data)
 	handlersByType[intent.ReceivePaymentsPublicly] = NewReceivePaymentsPubliclyIntentHandler(data)
+	handlersByType[intent.PublicDistribution] = NewPublicDistributionIntentHandler(data)
 	return handlersByType
 }

pkg/code/balance/lock.go

@@ -2,9 +2,11 @@ package balance
 
 import (
 	"context"
+	"errors"
 
 	"github.com/code-payments/code-server/pkg/code/common"
 	code_data "github.com/code-payments/code-server/pkg/code/data"
+	"github.com/code-payments/code-server/pkg/code/data/balance"
 )
 
 // OptimisticVersionLock is an optimistic version lock on an account's cached
@@ -28,7 +30,46 @@ func GetOptimisticVersionLock(ctx context.Context, data code_data.Provider, vaul
 	}, nil
 }
 
-// OnCommit is called in the DB transaction updating the account's cached balance
-func (l *OptimisticVersionLock) OnCommit(ctx context.Context, data code_data.Provider) error {
+// OnNewBalanceVersion is called in the DB transaction updating the account's
+// cached balance
+func (l *OptimisticVersionLock) OnNewBalanceVersion(ctx context.Context, data code_data.Provider) error {
 	return data.AdvanceCachedBalanceVersion(ctx, l.vault.PublicKey().ToBase58(), l.currentVersion)
 }
+
+// RequireSameBalanceVerion is called in the DB transaction requireing the
+// account's cached balance not be changed
+func (l *OptimisticVersionLock) RequireSameBalanceVerion(ctx context.Context, data code_data.Provider) error {
+	latestVersion, err := data.GetCachedBalanceVersion(ctx, l.vault.PublicKey().ToBase58())
+	if err != nil {
+		return err
+	}
+	if latestVersion < l.currentVersion {
+		return errors.New("unexpected balance version detected")
+	}
+	if l.currentVersion != latestVersion {
+		return balance.ErrStaleCachedBalanceVersion
+	}
+	return nil
+}
+
+// OpenCloseStatusLock is a lock on an account's open/close status
+type OpenCloseStatusLock struct {
+	vault *common.Account
+}
+
+func NewOpenCloseStatusLock(vault *common.Account) *OpenCloseStatusLock {
+	return &OpenCloseStatusLock{
+		vault: vault,
+	}
+}
+
+// OnPaymentToAccount is called in the DB transaction making a payment to the
+// account that may be closed
+func (l *OpenCloseStatusLock) OnPaymentToAccount(ctx context.Context, data code_data.Provider) error {
+	return data.CheckNotClosedForBalanceUpdate(ctx, l.vault.PublicKey().ToBase58())
+}
+
+// OnClose is called in the DB transaction closing the account
+func (l *OpenCloseStatusLock) OnClose(ctx context.Context, data code_data.Provider) error {
+	return data.MarkAsClosedForBalanceUpdate(ctx, l.vault.PublicKey().ToBase58())
+}

pkg/code/common/owner_test.go

@@ -136,6 +136,8 @@ func TestGetLatestTokenAccountRecordsForOwner(t *testing.T) {
 	authority3 := newRandomTestAccount(t)
 	authority4 := newRandomTestAccount(t)
 	authority5 := newRandomTestAccount(t)
+	authority6 := newRandomTestAccount(t)
+	authority7 := newRandomTestAccount(t)
 
 	for _, authorityAndType := range []struct {
 		account     *Account
@@ -195,9 +197,30 @@ func TestGetLatestTokenAccountRecordsForOwner(t *testing.T) {
 	}
 	require.NoError(t, data.CreateAccountInfo(ctx, swapAccountInfoRecord))
 
+	for i, authority := range []*Account{
+		authority6,
+		authority7,
+	} {
+		timelockAccounts, err := authority.GetTimelockAccounts(vmAccount, coreMintAccount)
+		require.NoError(t, err)
+
+		timelockRecord := timelockAccounts.ToDBRecord()
+		require.NoError(t, data.SaveTimelock(ctx, timelockRecord))
+
+		accountInfoRecord := &account.Record{
+			OwnerAccount:     owner.PublicKey().ToBase58(),
+			AuthorityAccount: timelockRecord.VaultOwner,
+			TokenAccount:     timelockRecord.VaultAddress,
+			MintAccount:      coreMintAccount.PublicKey().ToBase58(),
+			AccountType:      commonpb.AccountType_POOL,
+			Index:            uint64(i),
+		}
+		require.NoError(t, data.CreateAccountInfo(ctx, accountInfoRecord))
+	}
+
 	actual, err = GetLatestTokenAccountRecordsForOwner(ctx, data, owner)
 	require.NoError(t, err)
-	require.Len(t, actual, 4)
+	require.Len(t, actual, 5)
 
 	records, ok := actual[commonpb.AccountType_BUCKET_1_KIN]
 	require.True(t, ok)
@@ -238,4 +261,20 @@ func TestGetLatestTokenAccountRecordsForOwner(t *testing.T) {
 	assert.Equal(t, records[0].General.AuthorityAccount, authority5.PublicKey().ToBase58())
 	assert.Equal(t, records[0].General.TokenAccount, swapAta.PublicKey().ToBase58())
 	assert.Equal(t, records[0].General.AccountType, commonpb.AccountType_SWAP)
+
+	records, ok = actual[commonpb.AccountType_POOL]
+	require.True(t, ok)
+	require.Len(t, records, 2)
+
+	assert.Equal(t, records[0].General.AuthorityAccount, authority6.PublicKey().ToBase58())
+	assert.Equal(t, records[0].General.AccountType, commonpb.AccountType_POOL)
+	assert.Equal(t, records[0].Timelock.VaultOwner, authority6.PublicKey().ToBase58())
+	assert.Equal(t, records[0].General.TokenAccount, records[0].Timelock.VaultAddress)
+	assert.EqualValues(t, records[0].General.Index, 0)
+
+	assert.Equal(t, records[1].General.AuthorityAccount, authority7.PublicKey().ToBase58())
+	assert.Equal(t, records[1].General.AccountType, commonpb.AccountType_POOL)
+	assert.Equal(t, records[1].Timelock.VaultOwner, authority7.PublicKey().ToBase58())
+	assert.Equal(t, records[1].General.TokenAccount, records[1].Timelock.VaultAddress)
+	assert.EqualValues(t, records[1].General.Index, 1)
 }

pkg/code/data/account/acccount_info.go

@@ -24,6 +24,7 @@ var AllAccountTypes = []commonpb.AccountType{
 	commonpb.AccountType_REMOTE_SEND_GIFT_CARD,
 	commonpb.AccountType_RELATIONSHIP,
 	commonpb.AccountType_SWAP,
+	commonpb.AccountType_POOL,
 }
 
 type Record struct {
@@ -180,6 +181,10 @@ func (r *Record) Validate() error {
 		if r.OwnerAccount == r.AuthorityAccount {
 			return errors.New("owner cannot be authority for swap account")
 		}
+	case commonpb.AccountType_POOL:
+		if r.OwnerAccount == r.AuthorityAccount {
+			return errors.New("owner cannot be authority pool account")
+		}
 	default:
 		return errors.Errorf("unhandled account type: %s", r.AccountType.String())
 	}

pkg/code/data/account/memory/store.go

@@ -275,6 +275,10 @@ func (s *store) GetLatestByOwnerAddress(_ context.Context, address string) (map[
 
 	items := s.findByOwnerAddress(address)
 	for _, accountType := range account.AllAccountTypes {
+		if accountType == commonpb.AccountType_POOL {
+			continue
+		}
+
 		items := s.filterByType(items, accountType)
 		if len(items) == 0 {
 			continue
@@ -292,6 +296,11 @@ func (s *store) GetLatestByOwnerAddress(_ context.Context, address string) (map[
 		}
 	}
 
+	items = s.filterByType(items, commonpb.AccountType_POOL)
+	if len(items) > 0 {
+		res[commonpb.AccountType_POOL] = cloneRecords(items)
+	}
+
 	if len(res) == 0 {
 		return nil, account.ErrAccountInfoNotFound
 	}

pkg/code/data/account/postgres/model.go

@@ -210,27 +210,47 @@ func dbGetByAuthorityAddress(ctx context.Context, db *sqlx.DB, address string) (
 }
 
 func dbGetLatestByOwnerAddress(ctx context.Context, db *sqlx.DB, address string) ([]*model, error) {
-	var res []*model
+	var res1 []*model
 
-	query := `SELECT DISTINCT ON (account_type, relationship_to) id, owner_account, authority_account, token_account, mint_account, account_type, index, relationship_to, requires_deposit_sync, deposits_last_synced_at, requires_auto_return_check, requires_swap_retry, last_swap_retry_at, created_at FROM ` + tableName + `
-		WHERE owner_account = $1
+	query1 := `SELECT DISTINCT ON (account_type, relationship_to) id, owner_account, authority_account, token_account, mint_account, account_type, index, relationship_to, requires_deposit_sync, deposits_last_synced_at, requires_auto_return_check, requires_swap_retry, last_swap_retry_at, created_at FROM ` + tableName + `
+		WHERE owner_account = $1 AND account_type NOT IN ($2)
 		ORDER BY account_type, relationship_to, index DESC
 	`
 
 	err := db.SelectContext(
 		ctx,
-		&res,
-		query,
+		&res1,
+		query1,
 		address,
+		commonpb.AccountType_POOL,
 	)
-	if err != nil {
-		return nil, pgutil.CheckNoRows(err, account.ErrAccountInfoNotFound)
+	if err != nil && !pgutil.IsNoRows(err) {
+		return nil, err
+	}
+
+	var res2 []*model
+
+	query2 := `SELECT id, owner_account, authority_account, token_account, mint_account, account_type, index, relationship_to, requires_deposit_sync, deposits_last_synced_at, requires_auto_return_check, requires_swap_retry, last_swap_retry_at, created_at FROM ` + tableName + `
+		WHERE owner_account = $1 AND account_type IN ($2)
+		ORDER BY index ASC
+	`
+	err = db.SelectContext(
+		ctx,
+		&res2,
+		query2,
+		address,
+		commonpb.AccountType_POOL,
+	)
+	if err != nil && !pgutil.IsNoRows(err) {
+		return nil, err
 	}
 
+	var res []*model
+	res = append(res, res1...)
+	res = append(res, res2...)
 	if len(res) == 0 {
 		return nil, account.ErrAccountInfoNotFound
 	}
-
 	return res, nil
 }