From a81214219dc57caa6a812ef5b656e910c1df37e3 Mon Sep 17 00:00:00 2001
From: Nishant Das <nishdas93@gmail.com>
Date: Mon, 7 Sep 2020 02:05:16 +0800
Subject: [PATCH] Hash Pubsub Messages Correctly (#7177)

* check message id

* gaz

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>
---
 beacon-chain/p2p/BUILD.bazel    |  2 ++
 beacon-chain/p2p/pubsub.go      |  6 ++++--
 beacon-chain/p2p/pubsub_test.go | 11 +++++++++++
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/beacon-chain/p2p/BUILD.bazel b/beacon-chain/p2p/BUILD.bazel
index 02946d9dd6a4..c28be5cd9d56 100644
--- a/beacon-chain/p2p/BUILD.bazel
+++ b/beacon-chain/p2p/BUILD.bazel
@@ -120,6 +120,7 @@ go_test(
         "//proto/testing:go_default_library",
         "//shared/bytesutil:go_default_library",
         "//shared/featureconfig:go_default_library",
+        "//shared/hashutil:go_default_library",
         "//shared/iputils:go_default_library",
         "//shared/p2putils:go_default_library",
         "//shared/params:go_default_library",
@@ -138,6 +139,7 @@ go_test(
         "@com_github_libp2p_go_libp2p_core//network:go_default_library",
         "@com_github_libp2p_go_libp2p_core//peer:go_default_library",
         "@com_github_libp2p_go_libp2p_pubsub//:go_default_library",
+        "@com_github_libp2p_go_libp2p_pubsub//pb:go_default_library",
         "@com_github_libp2p_go_libp2p_swarm//testing:go_default_library",
         "@com_github_multiformats_go_multiaddr//:go_default_library",
         "@com_github_prysmaticlabs_ethereumapis//eth/v1alpha1:go_default_library",
diff --git a/beacon-chain/p2p/pubsub.go b/beacon-chain/p2p/pubsub.go
index b71b6ed5fd37..75f6795a97f9 100644
--- a/beacon-chain/p2p/pubsub.go
+++ b/beacon-chain/p2p/pubsub.go
@@ -81,9 +81,11 @@ func (s *Service) SubscribeToTopic(topic string, opts ...pubsub.SubOpt) (*pubsub
 //
 // ETH2 spec defines the message ID as:
 //    message-id: base64(SHA256(message.data))
+// where base64 is the URL-safe base64 alphabet with
+// padding characters omitted.
 func msgIDFunction(pmsg *pubsub_pb.Message) string {
-	h := hashutil.FastSum256(pmsg.Data)
-	return base64.URLEncoding.EncodeToString(h[:])
+	h := hashutil.Hash(pmsg.Data)
+	return base64.RawURLEncoding.EncodeToString(h[:])
 }
 
 func setPubSubParameters() {
diff --git a/beacon-chain/p2p/pubsub_test.go b/beacon-chain/p2p/pubsub_test.go
index dee4bb781e0d..d55e73175ea4 100644
--- a/beacon-chain/p2p/pubsub_test.go
+++ b/beacon-chain/p2p/pubsub_test.go
@@ -2,10 +2,13 @@ package p2p
 
 import (
 	"context"
+	"encoding/base64"
 	"fmt"
 	"sync"
 	"testing"
 
+	pubsubpb "github.com/libp2p/go-libp2p-pubsub/pb"
+	"github.com/prysmaticlabs/prysm/shared/hashutil"
 	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
 	"github.com/prysmaticlabs/prysm/shared/testutil/require"
 )
@@ -26,3 +29,11 @@ func TestService_PublishToTopicConcurrentMapWrite(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+func TestMessageIDFunction_HashesCorrectly(t *testing.T) {
+	msg := [32]byte{'J', 'U', 'N', 'K'}
+	pMsg := &pubsubpb.Message{Data: msg[:]}
+	hashedData := hashutil.Hash(pMsg.Data)
+	msgID := base64.RawURLEncoding.EncodeToString(hashedData[:])
+	assert.Equal(t, msgID, msgIDFunction(pMsg), "Got incorrect msg id")
+}