From 75a5ab9441227b1d4a25b56c2bdfd544c5d28816 Mon Sep 17 00:00:00 2001
From: HelloZeroNet <hello@noloop.me>
Date: Mon, 25 Apr 2016 02:23:06 +0200
Subject: [PATCH] Track used site per connection

---
 src/Connection/Connection.py | 4 +++-
 src/File/FileRequest.py      | 2 ++
 src/Peer/Peer.py             | 2 ++
 3 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/Connection/Connection.py b/src/Connection/Connection.py
index d61d4613..4843da58 100644
--- a/src/Connection/Connection.py
+++ b/src/Connection/Connection.py
@@ -15,7 +15,7 @@ class Connection(object):
         "sock", "sock_wrapped", "ip", "port", "cert_pin", "site_lock", "id", "protocol", "type", "server", "unpacker", "req_id",
         "handshake", "crypt", "connected", "event_connected", "closed", "start_time", "last_recv_time",
         "last_message_time", "last_send_time", "last_sent_time", "incomplete_buff_recv", "bytes_recv", "bytes_sent",
-        "last_ping_delay", "last_req_time", "last_cmd", "bad_actions", "name", "updateName", "waiting_requests", "waiting_streams"
+        "last_ping_delay", "last_req_time", "last_cmd", "bad_actions", "sites", "name", "updateName", "waiting_requests", "waiting_streams"
     )
 
     def __init__(self, server, ip, port, sock=None, site_lock=None):
@@ -55,6 +55,7 @@ class Connection(object):
         self.last_req_time = 0
         self.last_cmd = None
         self.bad_actions = 0
+        self.sites = 0
 
         self.name = None
         self.updateName()
@@ -433,6 +434,7 @@ class Connection(object):
             request.set(False)
         self.waiting_requests = {}
         self.waiting_streams = {}
+        self.sites = 0
         self.server.removeConnection(self)  # Remove connection from server registry
         try:
             if self.sock:
diff --git a/src/File/FileRequest.py b/src/File/FileRequest.py
index 080e6431..b9a04053 100644
--- a/src/File/FileRequest.py
+++ b/src/File/FileRequest.py
@@ -134,6 +134,8 @@ class FileRequest(object):
             else:
                 peer = site.addPeer(self.connection.ip, self.connection.port, return_peer=True)  # Add or get peer
             if peer:
+                if not peer.connection:
+                    peer.connect(self.connection)  # Assign current connection to peer
                 peer.last_content_json_update = site.content_manager.contents[params["inner_path"]]["modified"]
                 if config.verbose:
                     self.log.debug(
diff --git a/src/Peer/Peer.py b/src/Peer/Peer.py
index 5d4b061f..00c21e9b 100644
--- a/src/Peer/Peer.py
+++ b/src/Peer/Peer.py
@@ -61,6 +61,7 @@ class Peer(object):
 
         if connection:  # Connection specified
             self.connection = connection
+            self.connection.sites += 1
         else:  # Try to find from connection pool or create new connection
             self.connection = None
 
@@ -71,6 +72,7 @@ class Peer(object):
                     self.connection = self.site.connection_server.getConnection(self.ip, self.port, site=self.site)
                 else:
                     self.connection = sys.modules["main"].file_server.getConnection(self.ip, self.port, site=self.site)
+                self.connection.sites += 1
 
             except Exception, err:
                 self.onConnectionError()