diff --git a/protocol/blockfetch/blockfetch.go b/protocol/blockfetch/blockfetch.go index 0762a8a9..6ca34ea3 100644 --- a/protocol/blockfetch/blockfetch.go +++ b/protocol/blockfetch/blockfetch.go @@ -88,6 +88,7 @@ type BlockFetch struct { type Config struct { BlockFunc BlockFunc + BlockRawFunc BlockRawFunc BatchDoneFunc BatchDoneFunc RequestRangeFunc RequestRangeFunc BatchStartTimeout time.Duration @@ -103,6 +104,7 @@ type CallbackContext struct { // Callback function types type BlockFunc func(CallbackContext, uint, ledger.Block) error +type BlockRawFunc func(CallbackContext, uint, []byte) error type BatchDoneFunc func(CallbackContext) error type RequestRangeFunc func(CallbackContext, common.Point, common.Point) error @@ -134,6 +136,12 @@ func WithBlockFunc(blockFunc BlockFunc) BlockFetchOptionFunc { } } +func WithBlockRawFunc(blockRawFunc BlockRawFunc) BlockFetchOptionFunc { + return func(c *Config) { + c.BlockRawFunc = blockRawFunc + } +} + func WithBatchDoneFunc(batchDoneFunc BatchDoneFunc) BlockFetchOptionFunc { return func(c *Config) { c.BatchDoneFunc = batchDoneFunc diff --git a/protocol/blockfetch/client.go b/protocol/blockfetch/client.go index 4bedf0c7..038f7783 100644 --- a/protocol/blockfetch/client.go +++ b/protocol/blockfetch/client.go @@ -254,12 +254,16 @@ func (c *Client) handleBlock(msgGeneric protocol.Message) error { if _, err := cbor.Decode(msg.WrappedBlock, &wrappedBlock); err != nil { return fmt.Errorf("%s: decode error: %s", ProtocolName, err) } - blk, err := ledger.NewBlockFromCbor( - wrappedBlock.Type, - wrappedBlock.RawBlock, - ) - if err != nil { - return err + var block ledger.Block + if !c.blockUseCallback || c.config.BlockFunc != nil { + var err error + block, err = ledger.NewBlockFromCbor( + wrappedBlock.Type, + wrappedBlock.RawBlock, + ) + if err != nil { + return err + } } // Check for shutdown select { @@ -269,11 +273,21 @@ 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, wrappedBlock.Type, blk); err != nil { - return err + if c.config.BlockRawFunc != nil { + if err := c.config.BlockRawFunc(c.callbackContext, wrappedBlock.Type, wrappedBlock.RawBlock); err != nil { + return err + } + } else if c.config.BlockFunc != nil { + if err := c.config.BlockFunc(c.callbackContext, wrappedBlock.Type, block); err != nil { + return err + } + } else { + return fmt.Errorf( + "received block-fetch Block message but no callback function is defined", + ) } } else { - c.blockChan <- blk + c.blockChan <- block } return nil }