Use stream chosen for connection in version message

This commit is contained in:
Dmitri Bogomolov 2020-05-15 09:37:59 +03:00
parent 920b5dc25b
commit 242a161585
Signed by untrusted user: g1itch
GPG Key ID: 720A756F18DEED13
3 changed files with 18 additions and 10 deletions

View File

@ -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"""

View File

@ -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

View File

@ -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