From b4a451a0c59f86c31d02885e342533b44b1bf4d1 Mon Sep 17 00:00:00 2001 From: aman035 Date: Fri, 21 Aug 2026 22:02:43 +0530 Subject: [PATCH] fix(tss): select fund migration signers from the old key's shareholders --- .../tss/coordinator/coordinator.go | 91 +++++- .../tss/coordinator/coordinator_test.go | 11 + .../fund_migrate_participants_test.go | 276 ++++++++++++++++++ universalClient/tss/coordinator/utils.go | 19 +- .../tss/sessionmanager/sessionmanager_test.go | 4 + 5 files changed, 393 insertions(+), 8 deletions(-) create mode 100644 universalClient/tss/coordinator/fund_migrate_participants_test.go diff --git a/universalClient/tss/coordinator/coordinator.go b/universalClient/tss/coordinator/coordinator.go index 55627b11..84406236 100644 --- a/universalClient/tss/coordinator/coordinator.go +++ b/universalClient/tss/coordinator/coordinator.go @@ -33,6 +33,7 @@ import ( type PushCoreClient interface { GetLatestBlock(ctx context.Context) (uint64, error) GetCurrentKey(ctx context.Context) (*utsstypes.TssKey, error) + GetKeyByID(ctx context.Context, keyID string) (*utsstypes.TssKey, error) GetAllUniversalValidators(ctx context.Context) ([]*types.UniversalValidator, error) } @@ -472,11 +473,13 @@ func (c *Coordinator) processConfirmedEvents(ctx context.Context) error { // For SIGN/FUND_MIGRATE: pick a random threshold subset (>2/3 of eligible) rather than all eligible. // A threshold subset suffices for signing and is more resilient when some nodes are offline. // For all other protocols (keygen, keyrefresh, quorum_change), all eligible must participate. - var participants []*types.UniversalValidator - if event.Type == store.EventTypeSignOutbound || event.Type == store.EventTypeSignFundMigrate { - participants = getSignParticipants(allValidators) - } else { - participants = getEligibleForProtocol(event.Type, allValidators) + participants, err := c.selectParticipants(ctx, event, allValidators) + if err != nil { + c.logger.Error().Err(err). + Str("event_id", event.EventID). + Str("type", event.Type). + Msg("cannot select participants for event") + continue } if participants == nil { c.logger.Debug().Str("event_id", event.EventID).Str("type", event.Type).Msg("unknown protocol type") @@ -1156,3 +1159,81 @@ func (c *Coordinator) assignFundMigrateNonce(ctx context.Context, event store.Ev return builder.GetNextNonce(ctx, oldTSSAddr, true) } + +// selectParticipants picks who takes part in an event. +// +// For SIGN a random threshold subset (>2/3 of eligible) suffices and is more +// resilient when some nodes are offline. For all other protocols (keygen, +// keyrefresh, quorum change) every eligible validator must participate. +func (c *Coordinator) selectParticipants( + ctx context.Context, + event store.Event, + allValidators []*types.UniversalValidator, +) ([]*types.UniversalValidator, error) { + switch event.Type { + case store.EventTypeSignOutbound: + return getSignParticipants(allValidators), nil + case store.EventTypeSignFundMigrate: + // Signed with the old key's shares, so the signers must be drawn from + // the validators that hold them rather than from whoever is eligible + // now. A newcomer selected here has no such share and never ACKs, so + // the session stalls waiting for a party that cannot take part. + return c.fundMigrateParticipants(ctx, event, allValidators) + default: + return getEligibleForProtocol(event.Type, allValidators), nil + } +} + +// fundMigrateParticipants selects signers for a fund migration from the +// validators that hold the old key's shares. +// +// The signature is produced with the old keyshare, so eligibility is decided by +// the historical shareholder set recorded on chain, not by who is a validator +// today. The required count is the old key's threshold for the same reason: it +// is the quorum that key was created under. +// +// Fails rather than returning a short set. Too few surviving shareholders means +// no subset can sign, and returning one anyway would stall the session on an +// ACK that is never coming instead of reporting why. +func (c *Coordinator) fundMigrateParticipants( + ctx context.Context, + event store.Event, + allValidators []*types.UniversalValidator, +) ([]*types.UniversalValidator, error) { + var migrationData utsstypes.FundMigrationInitiatedEventData + if err := json.Unmarshal(event.EventData, &migrationData); err != nil { + return nil, fmt.Errorf("parse fund migration data: %w", err) + } + if migrationData.OldKeyID == "" { + return nil, fmt.Errorf("fund migration event carries no old key id") + } + + oldKey, err := c.pushCore.GetKeyByID(ctx, migrationData.OldKeyID) + if err != nil { + return nil, fmt.Errorf("fetch old key %s: %w", migrationData.OldKeyID, err) + } + if oldKey == nil || len(oldKey.Participants) == 0 { + return nil, fmt.Errorf("old key %s records no participants", migrationData.OldKeyID) + } + + shareholders := make(map[string]bool, len(oldKey.Participants)) + for _, p := range oldKey.Participants { + shareholders[p] = true + } + + var holders []*types.UniversalValidator + for _, v := range getSignEligible(allValidators) { + if v.IdentifyInfo != nil && shareholders[v.IdentifyInfo.CoreValidatorAddress] { + holders = append(holders, v) + } + } + + required := CalculateThreshold(len(oldKey.Participants)) + if len(holders) < required { + return nil, fmt.Errorf( + "key %s needs %d of its %d shareholders to sign, only %d are still eligible", + migrationData.OldKeyID, required, len(oldKey.Participants), len(holders)) + } + + return selectRandomSubset(holders, required), nil +} diff --git a/universalClient/tss/coordinator/coordinator_test.go b/universalClient/tss/coordinator/coordinator_test.go index 72daca8e..52753f3d 100644 --- a/universalClient/tss/coordinator/coordinator_test.go +++ b/universalClient/tss/coordinator/coordinator_test.go @@ -1123,6 +1123,10 @@ type stalenessMockPushCore struct { block uint64 validators []*types.UniversalValidator failGetAll bool + + // Old key history, consulted when selecting fund migration signers. + keysByID map[string]*utsstypes.TssKey + keyErr error } func (m *stalenessMockPushCore) GetLatestBlock(_ context.Context) (uint64, error) { @@ -1133,6 +1137,13 @@ func (m *stalenessMockPushCore) GetCurrentKey(_ context.Context) (*utsstypes.Tss return &utsstypes.TssKey{KeyId: "test-key"}, nil } +func (m *stalenessMockPushCore) GetKeyByID(_ context.Context, keyID string) (*utsstypes.TssKey, error) { + if m.keyErr != nil { + return nil, m.keyErr + } + return m.keysByID[keyID], nil +} + func (m *stalenessMockPushCore) GetAllUniversalValidators(_ context.Context) ([]*types.UniversalValidator, error) { if m.failGetAll { return nil, fmt.Errorf("simulated GetAllUniversalValidators RPC failure") diff --git a/universalClient/tss/coordinator/fund_migrate_participants_test.go b/universalClient/tss/coordinator/fund_migrate_participants_test.go new file mode 100644 index 00000000..c9ffebae --- /dev/null +++ b/universalClient/tss/coordinator/fund_migrate_participants_test.go @@ -0,0 +1,276 @@ +package coordinator + +import ( + "context" + "encoding/json" + "fmt" + "testing" + + "github.com/rs/zerolog" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/pushchain/push-chain-node/universalClient/store" + utsstypes "github.com/pushchain/push-chain-node/x/utss/types" + "github.com/pushchain/push-chain-node/x/uvalidator/types" +) + +func activeValidator(addr string) *types.UniversalValidator { + return &types.UniversalValidator{ + IdentifyInfo: &types.IdentityInfo{CoreValidatorAddress: addr}, + LifecycleInfo: &types.LifecycleInfo{CurrentStatus: types.UVStatus_UV_STATUS_ACTIVE}, + } +} + +func validatorWithStatus(addr string, status types.UVStatus) *types.UniversalValidator { + return &types.UniversalValidator{ + IdentifyInfo: &types.IdentityInfo{CoreValidatorAddress: addr}, + LifecycleInfo: &types.LifecycleInfo{CurrentStatus: status}, + } +} + +func validatorSet(addrs ...string) []*types.UniversalValidator { + set := make([]*types.UniversalValidator, 0, len(addrs)) + for _, a := range addrs { + set = append(set, activeValidator(a)) + } + return set +} + +func addressesOf(vs []*types.UniversalValidator) []string { + addrs := make([]string, 0, len(vs)) + for _, v := range vs { + addrs = append(addrs, v.IdentifyInfo.CoreValidatorAddress) + } + return addrs +} + +func fundMigrateEvent(t *testing.T, oldKeyID string) store.Event { + t.Helper() + data, err := json.Marshal(utsstypes.FundMigrationInitiatedEventData{OldKeyID: oldKeyID}) + require.NoError(t, err) + return store.Event{ + EventID: "fm-1", + Type: store.EventTypeSignFundMigrate, + EventData: data, + } +} + +func coordinatorWithKeys(keys map[string]*utsstypes.TssKey) *Coordinator { + return &Coordinator{ + pushCore: &stalenessMockPushCore{keysByID: keys}, + logger: zerolog.Nop(), + } +} + +// The finding's scenario: the old key has three shareholders, the validator set +// has since grown to ten. Selecting from the current set draws newcomers who +// hold no share of that key. +func TestFundMigrateParticipants_DrawsOnlyFromOldKeyShareholders(t *testing.T) { + keys := map[string]*utsstypes.TssKey{ + "old-key": {KeyId: "old-key", Participants: []string{"v1", "v2", "v3"}}, + } + c := coordinatorWithKeys(keys) + + all := validatorSet("v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10") + + // Selection is randomised, so repeat to catch a newcomer slipping in. + for i := 0; i < 200; i++ { + got, err := c.fundMigrateParticipants(context.Background(), fundMigrateEvent(t, "old-key"), all) + require.NoError(t, err) + + // Old key threshold is 3 of 3, not 7 of 10. + require.Len(t, got, 3) + assert.ElementsMatch(t, []string{"v1", "v2", "v3"}, addressesOf(got)) + } +} + +// A subset of shareholders large enough to sign, alongside a much larger +// current set. Every signer must still be a shareholder. +func TestFundMigrateParticipants_UsesOldKeyThreshold(t *testing.T) { + keys := map[string]*utsstypes.TssKey{ + "old-key": {KeyId: "old-key", Participants: []string{"v1", "v2", "v3", "v4", "v5", "v6"}}, + } + c := coordinatorWithKeys(keys) + + all := validatorSet("v1", "v2", "v3", "v4", "v5", "v6", "n1", "n2", "n3", "n4", "n5") + + shareholders := map[string]bool{"v1": true, "v2": true, "v3": true, "v4": true, "v5": true, "v6": true} + for i := 0; i < 200; i++ { + got, err := c.fundMigrateParticipants(context.Background(), fundMigrateEvent(t, "old-key"), all) + require.NoError(t, err) + + // CalculateThreshold(6) is 5, and it is the old key's size that decides. + require.Len(t, got, CalculateThreshold(6)) + for _, addr := range addressesOf(got) { + assert.True(t, shareholders[addr], "selected %s which holds no share of the old key", addr) + } + } +} + +// Fail closed rather than hand back a set that cannot reach the old key's +// threshold. A short set would stall the session on an ACK that never arrives. +func TestFundMigrateParticipants_FailsWhenTooFewShareholdersRemain(t *testing.T) { + keys := map[string]*utsstypes.TssKey{ + "old-key": {KeyId: "old-key", Participants: []string{"v1", "v2", "v3", "v4", "v5", "v6"}}, + } + c := coordinatorWithKeys(keys) + + // Only 4 of the 6 shareholders remain, one short of the threshold of 5, + // while the current set is comfortably large. + all := validatorSet("v1", "v2", "v3", "v4", "n1", "n2", "n3", "n4", "n5", "n6") + + got, err := c.fundMigrateParticipants(context.Background(), fundMigrateEvent(t, "old-key"), all) + require.Error(t, err) + assert.Nil(t, got) + assert.Contains(t, err.Error(), "only 4 are still eligible") +} + +// Pending leave keeps signing; anything else is not a usable signer even when +// it holds a share. +func TestFundMigrateParticipants_ExcludesIneligibleShareholders(t *testing.T) { + keys := map[string]*utsstypes.TssKey{ + "old-key": {KeyId: "old-key", Participants: []string{"v1", "v2", "v3"}}, + } + c := coordinatorWithKeys(keys) + + all := []*types.UniversalValidator{ + validatorWithStatus("v1", types.UVStatus_UV_STATUS_ACTIVE), + validatorWithStatus("v2", types.UVStatus_UV_STATUS_PENDING_LEAVE), + validatorWithStatus("v3", types.UVStatus_UV_STATUS_ACTIVE), + } + + got, err := c.fundMigrateParticipants(context.Background(), fundMigrateEvent(t, "old-key"), all) + require.NoError(t, err) + assert.ElementsMatch(t, []string{"v1", "v2", "v3"}, addressesOf(got)) + + // The same set with one shareholder no longer signing is one short. + all[1] = validatorWithStatus("v2", types.UVStatus_UV_STATUS_INACTIVE) + got, err = c.fundMigrateParticipants(context.Background(), fundMigrateEvent(t, "old-key"), all) + require.Error(t, err) + assert.Nil(t, got) +} + +func TestFundMigrateParticipants_RejectsUnusableEventData(t *testing.T) { + c := coordinatorWithKeys(map[string]*utsstypes.TssKey{ + "old-key": {KeyId: "old-key", Participants: []string{"v1", "v2", "v3"}}, + }) + all := validatorSet("v1", "v2", "v3") + + t.Run("malformed event data", func(t *testing.T) { + event := store.Event{EventID: "fm-1", Type: store.EventTypeSignFundMigrate, EventData: []byte("not json")} + _, err := c.fundMigrateParticipants(context.Background(), event, all) + require.Error(t, err) + assert.Contains(t, err.Error(), "parse fund migration data") + }) + + t.Run("no old key id", func(t *testing.T) { + _, err := c.fundMigrateParticipants(context.Background(), fundMigrateEvent(t, ""), all) + require.Error(t, err) + assert.Contains(t, err.Error(), "no old key id") + }) + + t.Run("unknown old key", func(t *testing.T) { + _, err := c.fundMigrateParticipants(context.Background(), fundMigrateEvent(t, "missing-key"), all) + require.Error(t, err) + assert.Contains(t, err.Error(), "records no participants") + }) + + t.Run("key with empty participants", func(t *testing.T) { + c := coordinatorWithKeys(map[string]*utsstypes.TssKey{"old-key": {KeyId: "old-key"}}) + _, err := c.fundMigrateParticipants(context.Background(), fundMigrateEvent(t, "old-key"), all) + require.Error(t, err) + assert.Contains(t, err.Error(), "records no participants") + }) + + t.Run("lookup failure", func(t *testing.T) { + c := &Coordinator{ + pushCore: &stalenessMockPushCore{keyErr: fmt.Errorf("rpc down")}, + logger: zerolog.Nop(), + } + _, err := c.fundMigrateParticipants(context.Background(), fundMigrateEvent(t, "old-key"), all) + require.Error(t, err) + assert.Contains(t, err.Error(), "fetch old key") + }) +} + +// A shareholder that has since dropped its identity record must not be counted +// towards the threshold, since it cannot be addressed as a party. +func TestFundMigrateParticipants_SkipsValidatorWithoutIdentity(t *testing.T) { + c := coordinatorWithKeys(map[string]*utsstypes.TssKey{ + "old-key": {KeyId: "old-key", Participants: []string{"v1", "v2", "v3"}}, + }) + + all := []*types.UniversalValidator{ + activeValidator("v1"), + {LifecycleInfo: &types.LifecycleInfo{CurrentStatus: types.UVStatus_UV_STATUS_ACTIVE}}, + activeValidator("v3"), + } + + _, err := c.fundMigrateParticipants(context.Background(), fundMigrateEvent(t, "old-key"), all) + require.Error(t, err) + assert.Contains(t, err.Error(), "only 2 are still eligible") +} + +// The routing itself: a fund migration must not be selected the way an +// outbound is, which is the defect this change fixes. +func TestSelectParticipants_RoutesFundMigrateToShareholders(t *testing.T) { + c := coordinatorWithKeys(map[string]*utsstypes.TssKey{ + "old-key": {KeyId: "old-key", Participants: []string{"v1", "v2", "v3"}}, + }) + + all := validatorSet("v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10") + + t.Run("fund migrate is confined to the old key", func(t *testing.T) { + for i := 0; i < 100; i++ { + got, err := c.selectParticipants(context.Background(), fundMigrateEvent(t, "old-key"), all) + require.NoError(t, err) + assert.ElementsMatch(t, []string{"v1", "v2", "v3"}, addressesOf(got)) + } + }) + + t.Run("outbound still uses the current set", func(t *testing.T) { + event := store.Event{EventID: "ob-1", Type: store.EventTypeSignOutbound} + got, err := c.selectParticipants(context.Background(), event, all) + require.NoError(t, err) + assert.Len(t, got, CalculateThreshold(len(all))) + }) + + t.Run("fund migrate reports rather than returning a short set", func(t *testing.T) { + _, err := c.selectParticipants(context.Background(), fundMigrateEvent(t, "gone"), all) + require.Error(t, err) + }) + + t.Run("other protocols take every eligible validator", func(t *testing.T) { + event := store.Event{EventID: "kg-1", Type: store.EventTypeKeygen} + got, err := c.selectParticipants(context.Background(), event, all) + require.NoError(t, err) + assert.Len(t, got, len(all)) + }) +} + +// selectRandomSubset is what keeps the count tied to the old key rather than to +// the surviving holders. +func TestSelectRandomSubset(t *testing.T) { + all := validatorSet("v1", "v2", "v3", "v4", "v5") + + assert.Nil(t, selectRandomSubset(nil, 3)) + assert.Nil(t, selectRandomSubset(all, 0)) + assert.Nil(t, selectRandomSubset(all, -1)) + assert.Len(t, selectRandomSubset(all, 5), 5) + assert.Len(t, selectRandomSubset(all, 9), 5) + + // Picks vary across calls and never repeat a validator within one pick. + seen := map[string]bool{} + for i := 0; i < 200; i++ { + got := selectRandomSubset(all, 3) + require.Len(t, got, 3) + unique := map[string]bool{} + for _, addr := range addressesOf(got) { + assert.False(t, unique[addr], "duplicate %s in one selection", addr) + unique[addr] = true + seen[addr] = true + } + } + assert.Len(t, seen, 5, "selection never reached some validators") +} diff --git a/universalClient/tss/coordinator/utils.go b/universalClient/tss/coordinator/utils.go index b64371ba..dfcfee77 100644 --- a/universalClient/tss/coordinator/utils.go +++ b/universalClient/tss/coordinator/utils.go @@ -80,13 +80,26 @@ func selectRandomThreshold(eligible []*types.UniversalValidator) []*types.Univer return eligible } - // Randomly select at least minRequired participants - // Shuffle and take first minRequired + return selectRandomSubset(eligible, minRequired) +} + +// selectRandomSubset returns a random n of eligible, or all of them when there +// are no more than n. Used where the required count is not derived from the +// input, such as fund migration, where the quorum belongs to the old key rather +// than to the set of validators still holding its shares. +func selectRandomSubset(eligible []*types.UniversalValidator, n int) []*types.UniversalValidator { + if len(eligible) == 0 || n <= 0 { + return nil + } + if len(eligible) <= n { + return eligible + } + shuffled := make([]*types.UniversalValidator, len(eligible)) copy(shuffled, eligible) rand.Shuffle(len(shuffled), func(i, j int) { shuffled[i], shuffled[j] = shuffled[j], shuffled[i] }) - return shuffled[:minRequired] + return shuffled[:n] } diff --git a/universalClient/tss/sessionmanager/sessionmanager_test.go b/universalClient/tss/sessionmanager/sessionmanager_test.go index 73f0dd2a..0cbd5bdd 100644 --- a/universalClient/tss/sessionmanager/sessionmanager_test.go +++ b/universalClient/tss/sessionmanager/sessionmanager_test.go @@ -63,6 +63,10 @@ func (m *mockPushCore) GetCurrentKey(_ context.Context) (*utsstypes.TssKey, erro return &utsstypes.TssKey{KeyId: "test-key"}, nil } +func (m *mockPushCore) GetKeyByID(_ context.Context, keyID string) (*utsstypes.TssKey, error) { + return &utsstypes.TssKey{KeyId: keyID}, nil +} + func (m *mockPushCore) GetAllUniversalValidators(_ context.Context) ([]*types.UniversalValidator, error) { return nil, nil }