From 144d124195f17421849b8a085be43b8589a4ae87 Mon Sep 17 00:00:00 2001
From: Hussam <hussam@dominantstrategies.io>
Date: Tue, 14 May 2024 11:51:32 -0500
Subject: [PATCH] Workobject for requests

---
 p2p/node/api.go          | 4 ++--
 p2p/node/p2p_services.go | 5 +++++
 quai/handler.go          | 4 ++--
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/p2p/node/api.go b/p2p/node/api.go
index f36d73f041..39a1572e40 100644
--- a/p2p/node/api.go
+++ b/p2p/node/api.go
@@ -116,7 +116,7 @@ func (p *P2PNode) Stop() error {
 	}
 }
 
-func (p *P2PNode) requestFromPeers(topic *pubsubManager.Topic, location common.Location, requestData interface{}, responseDataType interface{}, resultChan chan interface{}) {
+func (p *P2PNode) requestFromPeers(topic *pubsubManager.Topic, requestData interface{}, resultChan chan interface{}) {
 	go func() {
 		defer func() {
 			if r := recover(); r != nil {
@@ -215,7 +215,7 @@ func (p *P2PNode) Request(location common.Location, requestData interface{}, res
 		}
 	}
 
-	p.requestFromPeers(topic, location, requestData, responseDataType, resultChan)
+	p.requestFromPeers(topic, requestData, resultChan)
 	// TODO: optimize with waitgroups or a doneChan to only query if no peers responded
 	// Right now this creates too many streams, so don't call this until we have a better solution
 	// p.queryDHT(location, requestData, responseDataType, resultChan)
diff --git a/p2p/node/p2p_services.go b/p2p/node/p2p_services.go
index 6482394fe1..f182a284dd 100644
--- a/p2p/node/p2p_services.go
+++ b/p2p/node/p2p_services.go
@@ -86,6 +86,11 @@ func (p *P2PNode) requestFromPeer(peerID peer.ID, topic *pubsubManager.Topic, re
 	case *types.WorkObject:
 		if block, ok := recvdType.(*types.WorkObject); ok {
 			switch data := reqData.(type) {
+			case *types.WorkObjectBlockView, *types.WorkObjectHeaderView:
+				if block.Hash() == data.(*types.WorkObject).Hash() {
+					return data, nil
+				}
+				return nil, errors.Errorf("invalid response: expected block with hash %s, got %s", data.(*types.WorkObject).Hash(), block.Hash())
 			case common.Hash:
 				if block.Hash() == data {
 					return block, nil
diff --git a/quai/handler.go b/quai/handler.go
index abce4576f6..ef2e739d1c 100644
--- a/quai/handler.go
+++ b/quai/handler.go
@@ -238,10 +238,10 @@ func (h *handler) GetNextPrimeBlock(number *big.Int) {
 				// If the blockHash for the asked number is not present in the
 				// appended database we ask the peer for the block with this hash
 				if block == nil {
-					resultCh := h.p2pBackend.Request(h.nodeLocation, blockHash, &types.WorkObject{})
+					resultCh := h.p2pBackend.Request(h.nodeLocation, blockHash, &types.WorkObjectBlockView{})
 					block := <-resultCh
 					if block != nil {
-						h.core.WriteBlock(block.(*types.WorkObject))
+						h.core.WriteBlock(block.(*types.WorkObjectBlockView).WorkObject)
 					}
 				}
 			}