From 4eac8577cd10f7263e168aa930cc0a5404c8af43 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Wed, 11 Aug 2021 16:04:20 -0700
Subject: [PATCH] chainntnfs: update impl independent functions to populate new
 block header

---
 chainntnfs/interface.go | 38 ++++++++++++++++++++++++++++++--------
 1 file changed, 30 insertions(+), 8 deletions(-)

diff --git a/chainntnfs/interface.go b/chainntnfs/interface.go
index fc365e70..724d7709 100644
--- a/chainntnfs/interface.go
+++ b/chainntnfs/interface.go
@@ -492,8 +492,9 @@ func RewindChain(chainConn ChainConn, txNotifier *TxNotifier,
 	currBestBlock BlockEpoch, targetHeight int32) (BlockEpoch, error) {
 
 	newBestBlock := BlockEpoch{
-		Height: currBestBlock.Height,
-		Hash:   currBestBlock.Hash,
+		Height:      currBestBlock.Height,
+		Hash:        currBestBlock.Hash,
+		BlockHeader: currBestBlock.BlockHeader,
 	}
 
 	for height := currBestBlock.Height; height > targetHeight; height-- {
@@ -503,6 +504,11 @@ func RewindChain(chainConn ChainConn, txNotifier *TxNotifier,
 				"find blockhash for disconnected height=%d: %v",
 				height, err)
 		}
+		header, err := chainConn.GetBlockHeader(hash)
+		if err != nil {
+			return newBestBlock, fmt.Errorf("unable to get block "+
+				"header for height=%v", height-1)
+		}
 
 		Log.Infof("Block disconnected from main chain: "+
 			"height=%v, sha=%v", height, newBestBlock.Hash)
@@ -515,7 +521,9 @@ func RewindChain(chainConn ChainConn, txNotifier *TxNotifier,
 		}
 		newBestBlock.Height = height - 1
 		newBestBlock.Hash = hash
+		newBestBlock.BlockHeader = header
 	}
+
 	return newBestBlock, nil
 }
 
@@ -539,8 +547,9 @@ func HandleMissedBlocks(chainConn ChainConn, txNotifier *TxNotifier,
 		// If a reorg causes our best hash to be incorrect, rewind the
 		// chain so our best block is set to the closest common
 		// ancestor, then dispatch notifications from there.
-		hashAtBestHeight, err :=
-			chainConn.GetBlockHash(int64(currBestBlock.Height))
+		hashAtBestHeight, err := chainConn.GetBlockHash(
+			int64(currBestBlock.Height),
+		)
 		if err != nil {
 			return currBestBlock, nil, fmt.Errorf("unable to find "+
 				"blockhash for height=%d: %v",
@@ -555,8 +564,9 @@ func HandleMissedBlocks(chainConn ChainConn, txNotifier *TxNotifier,
 				"common ancestor: %v", err)
 		}
 
-		currBestBlock, err = RewindChain(chainConn, txNotifier,
-			currBestBlock, startingHeight)
+		currBestBlock, err = RewindChain(
+			chainConn, txNotifier, currBestBlock, startingHeight,
+		)
 		if err != nil {
 			return currBestBlock, nil, fmt.Errorf("unable to "+
 				"rewind chain: %v", err)
@@ -592,8 +602,20 @@ func getMissedBlocks(chainConn ChainConn, startingHeight,
 			return nil, fmt.Errorf("unable to find blockhash for "+
 				"height=%d: %v", height, err)
 		}
-		missedBlocks = append(missedBlocks,
-			BlockEpoch{Hash: hash, Height: height})
+		header, err := chainConn.GetBlockHeader(hash)
+		if err != nil {
+			return nil, fmt.Errorf("unable to find block header "+
+				"for height=%d: %v", height, err)
+		}
+
+		missedBlocks = append(
+			missedBlocks,
+			BlockEpoch{
+				Hash:        hash,
+				Height:      height,
+				BlockHeader: header,
+			},
+		)
 	}
 
 	return missedBlocks, nil