Simplify getting address from peer class (#1060)

* Simplify getting address from peer class

* Improved getting caddress field in version messages

* Simplified getting caddress even more

* Include local address and remote address as params for peer constructor

* Remove old comments from peer class

squeaknode/admin/messages.py

@@ -161,7 +161,7 @@ def payment_summary_to_message(
 
 def connected_peer_to_message(connected_peer: Peer) -> squeak_admin_pb2.ConnectedPeer:
     return squeak_admin_pb2.ConnectedPeer(
-        peer_address=peer_address_to_message(connected_peer.peer_address),
+        peer_address=peer_address_to_message(connected_peer.remote_address),
         connect_time_s=connected_peer.connect_time,
     )
 

squeaknode/network/connection_manager.py

@@ -51,28 +51,28 @@ class ConnectionManager(object):
                     return True
         return False
 
-    def add_peer(self, peer):
+    def add_peer(self, peer: Peer):
         """Add a peer.
         """
         with self.peers_lock:
             if self._is_duplicate_nonce(peer):
                 logger.debug('Failed to add peer {}'.format(peer))
                 raise DuplicateNonceError()
-            if self.has_connection(peer.address):
+            if self.has_connection(peer.remote_address):
                 logger.debug('Failed to add peer {}'.format(peer))
                 raise DuplicatePeerError()
-            self._peers[peer.address] = peer
+            self._peers[peer.remote_address] = peer
             logger.debug('Added peer {}'.format(peer))
             self.on_peers_changed()
 
-    def remove_peer(self, peer):
+    def remove_peer(self, peer: Peer):
         """Add a peer.
         """
         with self.peers_lock:
-            if not self.has_connection(peer.address):
+            if not self.has_connection(peer.remote_address):
                 logger.debug('Failed to remove peer {}'.format(peer))
                 raise MissingPeerError()
-            del self._peers[peer.address]
+            del self._peers[peer.remote_address]
             logger.debug('Removed peer {}'.format(peer))
             self.on_peers_changed()
 

squeaknode/network/network_manager.py

@@ -28,6 +28,8 @@ class NetworkManager(object):
 
     def __init__(self, config):
         self.config = config
+        self.local_ip = socket.gethostbyname('localhost')
+        self.local_port = self.config.server.rpc_port or squeak.params.params.DEFAULT_PORT
         self.peer_server = None
         self.peer_client = None
         self.connection_manager = ConnectionManager()
@@ -39,7 +41,7 @@ class NetworkManager(object):
         )
         self.peer_server = PeerServer(
             peer_handler,
-            self.config.server.rpc_port,
+            self.local_port,
         )
         self.peer_client = PeerClient(
             peer_handler,
@@ -78,14 +80,25 @@ class NetworkManager(object):
                     peer,
                 ))
 
-    def handle_connection(self, squeak_controller, peer_socket, address, outgoing):
+    def handle_connection(
+            self,
+            squeak_controller,
+            peer_socket: socket.socket,
+            address: PeerAddress,
+            outgoing: bool,
+    ):
         """Handles all sending and receiving of messages for the given peer.
 
         This method blocks until the peer connection has stopped.
         """
         logger.debug(
             'Setting up controller for peer address {} ...'.format(address))
-        with Peer(peer_socket, address, outgoing).open_connection(squeak_controller) as peer:
+        with Peer(
+                peer_socket,
+                self.local_address,
+                address,
+                outgoing,
+        ).open_connection(squeak_controller) as peer:
             self.connection_manager.add_peer(peer)
             try:
                 peer.handle_messages(squeak_controller)
@@ -95,15 +108,12 @@ class NetworkManager(object):
                 self.connection_manager.remove_peer(peer)
         logger.debug('Stopped controller for peer address {}.'.format(address))
 
-    def get_address(self):
-        # TODO: Add return type.
-        return (self.peer_server.ip, self.peer_server.port)
-
-    def get_remote_address(self, address):
-        # TODO: Add return type.
-        hostname, port = address
-        ip = socket.gethostbyname(hostname)
-        return (ip, port)
+    @property
+    def local_address(self) -> PeerAddress:
+        return PeerAddress(
+            self.local_ip,
+            self.local_port,
+        )
 
     # def register_peers_changed_callback(self, callback, stopped: threading.Event):
     #     """Registers a callback that gets called when connected peers changes.

squeaknode/network/peer.py

