diff --git a/internal/db/delegation.go b/internal/db/delegation.go index cfa3597..40769c6 100644 --- a/internal/db/delegation.go +++ b/internal/db/delegation.go @@ -36,13 +36,22 @@ func (db *Database) SaveNewBTCDelegation( } func (db *Database) UpdateBTCDelegationState( - ctx context.Context, stakingTxHash string, newState types.DelegationState, + ctx context.Context, + stakingTxHash string, + newState types.DelegationState, + subState *types.DelegationSubState, ) error { filter := bson.M{"_id": stakingTxHash} + updateFields := bson.M{ + "state": newState.String(), + } + + if subState != nil { + updateFields["sub_state"] = subState.String() + } + update := bson.M{ - "$set": bson.M{ - "state": newState.String(), - }, + "$set": updateFields, } res := db.client.Database(db.dbName). diff --git a/internal/db/interface.go b/internal/db/interface.go index 70fd648..50888fc 100644 --- a/internal/db/interface.go +++ b/internal/db/interface.go @@ -98,7 +98,10 @@ type DbInterface interface { * @return An error if the operation failed */ UpdateBTCDelegationState( - ctx context.Context, stakingTxHash string, newState types.DelegationState, + ctx context.Context, + stakingTxHash string, + newState types.DelegationState, + subState *types.DelegationSubState, ) error /** * SaveBTCDelegationUnbondingCovenantSignature saves a BTC delegation @@ -160,7 +163,10 @@ type DbInterface interface { * @return An error if the operation failed */ SaveNewTimeLockExpire( - ctx context.Context, stakingTxHashHex string, expireHeight uint32, txType string, + ctx context.Context, + stakingTxHashHex string, + expireHeight uint32, + subState types.DelegationSubState, ) error /** * FindExpiredDelegations finds the expired delegations. diff --git a/internal/db/model/delegation.go b/internal/db/model/delegation.go index c504920..021c7cc 100644 --- a/internal/db/model/delegation.go +++ b/internal/db/model/delegation.go @@ -32,6 +32,7 @@ type BTCDelegationDetails struct { StartHeight uint32 `bson:"start_height"` EndHeight uint32 `bson:"end_height"` State types.DelegationState `bson:"state"` + SubState types.DelegationSubState `bson:"sub_state,omitempty"` ParamsVersion uint32 `bson:"params_version"` UnbondingTime uint32 `bson:"unbonding_time"` UnbondingTx string `bson:"unbonding_tx"` diff --git a/internal/db/model/timelock.go b/internal/db/model/timelock.go index f8c41f1..2b4ece3 100644 --- a/internal/db/model/timelock.go +++ b/internal/db/model/timelock.go @@ -1,15 +1,19 @@ package model +import "github.com/babylonlabs-io/babylon-staking-indexer/internal/types" + type TimeLockDocument struct { - StakingTxHashHex string `bson:"_id"` // Primary key - ExpireHeight uint32 `bson:"expire_height"` - TxType string `bson:"tx_type"` + StakingTxHashHex string `bson:"_id"` // Primary key + ExpireHeight uint32 `bson:"expire_height"` + DelegationSubState types.DelegationSubState `bson:"delegation_sub_state"` } -func NewTimeLockDocument(stakingTxHashHex string, expireHeight uint32, txType string) *TimeLockDocument { +func NewTimeLockDocument( + stakingTxHashHex string, expireHeight uint32, subState types.DelegationSubState, +) *TimeLockDocument { return &TimeLockDocument{ - StakingTxHashHex: stakingTxHashHex, - ExpireHeight: expireHeight, - TxType: txType, + StakingTxHashHex: stakingTxHashHex, + ExpireHeight: expireHeight, + DelegationSubState: subState, } } diff --git a/internal/db/timelock.go b/internal/db/timelock.go index 866a851..c94e128 100644 --- a/internal/db/timelock.go +++ b/internal/db/timelock.go @@ -6,16 +6,19 @@ import ( "fmt" "github.com/babylonlabs-io/babylon-staking-indexer/internal/db/model" + "github.com/babylonlabs-io/babylon-staking-indexer/internal/types" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" ) func (db *Database) SaveNewTimeLockExpire( - ctx context.Context, stakingTxHashHex string, - expireHeight uint32, txType string, + ctx context.Context, + stakingTxHashHex string, + expireHeight uint32, + subState types.DelegationSubState, ) error { - tlDoc := model.NewTimeLockDocument(stakingTxHashHex, expireHeight, txType) + tlDoc := model.NewTimeLockDocument(stakingTxHashHex, expireHeight, subState) _, err := db.client.Database(db.dbName). Collection(model.TimeLockCollection). InsertOne(ctx, tlDoc) diff --git a/internal/services/delegation.go b/internal/services/delegation.go index 136854d..a9c9496 100644 --- a/internal/services/delegation.go +++ b/internal/services/delegation.go @@ -158,7 +158,10 @@ func (s *Service) processCovenantQuorumReachedEvent( } if dbErr := s.db.UpdateBTCDelegationState( - ctx, covenantQuorumReachedEvent.StakingTxHash, newState, + ctx, + covenantQuorumReachedEvent.StakingTxHash, + newState, + nil, ); dbErr != nil { return types.NewError( http.StatusInternalServerError, @@ -205,12 +208,14 @@ func (s *Service) processBTCDelegationInclusionProofReceivedEvent( } if dbErr := s.db.UpdateBTCDelegationDetails( - ctx, inclusionProofEvent.StakingTxHash, model.FromEventBTCDelegationInclusionProofReceived(inclusionProofEvent), + ctx, + inclusionProofEvent.StakingTxHash, + model.FromEventBTCDelegationInclusionProofReceived(inclusionProofEvent), ); dbErr != nil { return types.NewError( http.StatusInternalServerError, types.InternalServiceError, - fmt.Errorf("failed to update BTC delegation state: %w", dbErr), + fmt.Errorf("failed to update BTC delegation details: %w", dbErr), ) } @@ -260,13 +265,15 @@ func (s *Service) processBTCDelegationUnbondedEarlyEvent( ) } + subState := types.SubStateEarlyUnbonding + // Save timelock expire unbondingExpireHeight := uint32(unbondingStartHeight) + delegation.UnbondingTime if err := s.db.SaveNewTimeLockExpire( ctx, delegation.StakingTxHashHex, unbondingExpireHeight, - types.EarlyUnbondingTxType.String(), + subState, ); err != nil { return types.NewError( http.StatusInternalServerError, @@ -280,6 +287,7 @@ func (s *Service) processBTCDelegationUnbondedEarlyEvent( ctx, unbondedEarlyEvent.StakingTxHash, types.StateUnbonding, + &subState, ); err != nil { return types.NewError( http.StatusInternalServerError, @@ -330,12 +338,14 @@ func (s *Service) processBTCDelegationExpiredEvent( return err } + subState := types.SubStateTimelock + // Save timelock expire if err := s.db.SaveNewTimeLockExpire( ctx, delegation.StakingTxHashHex, delegation.EndHeight, - types.ExpiredTxType.String(), + subState, ); err != nil { return types.NewError( http.StatusInternalServerError, @@ -349,6 +359,7 @@ func (s *Service) processBTCDelegationExpiredEvent( ctx, delegation.StakingTxHashHex, types.StateUnbonding, + &subState, ); err != nil { return types.NewError( http.StatusInternalServerError, diff --git a/internal/services/expiry_checker.go b/internal/services/expiry_checker.go index 2405cba..804a8e2 100644 --- a/internal/services/expiry_checker.go +++ b/internal/services/expiry_checker.go @@ -58,7 +58,12 @@ func (s *Service) checkExpiry(ctx context.Context) *types.Error { return consumerErr } - if err := s.db.UpdateBTCDelegationState(ctx, delegation.StakingTxHashHex, types.StateWithdrawable); err != nil { + if err := s.db.UpdateBTCDelegationState( + ctx, + delegation.StakingTxHashHex, + types.StateWithdrawable, + &tlDoc.DelegationSubState, + ); err != nil { log.Error().Err(err).Msg("Error updating BTC delegation state to withdrawable") return types.NewInternalServiceError( fmt.Errorf("failed to update BTC delegation state to withdrawable: %w", err), diff --git a/internal/services/watch_btc_events.go b/internal/services/watch_btc_events.go index a360e1b..d2ce161 100644 --- a/internal/services/watch_btc_events.go +++ b/internal/services/watch_btc_events.go @@ -34,6 +34,7 @@ func (s *Service) watchForSpendStakingTx( quitCtx, spendDetail.SpendingTx, spendDetail.SpenderInputIndex, + uint32(spendDetail.SpendingHeight), delegation, ); err != nil { log.Error().Err(err).Msg("failed to handle spending staking transaction") @@ -62,6 +63,7 @@ func (s *Service) watchForSpendUnbondingTx( if err := s.handleSpendingUnbondingTransaction( quitCtx, spendDetail.SpendingTx, + uint32(spendDetail.SpendingHeight), spendDetail.SpenderInputIndex, delegation, ); err != nil { @@ -79,6 +81,7 @@ func (s *Service) watchForSpendUnbondingTx( func (s *Service) watchForSpendSlashingChange( spendEvent *notifier.SpendEvent, delegation *model.BTCDelegationDetails, + subState types.DelegationSubState, ) { defer s.wg.Done() quitCtx, cancel := s.quitContext() @@ -103,10 +106,12 @@ func (s *Service) watchForSpendSlashingChange( } // Update to withdrawn state + delegationSubState := subState if err := s.db.UpdateBTCDelegationState( quitCtx, delegation.StakingTxHashHex, - types.StateSlashedWithdrawn, + types.StateWithdrawn, + &delegationSubState, ); err != nil { log.Error().Err(err).Msg("failed to update delegation state") return @@ -121,8 +126,9 @@ func (s *Service) watchForSpendSlashingChange( func (s *Service) handleSpendingStakingTransaction( ctx context.Context, - tx *wire.MsgTx, + spendingTx *wire.MsgTx, spendingInputIdx uint32, + spendingHeight uint32, delegation *model.BTCDelegationDetails, ) error { params, err := s.db.GetStakingParams(ctx, delegation.ParamsVersion) @@ -131,7 +137,7 @@ func (s *Service) handleSpendingStakingTransaction( } // First try to validate as unbonding tx - isUnbonding, err := s.IsValidUnbondingTx(tx, delegation, params) + isUnbonding, err := s.IsValidUnbondingTx(spendingTx, delegation, params) if err != nil { return fmt.Errorf("failed to validate unbonding tx: %w", err) } @@ -141,10 +147,10 @@ func (s *Service) handleSpendingStakingTransaction( } // Try to validate as withdrawal transaction - withdrawalErr := s.validateWithdrawalTxFromStaking(tx, spendingInputIdx, delegation, params) + withdrawalErr := s.validateWithdrawalTxFromStaking(spendingTx, spendingInputIdx, delegation, params) if withdrawalErr == nil { // It's a valid withdrawal, process it - return s.handleWithdrawal(ctx, delegation) + return s.handleWithdrawal(ctx, delegation, types.SubStateTimelock) } // If it's not a valid withdrawal, check if it's a valid slashing @@ -153,7 +159,7 @@ func (s *Service) handleSpendingStakingTransaction( } // Try to validate as slashing transaction - if err := s.validateSlashingTxFromStaking(tx, spendingInputIdx, delegation, params); err != nil { + if err := s.validateSlashingTxFromStaking(spendingTx, spendingInputIdx, delegation, params); err != nil { if errors.Is(err, types.ErrInvalidSlashingTx) { // Neither withdrawal nor slashing - this is an invalid spend return fmt.Errorf("transaction is neither valid unbonding, withdrawal, nor slashing: %w", err) @@ -162,12 +168,19 @@ func (s *Service) handleSpendingStakingTransaction( } // It's a valid slashing tx, watch for spending change output - return s.startWatchingSlashingChange(ctx, tx, delegation) + return s.startWatchingSlashingChange( + ctx, + spendingTx, + spendingHeight, + delegation, + types.SubStateTimelockSlashing, + ) } func (s *Service) handleSpendingUnbondingTransaction( ctx context.Context, - tx *wire.MsgTx, + spendingTx *wire.MsgTx, + spendingHeight uint32, spendingInputIdx uint32, delegation *model.BTCDelegationDetails, ) error { @@ -177,10 +190,10 @@ func (s *Service) handleSpendingUnbondingTransaction( } // First try to validate as withdrawal transaction - withdrawalErr := s.validateWithdrawalTxFromUnbonding(tx, delegation, spendingInputIdx, params) + withdrawalErr := s.validateWithdrawalTxFromUnbonding(spendingTx, delegation, spendingInputIdx, params) if withdrawalErr == nil { // It's a valid withdrawal, process it - return s.handleWithdrawal(ctx, delegation) + return s.handleWithdrawal(ctx, delegation, types.SubStateEarlyUnbonding) } // If it's not a valid withdrawal, check if it's a valid slashing @@ -189,7 +202,7 @@ func (s *Service) handleSpendingUnbondingTransaction( } // Try to validate as slashing transaction - if err := s.validateSlashingTxFromUnbonding(tx, delegation, spendingInputIdx, params); err != nil { + if err := s.validateSlashingTxFromUnbonding(spendingTx, delegation, spendingInputIdx, params); err != nil { if errors.Is(err, types.ErrInvalidSlashingTx) { // Neither withdrawal nor slashing - this is an invalid spend return fmt.Errorf("transaction is neither valid withdrawal nor slashing: %w", err) @@ -198,10 +211,20 @@ func (s *Service) handleSpendingUnbondingTransaction( } // It's a valid slashing tx, watch for spending change output - return s.startWatchingSlashingChange(ctx, tx, delegation) + return s.startWatchingSlashingChange( + ctx, + spendingTx, + spendingHeight, + delegation, + types.SubStateEarlyUnbondingSlashing, + ) } -func (s *Service) handleWithdrawal(ctx context.Context, delegation *model.BTCDelegationDetails) error { +func (s *Service) handleWithdrawal( + ctx context.Context, + delegation *model.BTCDelegationDetails, + subState types.DelegationSubState, +) error { delegationState, err := s.db.GetBTCDelegationState(ctx, delegation.StakingTxHashHex) if err != nil { return fmt.Errorf("failed to get delegation state: %w", err) @@ -217,10 +240,17 @@ func (s *Service) handleWithdrawal(ctx context.Context, delegation *model.BTCDel ctx, delegation.StakingTxHashHex, types.StateWithdrawn, + &subState, ) } -func (s *Service) startWatchingSlashingChange(ctx context.Context, slashingTx *wire.MsgTx, delegation *model.BTCDelegationDetails) error { +func (s *Service) startWatchingSlashingChange( + ctx context.Context, + slashingTx *wire.MsgTx, + spendingHeight uint32, + delegation *model.BTCDelegationDetails, + subState types.DelegationSubState, +) error { // Create outpoint for the change output (index 1) changeOutpoint := wire.OutPoint{ Hash: slashingTx.TxHash(), @@ -237,8 +267,24 @@ func (s *Service) startWatchingSlashingChange(ctx context.Context, slashingTx *w return fmt.Errorf("failed to register spend ntfn for slashing change output: %w", err) } + stakingParams, err := s.db.GetStakingParams(ctx, delegation.ParamsVersion) + if err != nil { + return fmt.Errorf("failed to get staking params: %w", err) + } + slashingChangeTimelockExpireHeight := spendingHeight + stakingParams.MinUnbondingTimeBlocks + + // Save timelock expire to mark it as Withdrawn (sub state - timelock_slashing/early_unbonding_slashing) + if err := s.db.SaveNewTimeLockExpire( + ctx, + delegation.StakingTxHashHex, + slashingChangeTimelockExpireHeight, + subState, + ); err != nil { + return fmt.Errorf("failed to save timelock expire: %w", err) + } + s.wg.Add(1) - go s.watchForSpendSlashingChange(spendEv, delegation) + go s.watchForSpendSlashingChange(spendEv, delegation, subState) return nil } diff --git a/internal/types/state.go b/internal/types/state.go index 71dc5fb..0a8d4a8 100644 --- a/internal/types/state.go +++ b/internal/types/state.go @@ -6,14 +6,13 @@ import bbntypes "github.com/babylonlabs-io/babylon/x/btcstaking/types" type DelegationState string const ( - StatePending DelegationState = "PENDING" - StateVerified DelegationState = "VERIFIED" - StateActive DelegationState = "ACTIVE" - StateUnbonding DelegationState = "UNBONDING" - StateWithdrawable DelegationState = "WITHDRAWABLE" - StateWithdrawn DelegationState = "WITHDRAWN" - StateSlashed DelegationState = "SLASHED" - StateSlashedWithdrawn DelegationState = "SLASHED_WITHDRAWN" + StatePending DelegationState = "PENDING" + StateVerified DelegationState = "VERIFIED" + StateActive DelegationState = "ACTIVE" + StateUnbonding DelegationState = "UNBONDING" + StateWithdrawable DelegationState = "WITHDRAWABLE" + StateWithdrawn DelegationState = "WITHDRAWN" + StateSlashed DelegationState = "SLASHED" ) func (s DelegationState) String() string { @@ -66,3 +65,16 @@ func QualifiedStatesForWithdrawable() []DelegationState { func QualifiedStatesForSlashedWithdrawn() []DelegationState { return []DelegationState{StateSlashed} } + +type DelegationSubState string + +const ( + SubStateTimelock DelegationSubState = "TIMELOCK" + SubStateEarlyUnbonding DelegationSubState = "EARLY_UNBONDING" + SubStateTimelockSlashing DelegationSubState = "TIMELOCK_SLASHING" + SubStateEarlyUnbondingSlashing DelegationSubState = "EARLY_UNBONDING_SLASHING" +) + +func (p DelegationSubState) String() string { + return string(p) +} diff --git a/tests/mocks/mock_db_client.go b/tests/mocks/mock_db_client.go index 4c2ece4..1ae5142 100644 --- a/tests/mocks/mock_db_client.go +++ b/tests/mocks/mock_db_client.go @@ -305,17 +305,17 @@ func (_m *DbInterface) SaveNewFinalityProvider(ctx context.Context, fpDoc *model return r0 } -// SaveNewTimeLockExpire provides a mock function with given fields: ctx, stakingTxHashHex, expireHeight, txType -func (_m *DbInterface) SaveNewTimeLockExpire(ctx context.Context, stakingTxHashHex string, expireHeight uint32, txType string) error { - ret := _m.Called(ctx, stakingTxHashHex, expireHeight, txType) +// SaveNewTimeLockExpire provides a mock function with given fields: ctx, stakingTxHashHex, expireHeight, subState +func (_m *DbInterface) SaveNewTimeLockExpire(ctx context.Context, stakingTxHashHex string, expireHeight uint32, subState types.DelegationSubState) error { + ret := _m.Called(ctx, stakingTxHashHex, expireHeight, subState) if len(ret) == 0 { panic("no return value specified for SaveNewTimeLockExpire") } var r0 error - if rf, ok := ret.Get(0).(func(context.Context, string, uint32, string) error); ok { - r0 = rf(ctx, stakingTxHashHex, expireHeight, txType) + if rf, ok := ret.Get(0).(func(context.Context, string, uint32, types.DelegationSubState) error); ok { + r0 = rf(ctx, stakingTxHashHex, expireHeight, subState) } else { r0 = ret.Error(0) } @@ -359,17 +359,17 @@ func (_m *DbInterface) UpdateBTCDelegationDetails(ctx context.Context, stakingTx return r0 } -// UpdateBTCDelegationState provides a mock function with given fields: ctx, stakingTxHash, newState -func (_m *DbInterface) UpdateBTCDelegationState(ctx context.Context, stakingTxHash string, newState types.DelegationState) error { - ret := _m.Called(ctx, stakingTxHash, newState) +// UpdateBTCDelegationState provides a mock function with given fields: ctx, stakingTxHash, newState, subState +func (_m *DbInterface) UpdateBTCDelegationState(ctx context.Context, stakingTxHash string, newState types.DelegationState, subState *types.DelegationSubState) error { + ret := _m.Called(ctx, stakingTxHash, newState, subState) if len(ret) == 0 { panic("no return value specified for UpdateBTCDelegationState") } var r0 error - if rf, ok := ret.Get(0).(func(context.Context, string, types.DelegationState) error); ok { - r0 = rf(ctx, stakingTxHash, newState) + if rf, ok := ret.Get(0).(func(context.Context, string, types.DelegationState, *types.DelegationSubState) error); ok { + r0 = rf(ctx, stakingTxHash, newState, subState) } else { r0 = ret.Error(0) }