From 9e8633ded0caab24acc8a79eb34c85aa32fd6a2a Mon Sep 17 00:00:00 2001
From: Hussam <hussam@dominantstrategies.io>
Date: Fri, 27 Sep 2024 10:59:41 -0500
Subject: [PATCH] Verify send/receive by hash, fix emptyWo method

---
 core/types/block.go                      |  9 ++++++++-
 p2p/node/pubsubManager/gossipsub.go      |  6 +++---
 p2p/node/pubsubManager/gossipsub_test.go | 21 +++++++++++----------
 3 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/core/types/block.go b/core/types/block.go
index 0ce89888c..8230688f6 100644
--- a/core/types/block.go
+++ b/core/types/block.go
@@ -173,8 +173,8 @@ func EmptyWorkObject(nodeCtx int) *WorkObject {
 	wo.woHeader.SetDifficulty(big.NewInt(0))
 	wo.woHeader.SetPrimeTerminusNumber(big.NewInt(0))
 	wo.woHeader.SetTxHash(EmptyRootHash)
-	wo.woHeader.SetLocation(common.Location{})
 	wo.woHeader.SetNonce(EncodeNonce(0))
+	wo.woHeader.SetLocation(common.Location{})
 	wo.woHeader.SetTime(0)
 	wo.woBody.SetHeader(h)
 	wo.woBody.SetUncles([]*WorkObjectHeader{})
@@ -184,6 +184,13 @@ func EmptyWorkObject(nodeCtx int) *WorkObject {
 	return NewWorkObjectWithHeader(wo, NewEmptyTx(), nodeCtx, BlockObject)
 }
 
+func EmptyZoneWorkObject() *WorkObject {
+	emptyWo := EmptyWorkObject(common.ZONE_CTX)
+	emptyWo.woHeader.SetLocation(common.Location{0, 0})
+	emptyWo.woHeader.SetCoinbase(common.ZeroAddress(emptyWo.woHeader.location))
+	return emptyWo
+}
+
 // ProtoEncode serializes h into the Quai Proto Header format
 func (h *Header) ProtoEncode() (*ProtoHeader, error) {
 	if h == nil {
diff --git a/p2p/node/pubsubManager/gossipsub.go b/p2p/node/pubsubManager/gossipsub.go
index 731b56d74..def15a150 100644
--- a/p2p/node/pubsubManager/gossipsub.go
+++ b/p2p/node/pubsubManager/gossipsub.go
@@ -385,12 +385,12 @@ func (g *PubsubManager) Unsubscribe(location common.Location, datatype interface
 }
 
 // broadcasts data to subscribing peers
-func (g *PubsubManager) Broadcast(location common.Location, datatype interface{}) error {
-	topicName, err := NewTopic(g.genesis, location, datatype)
+func (g *PubsubManager) Broadcast(location common.Location, data interface{}) error {
+	topicName, err := NewTopic(g.genesis, location, data)
 	if err != nil {
 		return err
 	}
-	protoData, err := pb.ConvertAndMarshal(datatype)
+	protoData, err := pb.ConvertAndMarshal(data)
 	if err != nil {
 		return err
 	}
diff --git a/p2p/node/pubsubManager/gossipsub_test.go b/p2p/node/pubsubManager/gossipsub_test.go
index a5883f7c9..d117586dd 100644
--- a/p2p/node/pubsubManager/gossipsub_test.go
+++ b/p2p/node/pubsubManager/gossipsub_test.go
@@ -2,7 +2,6 @@ package pubsubManager
 
 import (
 	"context"
-	"math/big"
 	"sync"
 	"testing"
 
@@ -105,17 +104,20 @@ func TestPubsubManager(t *testing.T) {
 		ps.SetReceiveHandler(func(receivedFrom peer.ID, msgId string, msgTopic string, data interface{}, location common.Location) {
 			testCh <- data
 		})
-		// Success case
-		newWoHeader := types.NewWorkObjectHeader(common.Hash{1}, common.Hash{1}, big.NewInt(1), big.NewInt(1), big.NewInt(1),
-			common.Hash{1}, types.BlockNonce{}, 0, common.Location{}, common.ZeroAddress(common.Location{0, 0}))
-		newWo := types.NewWorkObjectWithHeaderAndTx(newWoHeader, types.NewEmptyTx())
-		broadcastedMessage := newWo.ConvertToHeaderView()
+
+		// Create and broadcast a new WorkObjectBlock
+		newWo := types.EmptyZoneWorkObject(common.ZONE_CTX)
+		broadcastedMessage := newWo.ConvertToBlockView()
 		err = ps.Broadcast(common.Location{0, 0}, broadcastedMessage)
 		require.NoError(t, err, "Failed to broadcast message")
 
-		// Verify if subscription received message
+		// Verify if subscription received correct message type
 		receivedMessage := <-testCh
-		require.Equal(t, broadcastedMessage, receivedMessage, "Received message is not the same as broadcasted message")
+		recvdWorkObject, ok := receivedMessage.(types.WorkObjectBlockView)
+		require.True(t, ok, "Unable to cast workobject")
+
+		// Verify equality of the send and receive
+		require.Equal(t, broadcastedMessage.Hash(), recvdWorkObject.Hash())
 	})
 
 	t.Run("Unsubscribe", func(t *testing.T) {
@@ -151,8 +153,7 @@ func TestMultipleRequests(t *testing.T) {
 	quaiBackend, _ := quai.NewQuaiBackend()
 	ps.SetQuaiBackend(quaiBackend)
 
-	wo := types.EmptyWorkObject(common.ZONE_CTX)
-	wo.WorkObjectHeader().SetCoinbase(common.BytesToAddress([]byte{1}, common.Location{0, 0}))
+	wo := types.EmptyZoneWorkObject(common.ZONE_CTX)
 
 	tx := types.NewEmptyTx()
 	txs := types.Transactions{tx}