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 self.network_group = None
# userAgent initialization # userAgent initialization
self.userAgent = '' self.userAgent = ''
self.stream = None
def bm_proto_reset(self): def bm_proto_reset(self):
"""Reset the bitmessage object parser""" """Reset the bitmessage object parser"""

View File

@ -299,8 +299,10 @@ class BMConnectionPool(object):
for i in range( for i in range(
state.maximumNumberOfHalfOpenConnections - pending): state.maximumNumberOfHalfOpenConnections - pending):
try: try:
target_stream = helper_random.randomchoice(
list(self.streams))
chosen = self.trustedPeer or chooseConnection( chosen = self.trustedPeer or chooseConnection(
helper_random.randomchoice(list(self.streams))) target_stream)
except (ValueError, IndexError): except (ValueError, IndexError):
continue continue
if chosen in self.outboundConnections: if chosen in self.outboundConnections:
@ -327,15 +329,18 @@ class BMConnectionPool(object):
try: try:
if chosen.host.endswith(".onion") and Proxy.onion_proxy: if chosen.host.endswith(".onion") and Proxy.onion_proxy:
if onionsocksproxytype == "SOCKS5": if onionsocksproxytype == "SOCKS5":
self.addConnection(Socks5BMConnection(chosen)) proto = Socks5BMConnection(chosen)
elif onionsocksproxytype == "SOCKS4a": elif onionsocksproxytype == "SOCKS4a":
self.addConnection(Socks4aBMConnection(chosen)) proto = Socks4aBMConnection(chosen)
elif socksproxytype == "SOCKS5": elif socksproxytype == "SOCKS5":
self.addConnection(Socks5BMConnection(chosen)) proto = Socks5BMConnection(chosen)
elif socksproxytype == "SOCKS4a": elif socksproxytype == "SOCKS4a":
self.addConnection(Socks4aBMConnection(chosen)) proto = Socks4aBMConnection(chosen)
else: else:
self.addConnection(TCPConnection(chosen)) proto = TCPConnection(chosen)
proto.stream = target_stream
self.addConnection(proto)
except socket.error as e: except socket.error as e:
if e.errno == errno.ENETUNREACH: if e.errno == errno.ENETUNREACH:
continue continue

View File

@ -267,8 +267,9 @@ class TCPConnection(BMProto, TLSDispatcher):
self.append_write_buf( self.append_write_buf(
protocol.assembleVersionMessage( protocol.assembleVersionMessage(
self.destination.host, self.destination.port, self.destination.host, self.destination.port,
connectionpool.BMConnectionPool().streams, [self.stream] if self.stream
False, nodeid=self.nodeid)) else connectionpool.BMConnectionPool().streams,
nodeid=self.nodeid))
self.connectedAt = time.time() self.connectedAt = time.time()
receiveDataQueue.put(self.destination) receiveDataQueue.put(self.destination)
@ -318,8 +319,9 @@ class Socks5BMConnection(Socks5Connection, TCPConnection):
self.append_write_buf( self.append_write_buf(
protocol.assembleVersionMessage( protocol.assembleVersionMessage(
self.destination.host, self.destination.port, self.destination.host, self.destination.port,
connectionpool.BMConnectionPool().streams, [self.stream] if self.stream
False, nodeid=self.nodeid)) else connectionpool.BMConnectionPool().streams,
nodeid=self.nodeid))
self.set_state("bm_header", expectBytes=protocol.Header.size) self.set_state("bm_header", expectBytes=protocol.Header.size)
return True return True