Fix handling of the restart during unbonding (#20)

* Fix handling of the restart during unbonding

itest/e2e_test.go

@@ -356,10 +356,23 @@ func (tm *TestManager) Stop(t *testing.T) {
 }
 
 func (tm *TestManager) RestartApp(t *testing.T) {
+	// Restart the app with no-op action
+	tm.RestartAppWithAction(t, func(t *testing.T) {})
+}
+
+// RestartAppWithAction:
+// 1. Stop the staker app
+// 2. Perform provided action. Warning:this action must not use staker app as
+// app is stopped at this point
+// 3. Start the staker app
+func (tm *TestManager) RestartAppWithAction(t *testing.T, action func(t *testing.T)) {
 	// First stop the app
 	tm.serverStopper.RequestShutdown()
 	tm.wg.Wait()
 
+	// Perform the action
+	action(t)
+
 	// Now reset all components and start again
 	logger := logrus.New()
 	logger.SetLevel(logrus.DebugLevel)
@@ -1600,3 +1613,73 @@ func TestSendingStakingTransaction_Restaking(t *testing.T) {
 	tm.insertCovenantSigForDelegation(t, pend[0])
 	tm.waitForStakingTxState(t, txHash, proto.TransactionState_DELEGATION_ACTIVE)
 }
+
+func TestRecoverAfterRestartDuringWithdrawal(t *testing.T) {
+	// need to have at least 300 block on testnet as only then segwit is activated.
+	// Mature output is out which has 100 confirmations, which means 200mature outputs
+	// will generate 300 blocks
+	numMatureOutputs := uint32(200)
+	tm := StartManager(t, numMatureOutputs)
+	defer tm.Stop(t)
+	tm.insertAllMinedBlocksToBabylon(t)
+
+	cl := tm.Sa.BabylonController()
+	params, err := cl.Params()
+	require.NoError(t, err)
+	stakingTime := uint16(staker.GetMinStakingTime(params))
+
+	testStakingData := tm.getTestStakingData(t, tm.WalletPrivKey.PubKey(), stakingTime, 10000, 1)
+
+	hashed, err := chainhash.NewHash(datagen.GenRandomByteArray(r, 32))
+	require.NoError(t, err)
+	scr, err := txscript.PayToTaprootScript(tm.CovenantPrivKeys[0].PubKey())
+	require.NoError(t, err)
+	_, st, erro := tm.Sa.Wallet().TxDetails(hashed, scr)
+	// query for exsisting tx is not an error, proper state should be returned
+	require.NoError(t, erro)
+	require.Equal(t, st, walletcontroller.TxNotFound)
+
+	tm.createAndRegisterFinalityProviders(t, testStakingData)
+
+	txHash := tm.sendStakingTxBTC(t, testStakingData)
+
+	go tm.mineNEmptyBlocks(t, params.ConfirmationTimeBlocks, true)
+	// must wait for all covenant signatures to be received, to be able to unbond
+	tm.waitForStakingTxState(t, txHash, proto.TransactionState_SENT_TO_BABYLON)
+
+	pend, err := tm.BabylonClient.QueryPendingBTCDelegations()
+	require.NoError(t, err)
+	require.Len(t, pend, 1)
+	// need to activate delegation to unbond
+	tm.insertCovenantSigForDelegation(t, pend[0])
+	tm.waitForStakingTxState(t, txHash, proto.TransactionState_DELEGATION_ACTIVE)
+
+	// Unbond staking transaction and wait for it to be included in mempool
+	feeRate := 2000
+	unbondResponse, err := tm.StakerClient.UnbondStaking(context.Background(), txHash.String(), &feeRate)
+	require.NoError(t, err)
+	unbondingTxHash, err := chainhash.NewHashFromStr(unbondResponse.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)
+
+	tm.RestartAppWithAction(t, func(t *testing.T) {
+		// unbodning tx got confirmed during the stop period
+		_ = tm.mineNEmptyBlocks(t, staker.UnbondingTxConfirmations+1, false)
+	})
+
+	tm.waitForStakingTxState(t, txHash, proto.TransactionState_UNBONDING_CONFIRMED_ON_BTC)
+	// it should be possible ot spend from unbonding tx
+	tm.spendStakingTxWithHash(t, txHash)
+}

staker/stakerapp.go

@@ -515,9 +515,16 @@ func (app *StakerApp) handleBtcTxInfo(
 	return nil
 }
 
+func (app *StakerApp) mustSetTxSpentOnBtc(hash *chainhash.Hash) {
+	if err := app.txTracker.SetTxSpentOnBtc(hash); err != nil {
+		app.logger.Fatalf("Error setting transaction spent on btc: %s", err)
+	}
+}
+
 // TODO: We should also handle case when btc node or babylon node lost data and start from scratch
 // i.e keep track what is last known block height on both chains and detect if after restart
 // for some reason they are behind staker
+// TODO: Refactor this functions after adding unit tests to stakerapp
 func (app *StakerApp) checkTransactionsStatus() error {
 	stakingParams, err := app.babylonClient.Params()
 
@@ -564,10 +571,16 @@ func (app *StakerApp) checkTransactionsStatus() error {
 			})
 			return nil
 		case proto.TransactionState_DELEGATION_ACTIVE:
-			// we recevied all necessary data from babylon nothing to do here
+			transactionsOnBabylon = append(transactionsOnBabylon, &stakingDbInfo{
+				stakingTxHash:  &stakingTxHash,
+				stakingTxState: tx.State,
+			})
 			return nil
 		case proto.TransactionState_UNBONDING_CONFIRMED_ON_BTC:
-			// unbonding tx was sent to babylon, received all signatures and was confirmed on btc, nothing to do here
+			transactionsOnBabylon = append(transactionsOnBabylon, &stakingDbInfo{
+				stakingTxHash:  &stakingTxHash,
+				stakingTxState: tx.State,
+			})
 			return nil
 		case proto.TransactionState_SPENT_ON_BTC:
 			// nothing to do, staking transaction is already spent
@@ -669,6 +682,98 @@ func (app *StakerApp) checkTransactionsStatus() error {
 			// we crashed after succesful send to babaylon, restart checking for unbonding signatures
 			app.wg.Add(1)
 			go app.checkForUnbondingTxSignaturesOnBabylon(stakingTxHash)
+		} else if localInfo.stakingTxState == proto.TransactionState_DELEGATION_ACTIVE {
+			// delegation was sent to Babylon and activated by covenants, check whether we:
+			// - did not spend tx before restart
+			// - did not send unbonding tx before restart
+			stakingTxHash := localInfo.stakingTxHash
+			tx, _ := app.mustGetTransactionAndStakerAddress(stakingTxHash)
+
+			// 1. First check if staking output is still unspent on BTC chain
+			stakingOutputSpent, err := app.wc.OutputSpent(stakingTxHash, tx.StakingOutputIndex)
+
+			if err != nil {
+				return err
+			}
+
+			if !stakingOutputSpent {
+				// If the staking output is unspent, then it means that delegation is
+				// sitll considered active. We can move forward without to next transaction
+				// and leave the state as it is.
+				continue
+			}
+
+			// 2. Staking output has been spent, we need to check whether this is unbonding
+			// or withdrawal transaction
+			unbondingTxHash := tx.UnbondingTxData.UnbondingTx.TxHash()
+
+			_, unbondingTxStatus, err := app.wc.TxDetails(
+				&unbondingTxHash,
+				// unbonding tx always have only one output
+				tx.UnbondingTxData.UnbondingTx.TxOut[0].PkScript,
+			)
+
+			if err != nil {
+				return err
+			}
+
+			if unbondingTxStatus == walletcontroller.TxNotFound {
+				// no unbonding tx on chain and staking output already spent, most probably
+				// staking transaction has been withdrawn, update state in db
+				app.mustSetTxSpentOnBtc(stakingTxHash)
+				continue
+			}
+
+			unbondingOutputSpent, err := app.wc.OutputSpent(&unbondingTxHash, 0)
+
+			if err != nil {
+				return err
+			}
+
+			if unbondingOutputSpent {
+				app.mustSetTxSpentOnBtc(stakingTxHash)
+				continue
+			}
+
+			// At this point:
+			// - staking output is spent
+			// - unbonding tx has been found in the btc chain
+			// - unbonding output is not spent
+			// we can start waiting for unbonding tx confirmation
+			ev, err := app.notifier.RegisterConfirmationsNtfn(
+				&unbondingTxHash,
+				tx.UnbondingTxData.UnbondingTx.TxOut[0].PkScript,
+				UnbondingTxConfirmations,
+				// unbonding transactions will for sure be included after staking tranasction
+				tx.StakingTxConfirmationInfo.Height,
+			)
+
+			if err != nil {
+				return err
+			}
+
+			// unbonding tx is in mempool, wait for confirmation and inform event
+			// loop about it
+			app.wg.Add(1)
+			go app.waitForUnbondingTxConfirmation(
+				ev,
+				tx.UnbondingTxData,
+				stakingTxHash,
+			)
+		} else if localInfo.stakingTxState == proto.TransactionState_UNBONDING_CONFIRMED_ON_BTC {
+			stakingTxHash := localInfo.stakingTxHash
+			tx, _ := app.mustGetTransactionAndStakerAddress(stakingTxHash)
+			unbondingTxHash := tx.UnbondingTxData.UnbondingTx.TxHash()
+
+			unbondingOutputSpent, err := app.wc.OutputSpent(&unbondingTxHash, 0)
+
+			if err != nil {
+				return err
+			}
+
+			if unbondingOutputSpent {
+				app.mustSetTxSpentOnBtc(stakingTxHash)
+			}
 		} else {
 			// we should not have any other state here, so kill app
 			return fmt.Errorf("unexpected local transaction state: %s, expected: %s", localInfo.stakingTxState, proto.TransactionState_SENT_TO_BABYLON)
@@ -930,11 +1035,13 @@ func (app *StakerApp) sendUnbondingTxToBtc(
 	return notificationEv, nil
 }
 
+// waitForUnbondingTxConfirmation blocks until unbonding tx is confirmed on btc chain.
 func (app *StakerApp) waitForUnbondingTxConfirmation(
 	waitEv *notifier.ConfirmationEvent,
 	unbondingData *stakerdb.UnbondingStoreData,
 	stakingTxHash *chainhash.Hash,
 ) {
+	defer app.wg.Done()
 	defer waitEv.Cancel()
 	unbondingTxHash := unbondingData.UnbondingTx.TxHash()
 
@@ -996,7 +1103,8 @@ func (app *StakerApp) sendUnbondingTxToBtcTask(
 		return
 	}
 
-	app.waitForUnbondingTxConfirmation(
+	app.wg.Add(1)
+	go app.waitForUnbondingTxConfirmation(
 		waitEv,
 		unbondingData,
 		stakingTxHash,

walletcontroller/client.go

@@ -304,6 +304,21 @@ func (w *RpcWalletController) SignBip322NativeSegwit(msg []byte, address btcutil
 	return signed.TxIn[0].Witness, nil
 }
 
+func (w *RpcWalletController) OutputSpent(
+	txHash *chainhash.Hash,
+	outputIdx uint32,
+) (bool, error) {
+	res, err := w.Client.GetTxOut(
+		txHash, outputIdx, true,
+	)
+
+	if err != nil {
+		return false, err
+	}
+
+	return res == nil, nil
+}
+
 // TODO: Temporary implementation to encapsulate signing of taproot spending transaction, it will be replaced with PSBT
 // signing in the future
 func (w *RpcWalletController) SignOneInputTaprootSpendingTransaction(req *TaprootSigningRequest) (*TaprootSigningResult, error) {

walletcontroller/interface.go

@@ -57,4 +57,8 @@ type WalletController interface {
 	// SignOneInputTaprootSpendingTransaction signs transactions with one taproot input that
 	// uses script spending path.
 	SignOneInputTaprootSpendingTransaction(req *TaprootSigningRequest) (*TaprootSigningResult, error)
+	OutputSpent(
+		txHash *chainhash.Hash,
+		outputIdx uint32,
+	) (bool, error)
 }