@@ -16,6 +16,7 @@ from squeak.messages import MsgSerializable
 from squeaknode.core.peer_address import PeerAddress
 from squeaknode.core.util import generate_version_nonce
 from squeaknode.network.peer_message_handler import PeerMessageHandler
+from squeaknode.network.util import time_now
 
 
 MAX_MESSAGE_LEN = 1048576
@@ -36,20 +37,26 @@ class Peer(object):
     """Maintains the internal state of a peer connection.
     """
 
-    def __init__(self, peer_socket, address, outgoing=False):
-        time_now = int(time.time())
+    def __init__(
+            self,
+            peer_socket: socket.socket,
+            local_address: PeerAddress,
+            remote_address: PeerAddress,
+            outgoing: bool,
+    ):
         self._peer_socket = peer_socket
         self._peer_socket_lock = threading.Lock()
-        self._address = address
+        self._local_address = local_address
+        self._remote_address = remote_address
         self._outgoing = outgoing
-        self._connect_time = time_now
+        self._connect_time = 0
         self._local_version = None
         self._remote_version = None
         self._last_msg_revc_time = None
         self._last_sent_ping_nonce = None
         self._last_sent_ping_time = None
         self._last_recv_ping_time = None
-        self._recv_msg_queue = queue.Queue()
+        self._recv_msg_queue: queue.Queue = queue.Queue()
 
         self._subscription = None
 
@@ -65,46 +72,31 @@ class Peer(object):
             return remote_version.nVersion
 
     @property
-    def address(self):
-        return self._address
+    def local_address(self):
+        return self._local_address
+
+    @property
+    def remote_address(self):
+        return self._remote_address
 
     @property
     def subscription(self):
         return self._subscription
 
     @property
-    def peer_address(self):
-        # TODO: Just return the peer address object
-        ip, port = self._address
-        return PeerAddress(
-            host=ip,
-            port=port,
-        )
-
-    @property
-    def address_string(self):
-        ip, port = self._address
-        return '{}:{}'.format(ip, port)
-
-    @property
-    def ip(self):
-        # TODO: Just return self._address.host
-        ip, _ = self._address
-        return ip
-
-    @property
-    def port(self):
-        # TODO: Just return self._address.port
-        _, port = self._address
-        return port
-
-    @property
-    def caddress(self):
-        ip, port = self._address
+    def local_caddress(self):
         caddress = CAddress()
         caddress.nTime = self.connect_time
-        caddress.ip = ip
-        caddress.port = port
+        caddress.ip = self.local_address.host
+        caddress.port = self.local_address.port
+        return caddress
+
+    @property
+    def remote_caddress(self):
+        caddress = CAddress()
+        caddress.nTime = self.connect_time
+        caddress.ip = socket.gethostbyname(self.remote_address.host)
+        caddress.port = self.remote_address.port
         return caddress
 
     @property
@@ -223,14 +215,9 @@ class Peer(object):
     def version_pkt(self, squeak_controller):
         """Get the version message for this peer."""
         msg = msg_version()
-        local_ip, local_port = squeak_controller.get_address()
-        server_ip, server_port = squeak_controller.get_remote_address(
-            self.address)
         msg.nVersion = HANDSHAKE_VERSION
-        msg.addrTo.ip = server_ip
-        msg.addrTo.port = server_port
-        msg.addrFrom.ip = local_ip
-        msg.addrFrom.port = local_port
+        msg.addrTo = self.remote_caddress
+        msg.addrFrom = self.local_caddress
         msg.nNonce = generate_version_nonce()
         return msg
 
@@ -241,6 +228,9 @@ class Peer(object):
         )
         self.send_msg(subscribe_msg)
 
+    def set_connected(self):
+        self._connect_time = time_now()
+
     def set_subscription(self, subscription):
         self._subscription = subscription
 
@@ -261,13 +251,14 @@ class Peer(object):
             ).start()
             self.handshake(squeak_controller)
             self.update_subscription(squeak_controller)
+            self.set_connected()
             yield self
         finally:
             self.close()
             logger.debug('Closed connection to peer {} ...'.format(self))
 
     def __repr__(self):
-        return "Peer(%s)" % (self.address_string)
+        return "Peer(%s)" % (str(self.remote_address))
 
 
 class MessageDecoder:

squeaknode/network/peer_client.py

@@ -2,6 +2,8 @@ import logging
 import socket
 import threading
 
+from squeaknode.core.peer_address import PeerAddress
+
 
 SOCKET_CONNECT_TIMEOUT = 5
 
@@ -16,7 +18,7 @@ class PeerClient(object):
     def __init__(self, peer_handler):
         self.peer_handler = peer_handler
 
