From 242a161585df976c12e7a290cbf34a337fbc3701 Mon Sep 17 00:00:00 2001
From: Dmitri Bogomolov <4glitch@gmail.com>
Date: Fri, 15 May 2020 09:37:59 +0300
Subject: [PATCH] Use stream chosen for connection in version message

---
 src/network/bmproto.py        |  1 +
 src/network/connectionpool.py | 17 +++++++++++------
 src/network/tcp.py            | 10 ++++++----
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/src/network/bmproto.py b/src/network/bmproto.py
index 008eadb0..fc6ee92d 100644
--- a/src/network/bmproto.py
+++ b/src/network/bmproto.py
@@ -69,6 +69,7 @@ class BMProto(AdvancedDispatcher, ObjectTracker):
         self.network_group = None
         # userAgent initialization
         self.userAgent = ''
+        self.stream = None
 
     def bm_proto_reset(self):
         """Reset the bitmessage object parser"""
diff --git a/src/network/connectionpool.py b/src/network/connectionpool.py
index e85b57e5..42d2afc4 100644
--- a/src/network/connectionpool.py
+++ b/src/network/connectionpool.py
@@ -299,8 +299,10 @@ class BMConnectionPool(object):
                 for i in range(
                         state.maximumNumberOfHalfOpenConnections - pending):
                     try:
+                        target_stream = helper_random.randomchoice(
+                            list(self.streams))
                         chosen = self.trustedPeer or chooseConnection(
-                            helper_random.randomchoice(list(self.streams)))
+                            target_stream)
                     except (ValueError, IndexError):
                         continue
                     if chosen in self.outboundConnections:
@@ -327,15 +329,18 @@ class BMConnectionPool(object):
                     try:
                         if chosen.host.endswith(".onion") and Proxy.onion_proxy:
                             if onionsocksproxytype == "SOCKS5":
-                                self.addConnection(Socks5BMConnection(chosen))
+                                proto = Socks5BMConnection(chosen)
                             elif onionsocksproxytype == "SOCKS4a":
-                                self.addConnection(Socks4aBMConnection(chosen))
+                                proto = Socks4aBMConnection(chosen)
                         elif socksproxytype == "SOCKS5":
-                            self.addConnection(Socks5BMConnection(chosen))
+                            proto = Socks5BMConnection(chosen)
                         elif socksproxytype == "SOCKS4a":
-                            self.addConnection(Socks4aBMConnection(chosen))
+                            proto = Socks4aBMConnection(chosen)
                         else:
-                            self.addConnection(TCPConnection(chosen))
+                            proto = TCPConnection(chosen)
+
+                        proto.stream = target_stream
+                        self.addConnection(proto)
                     except socket.error as e:
                         if e.errno == errno.ENETUNREACH:
                             continue
diff --git a/src/network/tcp.py b/src/network/tcp.py
index ff778378..e7eaa346 100644
--- a/src/network/tcp.py
+++ b/src/network/tcp.py
@@ -267,8 +267,9 @@ class TCPConnection(BMProto, TLSDispatcher):
         self.append_write_buf(
             protocol.assembleVersionMessage(
                 self.destination.host, self.destination.port,
-                connectionpool.BMConnectionPool().streams,
-                False, nodeid=self.nodeid))
+                [self.stream] if self.stream
+                else connectionpool.BMConnectionPool().streams,
+                nodeid=self.nodeid))
         self.connectedAt = time.time()
         receiveDataQueue.put(self.destination)
 
@@ -318,8 +319,9 @@ class Socks5BMConnection(Socks5Connection, TCPConnection):
         self.append_write_buf(
             protocol.assembleVersionMessage(
                 self.destination.host, self.destination.port,
-                connectionpool.BMConnectionPool().streams,
-                False, nodeid=self.nodeid))
+                [self.stream] if self.stream
+                else connectionpool.BMConnectionPool().streams,
+                nodeid=self.nodeid))
         self.set_state("bm_header", expectBytes=protocol.Header.size)
         return True