From 41cdb92be70fa504e87f758d0a4289ac4b32a4c4 Mon Sep 17 00:00:00 2001
From: Akhil <111925100+shash256@users.noreply.github.com>
Date: Wed, 15 May 2024 14:13:13 +0400
Subject: [PATCH] feat: Added message size check before relay for lightpush
 (#2695)

---
 tests/node/test_wakunode_lightpush.nim | 21 ++++++++++++++++++++-
 waku/node/waku_node.nim                |  4 ++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/tests/node/test_wakunode_lightpush.nim b/tests/node/test_wakunode_lightpush.nim
index d783792f9..92ef250eb 100644
--- a/tests/node/test_wakunode_lightpush.nim
+++ b/tests/node/test_wakunode_lightpush.nim
@@ -6,6 +6,7 @@ import
   testutils/unittests,
   chronos,
   chronicles,
+  std/strformat,
   os,
   libp2p/[peerstore, crypto/crypto]
 
@@ -23,7 +24,8 @@ import
     waku_lightpush/protocol_metrics,
     waku_lightpush/rpc,
   ],
-  ../testlib/[assertions, common, wakucore, wakunode, testasync, futures, testutils]
+  ../testlib/[assertions, common, wakucore, wakunode, testasync, futures, testutils],
+  ../resources/payloads
 
 suite "Waku Lightpush - End To End":
   var
@@ -85,3 +87,20 @@ suite "Waku Lightpush - End To End":
 
       # Then the message is relayed to the server
       assertResultOk publishResponse
+
+  suite "Waku LightPush Validation Tests":
+    asyncTest "Validate message size exceeds limit":
+      let 
+        msgOverLimit = fakeWakuMessage(
+          contentTopic = contentTopic,
+          payload = getByteSequence(DefaultMaxWakuMessageSize + 64 * 1024),
+        )
+      
+      # When the client publishes an over-limit message
+      let publishResponse = await client.lightpushPublish(
+        some(pubsubTopic), msgOverLimit, serverRemotePeerInfo
+      )
+
+      check:
+        publishResponse.isErr()
+        publishResponse.error == fmt"Message size exceeded maximum of {DefaultMaxWakuMessageSize} bytes"
\ No newline at end of file
diff --git a/waku/node/waku_node.nim b/waku/node/waku_node.nim
index 4773846d3..f43573169 100644
--- a/waku/node/waku_node.nim
+++ b/waku/node/waku_node.nim
@@ -917,6 +917,10 @@ proc mountLightPush*(
     pushHandler = proc(
         peer: PeerId, pubsubTopic: string, message: WakuMessage
     ): Future[WakuLightPushResult[void]] {.async.} =
+      let validationRes = await node.wakuRelay.validateMessage(pubSubTopic, message)
+      if validationRes.isErr():
+        return err(validationRes.error)
+
       let publishedCount =
         await node.wakuRelay.publish(pubsubTopic, message.encode().buffer)