From b165a6b4efb9bb55c5368a49babf80cc6186a02b Mon Sep 17 00:00:00 2001
From: Dmitri Bogomolov <4glitch@gmail.com>
Date: Wed, 10 Jun 2020 13:22:49 +0300
Subject: [PATCH] Do addKnownNode recursively for multiple streams

---
 src/knownnodes.py  | 16 ++++++++++++++--
 src/network/tcp.py | 11 ++++-------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/src/knownnodes.py b/src/knownnodes.py
index 77c52a25..07871c7c 100644
--- a/src/knownnodes.py
+++ b/src/knownnodes.py
@@ -8,12 +8,16 @@ import os
 import pickle
 import threading
 import time
+try:
+    from collections.abc import Iterable
+except ImportError:
+    from collections import Iterable
 
 import state
 from bmconfigparser import BMConfigParser
 from network.node import Peer
 
-knownNodesLock = threading.Lock()
+knownNodesLock = threading.RLock()
 """Thread lock for knownnodes modification"""
 knownNodes = {stream: {} for stream in range(1, 4)}
 """The dict of known nodes for each stream"""
@@ -95,9 +99,17 @@ def saveKnownNodes(dirName=None):
 
 def addKnownNode(stream, peer, lastseen=None, is_self=False):
     """
-    Add a new node to the dict or update lastseen if it already exists
+    Add a new node to the dict or update lastseen if it already exists.
+    Do it for each stream number if *stream* is `Iterable`.
+    Returns True if added a new node.
     """
     # pylint: disable=too-many-branches
+    if isinstance(stream, Iterable):
+        with knownNodesLock:
+            for s in stream:
+                addKnownNode(s, peer, lastseen, is_self)
+        return
+
     rating = 0.0
     if not lastseen:
         # FIXME: maybe about 28 days?
diff --git a/src/network/tcp.py b/src/network/tcp.py
index 320e776a..0bdd3269 100644
--- a/src/network/tcp.py
+++ b/src/network/tcp.py
@@ -148,9 +148,8 @@ class TCPConnection(BMProto, TLSDispatcher):
         # The connection having host suitable for knownnodes
         if self.isOutbound or not self.local and not state.socksIP:
             knownnodes.increaseRating(self.destination)
-            with knownnodes.knownNodesLock:
-                for s in self.streams:
-                    knownnodes.addKnownNode(s, self.destination, time.time())
+            knownnodes.addKnownNode(
+                self.streams, self.destination, time.time())
             Dandelion().maybeAddStem(self)
         self.sendAddr()
         self.sendBigInv()
@@ -273,10 +272,8 @@ class TCPConnection(BMProto, TLSDispatcher):
                 (self.isOutbound, False, self.destination)
             ))
             if host_is_global:
-                with knownnodes.knownNodesLock:
-                    for s in self.streams:
-                        knownnodes.addKnownNode(
-                            s, self.destination, time.time())
+                knownnodes.addKnownNode(
+                    self.streams, self.destination, time.time())
                 Dandelion().maybeRemoveStem(self)
         elif host_is_global:
             knownnodes.decreaseRating(self.destination)