From f2ff1e315e75c28016a48c8359760f0bd5ad92cb Mon Sep 17 00:00:00 2001
From: Diego <diego@status.im>
Date: Tue, 4 Jul 2023 15:26:10 +0200
Subject: [PATCH] fix metrics and add tests

---
 libp2p/protocols/pubsub/gossipsub/scoring.nim | 13 +++---
 libp2p/protocols/pubsub/rpc/messages.nim      |  6 ++-
 tests/pubsub/testmessage.nim                  | 44 +++++++++++++++++++
 3 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/libp2p/protocols/pubsub/gossipsub/scoring.nim b/libp2p/protocols/pubsub/gossipsub/scoring.nim
index 2edf0d2bbe..7d52b124ec 100644
--- a/libp2p/protocols/pubsub/gossipsub/scoring.nim
+++ b/libp2p/protocols/pubsub/gossipsub/scoring.nim
@@ -28,8 +28,9 @@ declareGauge(libp2p_gossipsub_peers_score_invalidMessageDeliveries, "Detailed go
 declareGauge(libp2p_gossipsub_peers_score_appScore, "Detailed gossipsub scoring metric", labels = ["agent"])
 declareGauge(libp2p_gossipsub_peers_score_behaviourPenalty, "Detailed gossipsub scoring metric", labels = ["agent"])
 declareGauge(libp2p_gossipsub_peers_score_colocationFactor, "Detailed gossipsub scoring metric", labels = ["agent"])
-declareGauge(libp2p_gossipsub_peers_score_averageInvalidIgnoredTrafficRatio, "Average Invalid Ignored Traffic Ratio", labels = ["agent"])
-declareGauge(libp2p_gossipsub_peers_score_averageInvalidTrafficRatio, "Average Invalid Traffic Ratio", labels = ["agent"])
+declareGauge(libp2p_gossipsub_peers_score_invalidIgnoredTrafficMB, "Invalid Ignored Traffic (MB)", labels = ["agent"])
+declareGauge(libp2p_gossipsub_peers_score_invalidTrafficMB, "Invalid Traffic (MB)", labels = ["agent"])
+declareGauge(libp2p_gossipsub_peers_score_totalTrafficMB, "Total Traffic (MB)", labels = ["agent"])
 
 proc init*(_: type[TopicParams]): TopicParams =
   TopicParams(
@@ -132,10 +133,10 @@ proc disconnectIfBadTrafficPeer(g: GossipSub, peer: PubSubPeer) =
     let invalidIgnoredTrafficRatio = float64(peer.invalidIgnoredTraffic) / float64(peer.totalTraffic)
     let totalInvalidTrafficRatio = invalidTrafficRatio + invalidIgnoredTrafficRatio
     let numberOfPeersForAgent = float64(numberOfPeersForAgent(g, agent))
-    libp2p_gossipsub_peers_score_averageInvalidTrafficRatio
-      .inc(invalidTrafficRatio / (if numberOfPeersForAgent != 0: numberOfPeersForAgent else: 1), labelValues = [agent])
-    libp2p_gossipsub_peers_score_averageInvalidIgnoredTrafficRatio
-      .inc(invalidIgnoredTrafficRatio / (if numberOfPeersForAgent != 0: numberOfPeersForAgent else: 1), labelValues = [agent])
+    libp2p_gossipsub_peers_score_invalidTrafficMB.inc(float64(peer.invalidTraffic) / 1_000_000, labelValues = [agent])
+    libp2p_gossipsub_peers_score_invalidIgnoredTrafficMB.inc(float64(peer.invalidIgnoredTraffic) / 1_000_000, labelValues = [agent])
+    libp2p_gossipsub_peers_score_totalTrafficMB.inc(float64(peer.totalTraffic) / 1_000_000, labelValues = [agent])
+
     discard g.disconnectIfBadPeer(peer, -totalInvalidTrafficRatio, -0.30'f64) #g.parameters.maxInvalidTrafficRatio)
 
 proc updateScores*(g: GossipSub) = # avoid async
diff --git a/libp2p/protocols/pubsub/rpc/messages.nim b/libp2p/protocols/pubsub/rpc/messages.nim
index c91caddb71..6693716db7 100644
--- a/libp2p/protocols/pubsub/rpc/messages.nim
+++ b/libp2p/protocols/pubsub/rpc/messages.nim
@@ -141,10 +141,12 @@ proc byteSize*(msg: RPCMsg): int =
     let ctrl = msg.control.get()
     for item in ctrl.ihave:
       total += item.topicId.len
-      total += item.messageIds.len * sizeof(byte)  # Assuming MessageId is seq[byte]
+      for msgId in item.messageIds:
+        total += msgId.len
 
     for item in ctrl.iwant:
-      total += item.messageIds.len * sizeof(byte)  # Assuming MessageId is seq[byte]
+      for msgId in item.messageIds:
+        total += msgId.len
 
     for item in ctrl.graft:
       total += item.topicId.len
diff --git a/tests/pubsub/testmessage.nim b/tests/pubsub/testmessage.nim
index 7bc4b267a3..99310e7fbe 100644
--- a/tests/pubsub/testmessage.nim
+++ b/tests/pubsub/testmessage.nim
@@ -73,3 +73,47 @@ suite "Message":
     check:
       msgIdResult.isErr
       msgIdResult.error == ValidationResult.Reject
+
+  test "byteSize for Message":
+    var msg = Message(
+      fromPeer: PeerId(data: @[]),  # Empty seq[byte]
+      data: @[1'u8, 2, 3],  # 3 bytes
+      seqno: @[1'u8],  # 1 byte
+      signature: @[],  # Empty seq[byte]
+      key: @[1'u8],  # 1 byte
+      topicIds: @["abc", "defgh"]  # 3 + 5 = 8 bytes
+    )
+
+    check byteSize(msg) == 3 + 1 + 1 + 8  # Total: 13 bytes
+
+  test "byteSize for RPCMsg":
+    var msg = RPCMsg(
+      subscriptions: @[
+        SubOpts(topic: "abc", subscribe: true),
+        SubOpts(topic: "def", subscribe: false)
+      ],  # 3 + 3 + 2 * sizeof(bool) bytes
+      messages: @[
+        Message(fromPeer: PeerId(data: @[]), data: @[1'u8, 2, 3], seqno: @[1'u8], signature: @[], key: @[1'u8], topicIds: @["abc", "defgh"]),
+        Message(fromPeer: PeerId(data: @[]), data: @[], seqno: @[], signature: @[], key: @[], topicIds: @["abc"])
+      ],  # byteSize: 13 + 3 = 16 bytes
+      control: some(ControlMessage(
+        ihave: @[
+          ControlIHave(topicId: "ghi", messageIds: @[@[1'u8, 2, 3]])
+        ],  # 3 + 3 bytes
+        iwant: @[
+          ControlIWant(messageIds: @[@[1'u8, 2]])
+        ],  # 2 bytes
+        graft: @[
+          ControlGraft(topicId: "jkl")
+        ],  # 3 bytes
+        prune: @[
+          ControlPrune(topicId: "mno", peers: @[PeerInfoMsg(peerId: PeerId(data: @[]), signedPeerRecord: @[])], backoff: 1)
+        ]  # 3 + sizeof(uint64) bytes
+      )),
+      ping: @[],  # Empty seq[byte]
+      pong: @[]  # Empty seq[byte]
+    )
+
+    let boolSize = sizeof(bool)
+    let uint64Size = sizeof(uint64)
+    check byteSize(msg) == (3 + 3 + 2 * boolSize) + 16 + (3 + 3 + 2 + 3 + 3 + uint64Size)
\ No newline at end of file