From 69f31dda8b780e6e10ad8b0e9598e2da3f74cdbd Mon Sep 17 00:00:00 2001 From: Jonathan Zernik Date: Sat, 25 Dec 2021 15:43:29 -0800 Subject: [PATCH] Refactor connection methods into network handler (#1892) * Refactor get unknown invs method to network handler * Rename get unknown invs methods * Rename get unknown invs methods * Move get unknown invs method to network handler --- squeaknode/network/connection.py | 34 +++++++++--------------------- squeaknode/node/network_handler.py | 22 +++++++++++++++++++ 2 files changed, 32 insertions(+), 24 deletions(-) diff --git a/squeaknode/network/connection.py b/squeaknode/network/connection.py index 96f0522e..958fe0aa 100644 --- a/squeaknode/network/connection.py +++ b/squeaknode/network/connection.py @@ -34,6 +34,7 @@ from squeak.messages import msg_pong from squeak.messages import MSG_SECRET_KEY from squeak.messages import MSG_SQUEAK from squeak.messages import msg_squeak +from squeak.messages import MsgSerializable from squeak.net import CInterested from squeak.net import CInv @@ -183,34 +184,16 @@ class Connection(object): def handle_inv(self, msg): invs = msg.inv - unknown_squeak_invs = [ - inv for inv in invs - if inv.type == MSG_SQUEAK - and self.network_handler.get_squeak(inv.hash) is None - ] - unknown_secret_key_invs = [ - inv for inv in invs - if inv.type == MSG_SECRET_KEY - and self.network_handler.get_squeak(inv.hash) is not None - and self.network_handler.get_squeak_secret_key(inv.hash) is None - ] - unknown_invs = unknown_squeak_invs + unknown_secret_key_invs + unknown_invs = self.network_handler.get_unknown_invs(invs) if unknown_invs: getdata_msg = msg_getdata(inv=unknown_invs) self.peer.send_msg(getdata_msg) def handle_getdata(self, msg): invs = msg.inv - not_found = [] for inv in invs: reply_msg = self._get_inv_reply(inv) - if reply_msg is not None: - self.peer.send_msg(reply_msg) - else: - not_found.append(inv) - if not_found: - notfound_msg = msg_notfound(inv=not_found) - self.peer.send_msg(notfound_msg) + self.peer.send_msg(reply_msg) def handle_notfound(self, msg): pass @@ -326,24 +309,27 @@ class Connection(object): reply_to_hash=reply_to_hash, ) - def _get_inv_reply(self, inv): + def _get_inv_reply(self, inv) -> MsgSerializable: if inv.type == MSG_SQUEAK: return self._get_inv_reply_for_squeak(inv) if inv.type == MSG_SECRET_KEY: return self._get_inv_reply_for_secret_key(inv) + raise Exception("Uknown inv type.") - def _get_inv_reply_for_squeak(self, inv): + def _get_inv_reply_for_squeak(self, inv) -> MsgSerializable: squeak = self.network_handler.get_squeak(inv.hash) if squeak is not None: return msg_squeak(squeak=squeak) + else: + return msg_notfound(inv=[inv]) - def _get_inv_reply_for_secret_key(self, inv): + def _get_inv_reply_for_secret_key(self, inv) -> MsgSerializable: resp = self.network_handler.get_secret_key_reply( inv.hash, self.peer.remote_address, ) if resp is None: - return None + return msg_notfound(inv=[inv]) return resp.get_msg() diff --git a/squeaknode/node/network_handler.py b/squeaknode/node/network_handler.py index 46457f64..d39a09db 100644 --- a/squeaknode/node/network_handler.py +++ b/squeaknode/node/network_handler.py @@ -26,6 +26,8 @@ from typing import Optional from squeak.core import CSqueak from squeak.core.signing import SqueakPublicKey from squeak.messages import msg_getdata +from squeak.messages import MSG_SECRET_KEY +from squeak.messages import MSG_SQUEAK from squeak.messages import MsgSerializable from squeak.net import CInv from squeak.net import CSqueakLocator @@ -68,6 +70,26 @@ class NetworkHandler: def get_squeak_secret_key(self, squeak_hash: bytes) -> Optional[bytes]: return self.squeak_store.get_squeak_secret_key(squeak_hash) + def get_unknown_invs(self, invs): + unknown_squeak_invs = self.get_unknown_squeaks(invs) + unknown_secret_key_invs = self.get_unknown_secret_keys(invs) + return unknown_squeak_invs + unknown_secret_key_invs + + def get_unknown_squeaks(self, invs): + return [ + inv for inv in invs + if inv.type == MSG_SQUEAK + and self.get_squeak(inv.hash) is None + ] + + def get_unknown_secret_keys(self, invs): + return [ + inv for inv in invs + if inv.type == MSG_SECRET_KEY + and self.get_squeak(inv.hash) is not None + and self.get_squeak_secret_key(inv.hash) is None + ] + def save_squeak(self, squeak: CSqueak) -> Optional[bytes]: return self.squeak_store.save_squeak(squeak)