From 71ead8e3b4412fbf8b941341cad35735f0957d89 Mon Sep 17 00:00:00 2001
From: Wilmer Paulino <wilmer.paulino@gmail.com>
Date: Mon, 5 Nov 2018 17:07:15 -0800
Subject: [PATCH] wallet/wallet: commit birthday blockstamp upon initial
 sync/recovery

---
 wallet/chainntfns.go |  2 +-
 wallet/wallet.go     | 41 ++++++++++++++++++++++++++++++-----------
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/wallet/chainntfns.go b/wallet/chainntfns.go
index f6b32f0..87715e6 100644
--- a/wallet/chainntfns.go
+++ b/wallet/chainntfns.go
@@ -31,7 +31,7 @@ func (w *Wallet) handleChainNotifications() {
 		// some reason, however, the wallet will not be marked synced
 		// and many methods will error early since the wallet is known
 		// to be out of date.
-		err := w.syncWithChain()
+		err := w.syncWithChain(birthdayStamp)
 		if err != nil && !w.ShuttingDown() {
 			log.Warnf("Unable to synchronize wallet to chain: %v", err)
 		}
diff --git a/wallet/wallet.go b/wallet/wallet.go
index b63b4a2..66dec96 100644
--- a/wallet/wallet.go
+++ b/wallet/wallet.go
@@ -318,9 +318,10 @@ func (w *Wallet) activeData(dbtx walletdb.ReadTx) ([]btcutil.Address, []wtxmgr.C
 }
 
 // syncWithChain brings the wallet up to date with the current chain server
-// connection.  It creates a rescan request and blocks until the rescan has
-// finished.
-func (w *Wallet) syncWithChain() error {
+// connection. It creates a rescan request and blocks until the rescan has
+// finished. The birthday block can be passed in, if set, to ensure we can
+// properly detect if it gets rolled back.
+func (w *Wallet) syncWithChain(birthdayStamp *waddrmgr.BlockStamp) error {
 	chainClient, err := w.requireChainClient()
 	if err != nil {
 		return err
@@ -351,12 +352,6 @@ func (w *Wallet) syncWithChain() error {
 	isRecovery := w.recoveryWindow > 0
 	birthday := w.Manager.Birthday()
 
-	// If an initial sync is attempted, we will try and find the block stamp
-	// of the first block past our birthday. This will be fed into the
-	// rescan to ensure we catch transactions that are sent while performing
-	// the initial sync.
-	var birthdayStamp *waddrmgr.BlockStamp
-
 	// TODO(jrick): How should this handle a synced height earlier than
 	// the chain server best block?
 
@@ -496,6 +491,19 @@ func (w *Wallet) syncWithChain() error {
 						Hash:      *hash,
 						Timestamp: timestamp,
 					}
+
+					log.Debugf("Found birthday block: "+
+						"height=%d, hash=%v",
+						birthdayStamp.Height,
+						birthdayStamp.Hash)
+
+					err := w.Manager.SetBirthdayBlock(
+						ns, *birthdayStamp,
+					)
+					if err != nil {
+						tx.Rollback()
+						return err
+					}
 				}
 
 				// If we are in recovery mode and the check
@@ -647,6 +655,18 @@ func (w *Wallet) syncWithChain() error {
 	// points at the new tip.
 	if birthdayStamp != nil && rollbackStamp.Height <= birthdayStamp.Height {
 		birthdayStamp = &rollbackStamp
+
+		log.Debugf("Found new birthday block after rollback: "+
+			"height=%d, hash=%v", birthdayStamp.Height,
+			birthdayStamp.Hash)
+
+		err := walletdb.Update(w.db, func(tx walletdb.ReadWriteTx) error {
+			ns := tx.ReadWriteBucket(waddrmgrNamespaceKey)
+			return w.Manager.SetBirthdayBlock(ns, *birthdayStamp)
+		})
+		if err != nil {
+			return nil
+		}
 	}
 
 	// Request notifications for connected and disconnected blocks.
@@ -657,8 +677,7 @@ func (w *Wallet) syncWithChain() error {
 	// as well.  I am leaning towards allowing off all rpcclient
 	// notification re-registrations, in which case the code here should be
 	// left as is.
-	err = chainClient.NotifyBlocks()
-	if err != nil {
+	if err := chainClient.NotifyBlocks(); err != nil {
 		return err
 	}