From ebd030923488c5554f7c74ffac324db7f25b5f7f Mon Sep 17 00:00:00 2001
From: Aurora Gaffney <aurora@agaffney.org>
Date: Thu, 15 Aug 2024 18:33:24 -0500
Subject: [PATCH] feat: add block type to blockfetch Block callback

---
 cmd/gouroboros/chainsync.go       | 1 +
 protocol/blockfetch/blockfetch.go | 2 +-
 protocol/blockfetch/client.go     | 2 +-
 3 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/cmd/gouroboros/chainsync.go b/cmd/gouroboros/chainsync.go
index cb0d44fe..7f8a761b 100644
--- a/cmd/gouroboros/chainsync.go
+++ b/cmd/gouroboros/chainsync.go
@@ -324,6 +324,7 @@ func chainSyncRollForwardHandler(
 
 func blockFetchBlockHandler(
 	ctx blockfetch.CallbackContext,
+	blockType uint,
 	blockData ledger.Block,
 ) error {
 	switch block := blockData.(type) {
diff --git a/protocol/blockfetch/blockfetch.go b/protocol/blockfetch/blockfetch.go
index c329d4a6..f1cd247e 100644
--- a/protocol/blockfetch/blockfetch.go
+++ b/protocol/blockfetch/blockfetch.go
@@ -101,7 +101,7 @@ type CallbackContext struct {
 }
 
 // Callback function types
-type BlockFunc func(CallbackContext, ledger.Block) error
+type BlockFunc func(CallbackContext, uint, ledger.Block) error
 type RequestRangeFunc func(CallbackContext, common.Point, common.Point) error
 
 func New(protoOptions protocol.ProtocolOptions, cfg *Config) *BlockFetch {
diff --git a/protocol/blockfetch/client.go b/protocol/blockfetch/client.go
index 1cc7beb3..b2930939 100644
--- a/protocol/blockfetch/client.go
+++ b/protocol/blockfetch/client.go
@@ -191,7 +191,7 @@ func (c *Client) handleBlock(msgGeneric protocol.Message) error {
 	}
 	// We use the callback when requesting ranges and the internal channel for a single block
 	if c.blockUseCallback {
-		if err := c.config.BlockFunc(c.callbackContext, blk); err != nil {
+		if err := c.config.BlockFunc(c.callbackContext, wrappedBlock.Type, blk); err != nil {
 			return err
 		}
 	} else {