Fix withdrawable transactions query (#78)

* Fix withdrawable transactions query

CHANGELOG.md

@@ -35,6 +35,12 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ## Unreleased
 
+### Bug fix
+
+* [#78](https://github.com/babylonlabs-io/btc-staker/pull/78) Fix
+`withdrawable-transactions` query bug, introduced when adding pre-approval
+transactions handling
+
 ## v0.8.0
 
 ### Improvements

itest/e2e_test.go

@@ -1320,6 +1320,48 @@ func TestSendingStakingTransactionWithPreApproval(t *testing.T) {
 	require.NoError(t, err)
 	tm.waitForStakingTxState(t, txHash, proto.TransactionState_DELEGATION_ACTIVE)
 
+	// check that there is not error when qury for withdrawable transactions
+	withdrawableTransactionsResp, err := tm.StakerClient.WithdrawableTransactions(context.Background(), nil, nil)
+	require.NoError(t, err)
+	require.Len(t, withdrawableTransactionsResp.Transactions, 0)
+
+	//  Unbond pre-approval stake
+	resp, err := tm.StakerClient.UnbondStaking(context.Background(), txHash.String())
+	require.NoError(t, err)
+
+	unbondingTxHash, err := chainhash.NewHashFromStr(resp.UnbondingTxHash)
+	require.NoError(t, err)
+
+	require.Eventually(t, func() bool {
+		tx, err := tm.TestRpcClient.GetRawTransaction(unbondingTxHash)
+		if err != nil {
+			return false
+		}
+
+		if tx == nil {
+			return false
+
+		}
+		return true
+	}, 1*time.Minute, eventuallyPollTime)
+
+	block := tm.mineBlock(t)
+	require.Equal(t, 2, len(block.Transactions))
+	require.Equal(t, block.Transactions[1].TxHash(), *unbondingTxHash)
+	go tm.mineNEmptyBlocks(t, staker.UnbondingTxConfirmations, false)
+	tm.waitForStakingTxState(t, txHash, proto.TransactionState_UNBONDING_CONFIRMED_ON_BTC)
+
+	// Spend unbonding tx of pre-approval stake
+	withdrawableTransactionsResp, err = tm.StakerClient.WithdrawableTransactions(context.Background(), nil, nil)
+	require.NoError(t, err)
+	require.Len(t, withdrawableTransactionsResp.Transactions, 1)
+
+	// We can spend unbonding tx immediately as in e2e test, finalization time is 4 blocks and we locked it
+	// finalization time + 1 i.e 5 blocks, but to consider unboning tx as confirmed we need to wait for 6 blocks
+	// so at this point time lock should already have passed
+	tm.spendStakingTxWithHash(t, txHash)
+	go tm.mineNEmptyBlocks(t, staker.SpendStakeTxConfirmations, false)
+	tm.waitForStakingTxState(t, txHash, proto.TransactionState_SPENT_ON_BTC)
 }
 
 func TestMultipleWithdrawableStakingTransactions(t *testing.T) {

staker/babylontypes.go

@@ -352,10 +352,32 @@ func (app *StakerApp) activateVerifiedDelegation(
 					"stakingTxHash": stakingTxHash,
 				}).Debug("Delegation has been activated on the Babylon chain")
 
-				utils.PushOrQuit[*delegationActiveOnBabylonEvent](
-					app.delegationActiveOnBabylonEvChan,
-					&delegationActiveOnBabylonEvent{
+				info, status, err := app.wc.TxDetails(stakingTxHash, stakingTransaction.TxOut[stakingOutputIndex].PkScript)
+
+				if err != nil {
+					app.logger.WithFields(logrus.Fields{
+						"stakingTxHash": stakingTxHash,
+						"err":           err,
+					}).Error("error getting staking transaction details from btc chain")
+
+					// failed to retrieve transaction details from bitcoind node, most probably
+					// connection error, we will try again in next iteration
+					continue
+				}
+
+				if status != walletcontroller.TxInChain {
+					app.logger.WithFields(logrus.Fields{
+						"stakingTxHash": stakingTxHash,
+					}).Debug("Staking transaction active on babylon, but not on btc chain. Waiting for btc node to catch up")
+					continue
+				}
+
+				utils.PushOrQuit[*delegationActivatedPreApprovalEvent](
+					app.delegationActivatedPreApprovalEvChan,
+					&delegationActivatedPreApprovalEvent{
 						stakingTxHash: *stakingTxHash,
+						blockHash:     *info.BlockHash,
+						blockHeight:   info.BlockHeight,
 					},
 					app.quit,
 				)

staker/events.go

@@ -15,7 +15,8 @@ type StakingEvent interface {
 
 var _ StakingEvent = (*stakingTxBtcConfirmedEvent)(nil)
 var _ StakingEvent = (*delegationSubmittedToBabylonEvent)(nil)
-var _ StakingEvent = (*delegationActiveOnBabylonEvent)(nil)
+var _ StakingEvent = (*delegationActivatedPostApprovalEvent)(nil)
+var _ StakingEvent = (*delegationActivatedPreApprovalEvent)(nil)
 var _ StakingEvent = (*unbondingTxSignaturesConfirmedOnBabylonEvent)(nil)
 var _ StakingEvent = (*unbondingTxConfirmedOnBtcEvent)(nil)
 var _ StakingEvent = (*spendStakeTxConfirmedOnBtcEvent)(nil)
@@ -121,14 +122,28 @@ func (app *StakerApp) logStakingEventProcessed(event StakingEvent) {
 	}).Debug("Processed staking event")
 }
 
-type delegationActiveOnBabylonEvent struct {
+type delegationActivatedPostApprovalEvent struct {
 	stakingTxHash chainhash.Hash
 }
 
-func (event *delegationActiveOnBabylonEvent) EventId() chainhash.Hash {
+func (event *delegationActivatedPostApprovalEvent) EventId() chainhash.Hash {
 	return event.stakingTxHash
 }
 
-func (event *delegationActiveOnBabylonEvent) EventDesc() string {
-	return "DELEGATION_ACTIVE_ON_BABYLON_EVENT"
+func (event *delegationActivatedPostApprovalEvent) EventDesc() string {
+	return "DELEGATION_ACTIVE_ON_BABYLON_POST_APPROVAL_EVENT"
+}
+
+type delegationActivatedPreApprovalEvent struct {
+	stakingTxHash chainhash.Hash
+	blockHash     chainhash.Hash
+	blockHeight   uint32
+}
+
+func (event *delegationActivatedPreApprovalEvent) EventId() chainhash.Hash {
+	return event.stakingTxHash
+}
+
+func (event *delegationActivatedPreApprovalEvent) EventDesc() string {
+	return "DELEGATION_ACTIVE_ON_BABYLON_PRE_APPROVAL_EVENT"
 }

staker/stakerapp.go

@@ -126,7 +126,8 @@ type StakerApp struct {
 	stakingRequestedCmdChan                       chan *stakingRequestCmd
 	stakingTxBtcConfirmedEvChan                   chan *stakingTxBtcConfirmedEvent
 	delegationSubmittedToBabylonEvChan            chan *delegationSubmittedToBabylonEvent
-	delegationActiveOnBabylonEvChan               chan *delegationActiveOnBabylonEvent
+	delegationActivatedPostApprovalEvChan         chan *delegationActivatedPostApprovalEvent
+	delegationActivatedPreApprovalEvChan          chan *delegationActivatedPreApprovalEvent
 	unbondingTxSignaturesConfirmedOnBabylonEvChan chan *unbondingTxSignaturesConfirmedOnBabylonEvent
 	unbondingTxConfirmedOnBtcEvChan               chan *unbondingTxConfirmedOnBtcEvent
 	spendStakeTxConfirmedOnBtcEvChan              chan *spendStakeTxConfirmedOnBtcEvent
@@ -234,18 +235,17 @@ func NewStakerAppFromDeps(
 
 		// event for when delegation is sent to babylon and included in babylon
 		delegationSubmittedToBabylonEvChan: make(chan *delegationSubmittedToBabylonEvent),
-		// event for when delegation is active on babylon after being verified
-		delegationActiveOnBabylonEvChan: make(chan *delegationActiveOnBabylonEvent),
+		// event for when delegation is active on babylon after going through post approval flow
+		delegationActivatedPostApprovalEvChan: make(chan *delegationActivatedPostApprovalEvent),
+		// event for when delegation is active on babylon after going through pre approval flow
+		delegationActivatedPreApprovalEvChan: make(chan *delegationActivatedPreApprovalEvent),
 		// event emitte	d upon transaction which spends staking transaction is confirmed on BTC
 		spendStakeTxConfirmedOnBtcEvChan: make(chan *spendStakeTxConfirmedOnBtcEvent),
-
 		// channel which receives unbonding signatures from covenant for unbonding
 		// transaction
 		unbondingTxSignaturesConfirmedOnBabylonEvChan: make(chan *unbondingTxSignaturesConfirmedOnBabylonEvent),
-
 		// channel which receives confirmation that unbonding transaction was confirmed on BTC
 		unbondingTxConfirmedOnBtcEvChan: make(chan *unbondingTxConfirmedOnBtcEvent),
-
 		// channel which receives critical errors, critical errors are errors which we do not know
 		// how to handle, so we just log them. It is up to user to investigate, what had happend
 		// and report the situation
@@ -1542,9 +1542,9 @@ func (app *StakerApp) handleStakingEvents() {
 				app.wg.Add(1)
 				go func(hash chainhash.Hash) {
 					defer app.wg.Done()
-					utils.PushOrQuit[*delegationActiveOnBabylonEvent](
-						app.delegationActiveOnBabylonEvChan,
-						&delegationActiveOnBabylonEvent{
+					utils.PushOrQuit[*delegationActivatedPostApprovalEvent](
+						app.delegationActivatedPostApprovalEvChan,
+						&delegationActivatedPostApprovalEvent{
 							stakingTxHash: hash,
 						},
 						app.quit,
@@ -1587,7 +1587,7 @@ func (app *StakerApp) handleStakingEvents() {
 			}
 			app.logStakingEventProcessed(ev)
 
-		case ev := <-app.delegationActiveOnBabylonEvChan:
+		case ev := <-app.delegationActivatedPostApprovalEvChan:
 			app.logStakingEventReceived(ev)
 			if err := app.txTracker.SetDelegationActiveOnBabylon(&ev.stakingTxHash); err != nil {
 				// TODO: handle this error somehow, it means we received spend stake confirmation for tx which we do not store
@@ -1597,6 +1597,18 @@ func (app *StakerApp) handleStakingEvents() {
 			app.m.DelegationsActivatedOnBabylon.Inc()
 			app.logStakingEventProcessed(ev)
 
+		case ev := <-app.delegationActivatedPreApprovalEvChan:
+			app.logStakingEventReceived(ev)
+			if err := app.txTracker.SetDelegationActiveOnBabylonAndConfirmedOnBtc(
+				&ev.stakingTxHash,
+				&ev.blockHash,
+				ev.blockHeight,
+			); err != nil {
+				app.logger.Fatalf("Error setting state for tx %s: %s", ev.stakingTxHash, err)
+			}
+
+			app.logStakingEventProcessed(ev)
+
 		case ev := <-app.criticalErrorEvChan:
 			// if error is context.Canceled, it means one of started child go-routines
 			// received quit signal and is shutting down. We just ignore it.

staker/types.go

@@ -242,7 +242,7 @@ func createSpendStakeTxFromStoredTx(
 	// This is to cover cases:
 	// - staker is unable to sent delegation to babylon
 	// - staking transaction on babylon fail to get covenant signatures
-	if storedtx.StakingTxConfirmedOnBtc() {
+	if storedtx.StakingTxConfirmedOnBtc() && !storedtx.UnbondingTxConfirmedOnBtc() {
 		stakingInfo, err := staking.BuildStakingInfo(
 			stakerBtcPk,
 			storedtx.FinalityProvidersBtcPks,
@@ -284,7 +284,7 @@ func createSpendStakeTxFromStoredTx(
 			fundingOutput:          storedtx.StakingTx.TxOut[storedtx.StakingOutputIndex],
 			calculatedFee:          *calculatedFee,
 		}, nil
-	} else if storedtx.IsUnbonded() {
+	} else if storedtx.StakingTxConfirmedOnBtc() && storedtx.UnbondingTxConfirmedOnBtc() {
 		data := storedtx.UnbondingTxData
 
 		unbondingInfo, err := staking.BuildUnbondingInfo(

stakerdb/trackedtranactionstore.go

@@ -133,14 +133,16 @@ type StoredTransaction struct {
 
 // StakingTxConfirmedOnBtc returns true only if staking transaction was sent and confirmed on bitcoin
 func (t *StoredTransaction) StakingTxConfirmedOnBtc() bool {
-	return t.State == proto.TransactionState_SENT_TO_BABYLON ||
-		t.State == proto.TransactionState_DELEGATION_ACTIVE ||
-		t.State == proto.TransactionState_CONFIRMED_ON_BTC
+	return t.StakingTxConfirmationInfo != nil
 }
 
-// IsUnbonded returns true only if unbonding transaction was sent and confirmed on bitcoin
-func (t *StoredTransaction) IsUnbonded() bool {
-	return t.State == proto.TransactionState_UNBONDING_CONFIRMED_ON_BTC
+// UnbondingTxConfirmedOnBtc returns true only if unbonding transaction was sent and confirmed on bitcoin
+func (t *StoredTransaction) UnbondingTxConfirmedOnBtc() bool {
+	if t.UnbondingTxData == nil {
+		return false
+	}
+
+	return t.UnbondingTxData.UnbondingTxConfirmationInfo != nil
 }
 
 type WatchedTransactionData struct {
@@ -1020,6 +1022,23 @@ func (c *TrackedTransactionStore) SetDelegationActiveOnBabylon(txHash *chainhash
 	return c.setTxState(txHash, setTxSpentOnBtc)
 }
 
+func (c *TrackedTransactionStore) SetDelegationActiveOnBabylonAndConfirmedOnBtc(
+	txHash *chainhash.Hash,
+	blockHash *chainhash.Hash,
+	blockHeight uint32,
+) error {
+	setDelegationActiveOnBabylon := func(tx *proto.TrackedTransaction) error {
+		tx.State = proto.TransactionState_DELEGATION_ACTIVE
+		tx.StakingTxBtcConfirmationInfo = &proto.BTCConfirmationInfo{
+			BlockHash:   blockHash.CloneBytes(),
+			BlockHeight: blockHeight,
+		}
+		return nil
+	}
+
+	return c.setTxState(txHash, setDelegationActiveOnBabylon)
+}
+
 func (c *TrackedTransactionStore) SetTxUnbondingSignaturesReceived(
 	txHash *chainhash.Hash,
 	covenantSignatures []PubKeySigPair,
@@ -1222,10 +1241,10 @@ func (c *TrackedTransactionStore) QueryStoredTransactions(q StoredTransactionQue
 					return false, nil
 				}
 
-				if txFromDb.StakingTxConfirmedOnBtc() {
+				if txFromDb.StakingTxConfirmedOnBtc() && !txFromDb.UnbondingTxConfirmedOnBtc() {
 					scriptTimeLock = txFromDb.StakingTime
 					confirmationHeight = txFromDb.StakingTxConfirmationInfo.Height
-				} else if txFromDb.IsUnbonded() {
+				} else if txFromDb.StakingTxConfirmedOnBtc() && txFromDb.UnbondingTxConfirmedOnBtc() {
 					scriptTimeLock = txFromDb.UnbondingTxData.UnbondingTime
 					confirmationHeight = txFromDb.UnbondingTxData.UnbondingTxConfirmationInfo.Height
 				} else {