From 96847374d6c61f3372a16185d9fff93e582505bb Mon Sep 17 00:00:00 2001
From: "fryorcraken.eth" <git@fryorcraken.xyz>
Date: Fri, 10 Mar 2023 12:43:29 +1100
Subject: [PATCH] feat(relay): validate waku message at gossip layer

---
 packages/core/src/lib/relay/index.ts          |  7 +++
 .../src/lib/relay/message_validator.spec.ts   | 59 +++++++++++++++++++
 .../core/src/lib/relay/message_validator.ts   | 35 +++++++++++
 3 files changed, 101 insertions(+)
 create mode 100644 packages/core/src/lib/relay/message_validator.spec.ts
 create mode 100644 packages/core/src/lib/relay/message_validator.ts

diff --git a/packages/core/src/lib/relay/index.ts b/packages/core/src/lib/relay/index.ts
index 6f2355b326..5878509327 100644
--- a/packages/core/src/lib/relay/index.ts
+++ b/packages/core/src/lib/relay/index.ts
@@ -23,6 +23,7 @@ import { TopicOnlyDecoder } from "../message/topic_only_message.js";
 import { pushOrInitMapSet } from "../push_or_init_map.js";
 
 import * as constants from "./constants.js";
+import { messageValidator } from "./message_validator.js";
 
 const log = debug("waku:relay");
 
@@ -170,9 +171,15 @@ class Relay extends GossipSub implements IRelay {
       }
     );
 
+    this.topicValidators.set(pubSubTopic, messageValidator);
     super.subscribe(pubSubTopic);
   }
 
+  unsubscribe(pubSubTopic: TopicStr): void {
+    super.unsubscribe(pubSubTopic);
+    this.topicValidators.delete(pubSubTopic);
+  }
+
   getMeshPeers(topic?: TopicStr): PeerIdStr[] {
     return super.getMeshPeers(topic ?? this.pubSubTopic);
   }
diff --git a/packages/core/src/lib/relay/message_validator.spec.ts b/packages/core/src/lib/relay/message_validator.spec.ts
new file mode 100644
index 0000000000..0f01c1e10b
--- /dev/null
+++ b/packages/core/src/lib/relay/message_validator.spec.ts
@@ -0,0 +1,59 @@
+import { TopicValidatorResult } from "@libp2p/interface-pubsub";
+import type { UnsignedMessage } from "@libp2p/interface-pubsub";
+import { createSecp256k1PeerId } from "@libp2p/peer-id-factory";
+import { expect } from "chai";
+import fc from "fast-check";
+
+import { createEncoder } from "../message/version_0.js";
+
+import { messageValidator } from "./message_validator.js";
+
+describe("Message Validator", () => {
+  it("Accepts a valid Waku Message", async () => {
+    await fc.assert(
+      fc.asyncProperty(
+        fc.uint8Array({ minLength: 1 }),
+        fc.string({ minLength: 1 }),
+        fc.string({ minLength: 1 }),
+        async (payload, pubSubTopic, contentTopic) => {
+          const peerId = await createSecp256k1PeerId();
+
+          const encoder = createEncoder({ contentTopic });
+          const bytes = await encoder.toWire({ payload });
+
+          const message: UnsignedMessage = {
+            type: "unsigned",
+            topic: pubSubTopic,
+            data: bytes,
+          };
+
+          const result = messageValidator(peerId, message);
+
+          expect(result).to.eq(TopicValidatorResult.Accept);
+        }
+      )
+    );
+  });
+
+  it("Rejects garbage", async () => {
+    await fc.assert(
+      fc.asyncProperty(
+        fc.uint8Array(),
+        fc.string(),
+        async (data, pubSubTopic) => {
+          const peerId = await createSecp256k1PeerId();
+
+          const message: UnsignedMessage = {
+            type: "unsigned",
+            topic: pubSubTopic,
+            data,
+          };
+
+          const result = messageValidator(peerId, message);
+
+          expect(result).to.eq(TopicValidatorResult.Reject);
+        }
+      )
+    );
+  });
+});
diff --git a/packages/core/src/lib/relay/message_validator.ts b/packages/core/src/lib/relay/message_validator.ts
new file mode 100644
index 0000000000..6a07d6165d
--- /dev/null
+++ b/packages/core/src/lib/relay/message_validator.ts
@@ -0,0 +1,35 @@
+import type { PeerId } from "@libp2p/interface-peer-id";
+import type { Message } from "@libp2p/interface-pubsub";
+import { TopicValidatorResult } from "@libp2p/interface-pubsub";
+import { proto_message as proto } from "@waku/proto";
+import debug from "debug";
+
+const log = debug("waku:relay");
+
+export function messageValidator(
+  peer: PeerId,
+  message: Message
+): TopicValidatorResult {
+  const startTime = performance.now();
+  log(`validating message from ${peer} received on ${message.topic}`);
+  let result = TopicValidatorResult.Accept;
+
+  try {
+    const protoMessage = proto.WakuMessage.decode(message.data);
+
+    if (
+      !protoMessage.contentTopic ||
+      !protoMessage.contentTopic.length ||
+      !protoMessage.payload ||
+      !protoMessage.payload.length
+    ) {
+      result = TopicValidatorResult.Reject;
+    }
+  } catch (e) {
+    result = TopicValidatorResult.Reject;
+  }
+
+  const endTime = performance.now();
+  log(`Validation time (must be <100ms): ${endTime - startTime}ms`);
+  return result;
+}