From b7dc3d7576e9444a5acbb036812c05cfccb25815 Mon Sep 17 00:00:00 2001
From: fryorcraken <110212804+fryorcraken@users.noreply.github.com>
Date: Thu, 26 Oct 2023 22:14:47 +1100
Subject: [PATCH] fix: measure total message size (#1643)

Network message limitations are imposed on the whole message, not just
the payload.

Co-authored-by: Sasha <118575614+weboko@users.noreply.github.com>
---
 packages/core/src/lib/light_push/index.ts     |  7 +++--
 packages/relay/src/index.ts                   | 18 ++++++-------
 .../tests/tests/light-push/index.node.spec.ts |  8 +-----
 .../tests/tests/relay/publish.node.spec.ts    | 11 --------
 packages/utils/src/common/is_size_valid.ts    | 27 ++++++++++++++-----
 5 files changed, 35 insertions(+), 36 deletions(-)

diff --git a/packages/core/src/lib/light_push/index.ts b/packages/core/src/lib/light_push/index.ts
index d0532013b8..bea69c6538 100644
--- a/packages/core/src/lib/light_push/index.ts
+++ b/packages/core/src/lib/light_push/index.ts
@@ -11,7 +11,10 @@ import {
   SendResult
 } from "@waku/interfaces";
 import { PushResponse } from "@waku/proto";
-import { ensurePubsubTopicIsConfigured, isSizeUnderCap } from "@waku/utils";
+import {
+  ensurePubsubTopicIsConfigured,
+  isMessageSizeUnderCap
+} from "@waku/utils";
 import { Logger } from "@waku/utils";
 import all from "it-all";
 import * as lp from "it-length-prefixed";
@@ -61,7 +64,7 @@ class LightPush extends BaseProtocol implements ILightPush {
         return { query: null, error: SendError.EMPTY_PAYLOAD };
       }
 
-      if (!isSizeUnderCap(message.payload)) {
+      if (!(await isMessageSizeUnderCap(encoder, message))) {
         log.error("Failed to send waku light push: message is bigger than 1MB");
         return { query: null, error: SendError.SIZE_TOO_BIG };
       }
diff --git a/packages/relay/src/index.ts b/packages/relay/src/index.ts
index ad0a7e19da..f9a2f35e92 100644
--- a/packages/relay/src/index.ts
+++ b/packages/relay/src/index.ts
@@ -25,7 +25,7 @@ import {
   SendError,
   SendResult
 } from "@waku/interfaces";
-import { isSizeUnderCap, toAsyncIterator } from "@waku/utils";
+import { isWireSizeUnderCap, toAsyncIterator } from "@waku/utils";
 import { pushOrInitMapSet } from "@waku/utils";
 import { Logger } from "@waku/utils";
 
@@ -112,14 +112,6 @@ class Relay implements IRelay {
       };
     }
 
-    if (!isSizeUnderCap(message.payload)) {
-      log.error("Failed to send waku relay: message is bigger that 1MB");
-      return {
-        recipients,
-        errors: [SendError.SIZE_TOO_BIG]
-      };
-    }
-
     const msg = await encoder.toWire(message);
     if (!msg) {
       log.error("Failed to encode message, aborting publish");
@@ -129,6 +121,14 @@ class Relay implements IRelay {
       };
     }
 
+    if (!isWireSizeUnderCap(msg)) {
+      log.error("Failed to send waku relay: message is bigger that 1MB");
+      return {
+        recipients,
+        errors: [SendError.SIZE_TOO_BIG]
+      };
+    }
+
     return this.gossipSub.publish(pubsubTopic, msg);
   }
 
diff --git a/packages/tests/tests/light-push/index.node.spec.ts b/packages/tests/tests/light-push/index.node.spec.ts
index 8c68462c62..2ea99149ab 100644
--- a/packages/tests/tests/light-push/index.node.spec.ts
+++ b/packages/tests/tests/light-push/index.node.spec.ts
@@ -204,14 +204,8 @@ describe("Waku Light Push", function () {
   });
 
   it("Push message equal or less that 1MB", async function () {
-    const oneMbPayload = generateRandomUint8Array(1024 ** 2);
-    let pushResponse = await waku.lightPush.send(TestEncoder, {
-      payload: oneMbPayload
-    });
-    expect(pushResponse.recipients.length).to.greaterThan(0);
-
     const bigPayload = generateRandomUint8Array(65536);
-    pushResponse = await waku.lightPush.send(TestEncoder, {
+    const pushResponse = await waku.lightPush.send(TestEncoder, {
       payload: bigPayload
     });
     expect(pushResponse.recipients.length).to.greaterThan(0);
diff --git a/packages/tests/tests/relay/publish.node.spec.ts b/packages/tests/tests/relay/publish.node.spec.ts
index f650fe19b6..02c6ef0ec0 100644
--- a/packages/tests/tests/relay/publish.node.spec.ts
+++ b/packages/tests/tests/relay/publish.node.spec.ts
@@ -141,17 +141,6 @@ describe("Waku Relay, Publish", function () {
     expect(await messageCollector.waitForMessages(1)).to.eq(false);
   });
 
-  it("Publish message with size of 1 MB", async function () {
-    const pushResponse = await waku1.relay.send(TestEncoder, {
-      payload: generateRandomUint8Array(1024 ** 2)
-    });
-    expect(pushResponse.recipients.length).to.eq(1);
-    expect(pushResponse.recipients[0].toString()).to.eq(
-      waku2.libp2p.peerId.toString()
-    );
-    expect(await messageCollector.waitForMessages(1)).to.eq(true);
-  });
-
   [1024 ** 2 + 65536, 2 * 1024 ** 2].forEach((testItem) => {
     it("Fails to publish message with size larger than 1 MB", async function () {
       const pushResponse = await waku1.relay.send(TestEncoder, {
diff --git a/packages/utils/src/common/is_size_valid.ts b/packages/utils/src/common/is_size_valid.ts
index 7c565b321b..011e5c129c 100644
--- a/packages/utils/src/common/is_size_valid.ts
+++ b/packages/utils/src/common/is_size_valid.ts
@@ -1,10 +1,23 @@
+import type { IEncoder, IMessage } from "@waku/interfaces";
+
 const MB = 1024 ** 2;
-const SIZE_CAP = 1; // 1 MB
+const SIZE_CAP_IN_MB = 1;
 
-export const isSizeUnderCap = (payload: Uint8Array): boolean => {
-  if (payload.length / MB > SIZE_CAP) {
-    return false;
-  }
+/**
+ * Return whether the size of the message is under the upper limit for the network.
+ * This performs a protobuf encoding! If you have access to the fully encoded message,
+ * use {@link isSizeUnderCapBuf} instead.
+ * @param message
+ * @param encoder
+ */
+export async function isMessageSizeUnderCap(
+  encoder: IEncoder,
+  message: IMessage
+): Promise<boolean> {
+  const buf = await encoder.toWire(message);
+  if (!buf) return false;
+  return isWireSizeUnderCap(buf);
+}
 
-  return true;
-};
+export const isWireSizeUnderCap = (buf: Uint8Array): boolean =>
+  buf.length / MB <= SIZE_CAP_IN_MB;