-    def make_connection(self, address):
+    def make_connection(self, address: PeerAddress):
         logger.info('Making connection to {}'.format(address))
         try:
             peer_socket = socket.socket()
@@ -29,7 +31,7 @@ class PeerClient(object):
         except Exception:
             logger.exception('Failed to make connection to {}'.format(address))
 
-    def connect_address(self, address):
+    def connect_address(self, address: PeerAddress):
         """Connect to new address."""
         logger.info('Connecting to peer with address {}'.format(address))
         threading.Thread(

squeaknode/network/peer_handler.py

@@ -1,6 +1,9 @@
 import logging
+import socket
 import threading
 
+from squeaknode.core.peer_address import PeerAddress
+
 
 logger = logging.getLogger(__name__)
 
@@ -18,7 +21,7 @@ class PeerHandler():
         self.squeak_controller = squeak_controller
         self.handle_connection_fn = handle_connection_fn
 
-    def handle_connection(self, peer_socket, address, outgoing):
+    def handle_connection(self, peer_socket: socket.socket, address: PeerAddress, outgoing: bool):
         threading.Thread(
             target=self.handle_connection_fn,
             args=(self.squeak_controller, peer_socket, address, outgoing,),

squeaknode/network/peer_message_handler.py

@@ -42,20 +42,6 @@ class PeerMessageHandler:
 
     def handle_peer_message(self, msg):
         """Handle messages from a peer with completed handshake."""
-
-        # # Only allow version and verack messages before handshake is complete.
-        # if not self.peer.is_handshake_complete and msg.command not in [
-        #         b'version',
-        #         b'verack',
-        # ]:
-        #     raise Exception(
-        #         'Received non-handshake message from un-handshaked peer.')
-
-        # if msg.command == b'version':
-        #     self.handle_version(msg)
-        # if msg.command == b'verack':
-        #     self.handle_verack(msg)
-
         if msg.command == b'ping':
             self.handle_ping(msg)
         if msg.command == b'pong':
@@ -122,7 +108,7 @@ class PeerMessageHandler:
             if inv.type == 2:
                 offer = self.squeak_controller.get_buy_offer(
                     squeak_hash=inv.hash,
-                    peer_address=self.peer.peer_address,
+                    peer_address=self.peer.remote_address,
                 )
                 if offer is None:
                     not_found.append(inv)
@@ -164,7 +150,7 @@ class PeerMessageHandler:
             decoded_offer = self.squeak_controller.get_offer(
                 squeak=squeak,
                 offer=offer,
-                peer_address=self.peer.peer_address,
+                peer_address=self.peer.remote_address,
             )
             self.squeak_controller.save_offer(decoded_offer)
 

squeaknode/network/peer_server.py

@@ -2,7 +2,7 @@ import logging
 import socket
 import threading
 
-import squeak.params
+from squeaknode.core.peer_address import PeerAddress
 
 
 MIN_PEERS = 5
@@ -17,10 +17,9 @@ class PeerServer(object):
     """Maintains connections to other peers in the network.
     """
 
-    def __init__(self, peer_handler, port=None):
+    def __init__(self, peer_handler, port):
         self.peer_handler = peer_handler
-        self.ip = socket.gethostbyname('localhost')
-        self.port = port or squeak.params.params.DEFAULT_PORT
+        self.port = port
         self.listen_socket = socket.socket()
 
     def start(self):
@@ -44,8 +43,13 @@ class PeerServer(object):
             self.listen_socket.listen()
             while True:
                 peer_socket, address = self.listen_socket.accept()
+                host, port = address
+                peer_address = PeerAddress(
+                    host=host,
+                    port=port,
+                )
                 peer_socket.setblocking(True)
                 self.peer_handler.handle_connection(
-                    peer_socket, address, outgoing=False)
+                    peer_socket, peer_address, outgoing=False)
         except Exception:
             logger.info("Stopped accepting incoming connections.")

squeaknode/network/util.py

@@ -0,0 +1,5 @@
+import time
+
+
+def time_now():
+    return int(time.time())

squeaknode/node/squeak_controller.py

@@ -473,12 +473,6 @@ class SqueakController:
                 peer.address,
             )
 
-    def get_address(self):
-        return self.network_manager.get_address()
-
-    def get_remote_address(self, address):
-        return self.network_manager.get_remote_address(address)
-
     def get_connected_peer(self, peer_address: PeerAddress):
         return self.network_manager.get_connected_peer(peer_address)