Pass along whether message is a response to the handler

protocol/blockfetch/blockfetch.go

@@ -102,7 +102,7 @@ func New(options protocol.ProtocolOptions, callbackConfig *BlockFetchCallbackCon
 	return b
 }
 
-func (b *BlockFetch) messageHandler(msg protocol.Message) error {
+func (b *BlockFetch) messageHandler(msg protocol.Message, isResponse bool) error {
 	var err error
 	switch msg.Type() {
 	case MESSAGE_TYPE_START_BATCH:

protocol/chainsync/chainsync.go

@@ -137,7 +137,7 @@ func New(options protocol.ProtocolOptions, callbackConfig *ChainSyncCallbackConf
 	return c
 }
 
-func (c *ChainSync) messageHandler(msg protocol.Message) error {
+func (c *ChainSync) messageHandler(msg protocol.Message, isResponse bool) error {
 	var err error
 	switch msg.Type() {
 	case MESSAGE_TYPE_AWAIT_REPLY:

protocol/handshake/handshake.go

@@ -75,7 +75,7 @@ func New(options protocol.ProtocolOptions, allowedVersions []uint16) *Handshake
 	return h
 }
 
-func (h *Handshake) handleMessage(msg protocol.Message) error {
+func (h *Handshake) handleMessage(msg protocol.Message, isResponse bool) error {
 	var err error
 	switch msg.Type() {
 	case MESSAGE_TYPE_PROPOSE_VERSIONS:

protocol/keepalive/keepalive.go

@@ -85,7 +85,7 @@ func New(options protocol.ProtocolOptions, callbackConfig *KeepAliveCallbackConf
 	return k
 }
 
-func (k *KeepAlive) messageHandler(msg protocol.Message) error {
+func (k *KeepAlive) messageHandler(msg protocol.Message, isResponse bool) error {
 	var err error
 	switch msg.Type() {
 	case MESSAGE_TYPE_KEEP_ALIVE:

protocol/localtxsubmission/localtxsubmission.go

@@ -82,7 +82,7 @@ func New(options protocol.ProtocolOptions, callbackConfig *CallbackConfig) *Loca
 	return l
 }
 
-func (l *LocalTxSubmission) messageHandler(msg protocol.Message) error {
+func (l *LocalTxSubmission) messageHandler(msg protocol.Message, isResponse bool) error {
 	var err error
 	switch msg.Type() {
 	case MESSAGE_TYPE_SUBMIT_TX:

protocol/protocol.go

@@ -54,7 +54,7 @@ type ProtocolOptions struct {
 	Role      ProtocolRole
 }
 
-type MessageHandlerFunc func(Message) error
+type MessageHandlerFunc func(Message, bool) error
 type MessageFromCborFunc func(uint, []byte) (Message, error)
 
 func New(config ProtocolConfig) *Protocol {
@@ -108,6 +108,7 @@ func (p *Protocol) SendError(err error) {
 
 func (p *Protocol) recvLoop() {
 	leftoverData := false
+	isResponse := false
 	for {
 		var err error
 		// Don't grab the next segment from the muxer if we still have data in the buffer
@@ -116,6 +117,8 @@ func (p *Protocol) recvLoop() {
 			segment := <-p.recvChan
 			// Add segment payload to buffer
 			p.recvBuffer.Write(segment.Payload)
+			// Save whether it's a response
+			isResponse = segment.IsResponse()
 		}
 		leftoverData = false
 		// Decode message into generic list until we can determine what type of message it is
@@ -140,7 +143,7 @@ func (p *Protocol) recvLoop() {
 			p.config.ErrorChan <- fmt.Errorf("%s: received unknown message type: %#v", p.config.Name, tmpMsg)
 		}
 		// Handle message
-		if err := p.handleMessage(msg); err != nil {
+		if err := p.handleMessage(msg, isResponse); err != nil {
 			p.config.ErrorChan <- err
 		}
 		if numBytesRead < p.recvBuffer.Len() {
@@ -183,7 +186,7 @@ func (p *Protocol) getNewState(msg Message) (State, error) {
 	return newState, nil
 }
 
-func (p *Protocol) handleMessage(msg Message) error {
+func (p *Protocol) handleMessage(msg Message, isResponse bool) error {
 	// Lock the state to prevent collisions
 	p.stateMutex.Lock()
 	if err := p.checkCurrentState(); err != nil {
@@ -197,5 +200,5 @@ func (p *Protocol) handleMessage(msg Message) error {
 	p.state = newState
 	p.stateMutex.Unlock()
 	// Call handler function
-	return p.config.MessageHandlerFunc(msg)
+	return p.config.MessageHandlerFunc(msg, isResponse)
 }