From b61f1da9a8b0325e48b90cc13ee435470e38d174 Mon Sep 17 00:00:00 2001
From: Giovanni Petrantoni <giovanni@fragcolor.xyz>
Date: Tue, 21 Jul 2020 11:47:46 +0900
Subject: [PATCH] fixes

---
 libp2p/protocols/pubsub/gossipsub.nim | 22 ++++++++++++++--------
 tests/pubsub/utils.nim                |  8 ++++++--
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/libp2p/protocols/pubsub/gossipsub.nim b/libp2p/protocols/pubsub/gossipsub.nim
index b18d5b941..fed05e66b 100644
--- a/libp2p/protocols/pubsub/gossipsub.nim
+++ b/libp2p/protocols/pubsub/gossipsub.nim
@@ -479,8 +479,7 @@ proc updateScores(g: GossipSub) = # avoid async
         continue
 
     # Per topic
-    for topic in peer.topics:
-      var topicParams = g.topicParams.mgetOrPut(topic, TopicParams.init())
+    for topic, topicParams in g.topicParams:
       var info = stats.topicInfos.getOrDefault(topic)
 
       # Scoring
@@ -734,6 +733,15 @@ proc handleIWant(g: GossipSub,
       if msg.isSome:
         result.add(msg.get())
 
+proc punishPeer(g: GossipSub, peer: PubSubPeer, msg: Message) =
+  for t in msg.topicIDs:
+    # ensure we init a new topic if unknown
+    let _ = g.topicParams.mgetOrPut(t, TopicParams.init())
+    # update stats
+    var tstats = g.peerStats[peer].topicInfos.getOrDefault(t)
+    tstats.invalidMessageDeliveries += 1
+    g.peerStats[peer].topicInfos[t] = tstats
+
 method rpcHandler*(g: GossipSub,
                   peer: PubSubPeer,
                   rpcMsgs: seq[RPCMsg]) {.async.} =
@@ -755,15 +763,13 @@ method rpcHandler*(g: GossipSub,
         g.seen.put(msgId)                        # add the message to the seen cache
 
         if g.verifySignature and not msg.verify(peer.peerInfo):
-          trace "dropping message due to failed signature verification"
+          trace "dropping message due to failed signature verification", peer
+          g.punishPeer(peer, msg)
           continue
 
         if not (await g.validate(msg)):
-          trace "dropping message due to failed validation"
-          for t in msg.topicIDs:
-            var tstats = g.peerStats[peer].topicInfos.getOrDefault(t)
-            tstats.invalidMessageDeliveries += 1
-            g.peerStats[peer].topicInfos[t] = tstats
+          trace "dropping message due to failed validation", peer
+          g.punishPeer(peer, msg)
           continue
 
         # this shouldn't happen
diff --git a/tests/pubsub/utils.nim b/tests/pubsub/utils.nim
index 337364cc8..217e6e70c 100644
--- a/tests/pubsub/utils.nim
+++ b/tests/pubsub/utils.nim
@@ -1,13 +1,17 @@
-import random
+import random, options
 import chronos
 import ../../libp2p/standard_setup
+import ../../libp2p/protocols/pubsub/gossipsub
 export standard_setup
 
 randomize()
 
 proc generateNodes*(num: Natural, gossip: bool = false): seq[Switch] =
   for i in 0..<num:
-    result.add(newStandardSwitch(gossip = gossip))
+    var switch = newStandardSwitch(gossip = gossip)
+    var gossip = GossipSub(switch.pubSub.get())
+    gossip.parameters.floodPublish = false
+    result.add(switch)
 
 proc subscribeNodes*(nodes: seq[Switch]) {.async.} =
   var dials: seq[Future[void]]