From 990d11b1d70126bb545e834724e51a5f8e46e64a Mon Sep 17 00:00:00 2001
From: Maddiaa <47148561+Maddiaa0@users.noreply.github.com>
Date: Fri, 6 Dec 2024 01:37:04 +0800
Subject: [PATCH] fix(p2p): override msg Id  (#10415)

---
 .../circuit-types/src/p2p/block_proposal.ts   |  4 +-
 .../circuit-types/src/p2p/interface.ts        | 12 +++++
 yarn-project/p2p/package.json                 |  3 +-
 yarn-project/p2p/src/service/encoding.ts      | 48 +++++++++++++++++++
 .../p2p/src/service/libp2p_service.ts         |  4 ++
 yarn-project/yarn.lock                        |  8 ++++
 6 files changed, 76 insertions(+), 3 deletions(-)
 create mode 100644 yarn-project/p2p/src/service/encoding.ts

diff --git a/yarn-project/circuit-types/src/p2p/block_proposal.ts b/yarn-project/circuit-types/src/p2p/block_proposal.ts
index 8e64f4c1fd9..5d2ad7f7f46 100644
--- a/yarn-project/circuit-types/src/p2p/block_proposal.ts
+++ b/yarn-project/circuit-types/src/p2p/block_proposal.ts
@@ -1,5 +1,5 @@
 import { Buffer32 } from '@aztec/foundation/buffer';
-import { recoverAddress } from '@aztec/foundation/crypto';
+import { keccak256, recoverAddress } from '@aztec/foundation/crypto';
 import { type EthAddress } from '@aztec/foundation/eth-address';
 import { Signature } from '@aztec/foundation/eth-signature';
 import { type Fr } from '@aztec/foundation/fields';
@@ -42,7 +42,7 @@ export class BlockProposal extends Gossipable {
   }
 
   override p2pMessageIdentifier(): Buffer32 {
-    return BlockProposalHash.fromField(this.payload.archive);
+    return new BlockProposalHash(keccak256(this.signature.toBuffer()));
   }
 
   get archive(): Fr {
diff --git a/yarn-project/circuit-types/src/p2p/interface.ts b/yarn-project/circuit-types/src/p2p/interface.ts
index 1e252250399..06a02794602 100644
--- a/yarn-project/circuit-types/src/p2p/interface.ts
+++ b/yarn-project/circuit-types/src/p2p/interface.ts
@@ -17,3 +17,15 @@ export const TopicTypeMap: Record<string, typeof Gossipable> = {
   [TopicType.block_attestation]: BlockAttestation as unknown as typeof Gossipable,
   [TopicType.epoch_proof_quote]: EpochProofQuote as unknown as typeof Gossipable,
 };
+
+/**
+ * Map from topic to deserialiser
+ *
+ * Used in msgIdFn libp2p to get the p2pMessageIdentifier from a message
+ */
+export const TopicToDeserializer = {
+  [Tx.p2pTopic]: Tx.fromBuffer,
+  [BlockProposal.p2pTopic]: BlockProposal.fromBuffer,
+  [BlockAttestation.p2pTopic]: BlockAttestation.fromBuffer,
+  [EpochProofQuote.p2pTopic]: EpochProofQuote.fromBuffer,
+};
diff --git a/yarn-project/p2p/package.json b/yarn-project/p2p/package.json
index a458dcc94a5..bc72fd926e0 100644
--- a/yarn-project/p2p/package.json
+++ b/yarn-project/p2p/package.json
@@ -91,7 +91,8 @@
     "libp2p": "1.5.0",
     "semver": "^7.6.0",
     "sha3": "^2.1.4",
-    "tslib": "^2.4.0"
+    "tslib": "^2.4.0",
+    "xxhash-wasm": "^1.1.0"
   },
   "devDependencies": {
     "@aztec/archiver": "workspace:^",
diff --git a/yarn-project/p2p/src/service/encoding.ts b/yarn-project/p2p/src/service/encoding.ts
new file mode 100644
index 00000000000..34a0c49a1db
--- /dev/null
+++ b/yarn-project/p2p/src/service/encoding.ts
@@ -0,0 +1,48 @@
+// Taken from lodestar: https://github.com/ChainSafe/lodestar
+import { sha256 } from '@aztec/foundation/crypto';
+
+import { type RPC } from '@chainsafe/libp2p-gossipsub/message';
+import { type Message } from '@libp2p/interface';
+import xxhashFactory from 'xxhash-wasm';
+
+// Load WASM
+const xxhash = await xxhashFactory();
+
+// Use salt to prevent msgId from being mined for collisions
+const h64Seed = BigInt(Math.floor(Math.random() * 1e9));
+
+// Shared buffer to convert msgId to string
+const sharedMsgIdBuf = Buffer.alloc(20);
+
+/**
+ * The function used to generate a gossipsub message id
+ * We use the first 8 bytes of SHA256(data) for content addressing
+ */
+export function fastMsgIdFn(rpcMsg: RPC.Message): string {
+  if (rpcMsg.data) {
+    return xxhash.h64Raw(rpcMsg.data, h64Seed).toString(16);
+  }
+  return '0000000000000000';
+}
+
+export function msgIdToStrFn(msgId: Uint8Array): string {
+  // This happens serially, no need to reallocate the buffer
+  sharedMsgIdBuf.set(msgId);
+  return `0x${sharedMsgIdBuf.toString('hex')}`;
+}
+
+/**
+ * Get the message identifier from a libp2p message
+ *
+ * Follows similarly to:
+ * https://github.com/ethereum/consensus-specs/blob/v1.1.0-alpha.7/specs/altair/p2p-interface.md#topics-and-messages
+ *
+ * @param message - The libp2p message
+ * @returns The message identifier
+ */
+export function getMsgIdFn(message: Message) {
+  const { topic } = message;
+
+  const vec = [Buffer.from(topic), message.data];
+  return sha256(Buffer.concat(vec)).subarray(0, 20);
+}
diff --git a/yarn-project/p2p/src/service/libp2p_service.ts b/yarn-project/p2p/src/service/libp2p_service.ts
index 18d2d180a4a..28251f9110c 100644
--- a/yarn-project/p2p/src/service/libp2p_service.ts
+++ b/yarn-project/p2p/src/service/libp2p_service.ts
@@ -43,6 +43,7 @@ import {
 } from '../tx_validator/index.js';
 import { type PubSubLibp2p, convertToMultiaddr } from '../util.js';
 import { AztecDatastore } from './data_store.js';
+import { fastMsgIdFn, getMsgIdFn, msgIdToStrFn } from './encoding.js';
 import { PeerManager } from './peer_manager.js';
 import { PeerErrorSeverity } from './peer_scoring.js';
 import { pingHandler, statusHandler } from './reqresp/handlers.js';
@@ -242,6 +243,9 @@ export class LibP2PService extends WithTracer implements P2PService {
           heartbeatInterval: config.gossipsubInterval,
           mcacheLength: config.gossipsubMcacheLength,
           mcacheGossip: config.gossipsubMcacheGossip,
+          msgIdFn: getMsgIdFn,
+          msgIdToStrFn: msgIdToStrFn,
+          fastMsgIdFn: fastMsgIdFn,
           metricsRegister: otelMetricsAdapter,
           metricsTopicStrToLabel: metricsTopicStrToLabels(),
           scoreParams: createPeerScoreParams({
diff --git a/yarn-project/yarn.lock b/yarn-project/yarn.lock
index c53310ab30d..732f75683f7 100644
--- a/yarn-project/yarn.lock
+++ b/yarn-project/yarn.lock
@@ -924,6 +924,7 @@ __metadata:
     typescript: ^5.0.4
     uint8arrays: ^5.0.3
     viem: ^2.7.15
+    xxhash-wasm: ^1.1.0
   languageName: unknown
   linkType: soft
 
@@ -18628,6 +18629,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"xxhash-wasm@npm:^1.1.0":
+  version: 1.1.0
+  resolution: "xxhash-wasm@npm:1.1.0"
+  checksum: 2ccecb3b1dac5fefe11002d5ff5d106bbb5b506f9ee817ecf1bda65e132ebff3c82701c6727df3cb90b94a6dc1d8b294337678606f2304bcb0fd6b8dc68afe0d
+  languageName: node
+  linkType: hard
+
 "y18n@npm:^5.0.5":
   version: 5.0.8
   resolution: "y18n@npm:5.0.8"