diff --git a/squeaknode/node/active_download_manager.py b/squeaknode/node/active_download_manager.py index 683851a8..3f8a4c68 100644 --- a/squeaknode/node/active_download_manager.py +++ b/squeaknode/node/active_download_manager.py @@ -29,9 +29,7 @@ from concurrent.futures import ThreadPoolExecutor from concurrent.futures import TimeoutError from typing import Dict from typing import Optional -from typing import Union -from squeak.core import CSqueak from squeak.messages import msg_getdata from squeak.messages import msg_getsqueaks from squeak.messages import MsgSerializable @@ -40,9 +38,7 @@ from squeak.net import CInv from squeak.net import CSqueakLocator from squeaknode.core.download_result import DownloadResult -from squeaknode.core.interests import squeak_matches_interest -from squeaknode.core.offer import Offer -from squeaknode.core.squeaks import get_hash +from squeaknode.node.downloaded_object import DownloadedObject logger = logging.getLogger(__name__) @@ -61,7 +57,7 @@ class ActiveDownload(ABC): self.start_time_ms: Optional[int] = None @abstractmethod - def is_interested(self, downloaded_object: Union[CSqueak, Offer]) -> bool: + def is_interested(self, downloaded_object: DownloadedObject) -> bool: """Return True if the given squeak matches the download interest.""" @abstractmethod @@ -111,10 +107,8 @@ class InterestDownload(ActiveDownload): self.interest = interest super().__init__(limit) - def is_interested(self, downloaded_object: Union[CSqueak, Offer]) -> bool: - if type(downloaded_object) is not CSqueak: - return False - return squeak_matches_interest(downloaded_object, self.interest) + def is_interested(self, downloaded_object: DownloadedObject) -> bool: + return downloaded_object.matches_requested_squeak_range(self.interest) def get_download_msg(self) -> MsgSerializable: locator = CSqueakLocator( @@ -131,10 +125,8 @@ class HashDownload(ActiveDownload): self.squeak_hash = squeak_hash super().__init__(1) - def is_interested(self, downloaded_object: Union[CSqueak, Offer]) -> bool: - if type(downloaded_object) is not CSqueak: - return False - return self.squeak_hash == get_hash(downloaded_object) + def is_interested(self, downloaded_object: DownloadedObject) -> bool: + return downloaded_object.matches_requested_squeak_hash(self.squeak_hash) def get_download_msg(self) -> MsgSerializable: invs = [ @@ -151,10 +143,8 @@ class OffersDownload(ActiveDownload): self.squeak_hash = squeak_hash super().__init__(limit) - def is_interested(self, downloaded_object: Union[CSqueak, Offer]) -> bool: - if type(downloaded_object) is not Offer: - return False - return downloaded_object.squeak_hash == self.squeak_hash + def is_interested(self, downloaded_object: DownloadedObject) -> bool: + return downloaded_object.matches_requested_offer_hash(self.squeak_hash) def get_download_msg(self) -> MsgSerializable: invs = [ @@ -182,7 +172,7 @@ class ActiveDownloadManager: self.executor.shutdown(wait=True) logger.info("Stopped Download Manager.") - def lookup_counter(self, downloaded_object: Union[CSqueak, Offer]) -> Optional[ActiveDownload]: + def lookup_counter(self, downloaded_object: DownloadedObject) -> Optional[ActiveDownload]: for name, interest in self.downloads.items(): if interest.is_interested(downloaded_object): return interest diff --git a/squeaknode/node/downloaded_object.py b/squeaknode/node/downloaded_object.py new file mode 100644 index 00000000..fa0ebc6e --- /dev/null +++ b/squeaknode/node/downloaded_object.py @@ -0,0 +1,72 @@ +# MIT License +# +# Copyright (c) 2020 Jonathan Zernik +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +from squeak.core import CSqueak +from squeak.net import CInterested + +from squeaknode.core.interests import squeak_matches_interest +from squeaknode.core.offer import Offer +from squeaknode.core.squeaks import get_hash + + +class DownloadedObject: + + def matches_requested_squeak_range(self, interest: CInterested) -> bool: + """Return True if the object matches the requested interest. + """ + + def matches_requested_squeak_hash(self, squeak_hash: bytes) -> bool: + """Return True if the object matches the requested offer hash. + """ + + def matches_requested_offer_hash(self, squeak_hash: bytes) -> bool: + """Return True if the object matches the requested squeak hash. + """ + + +class DownloadedSqueak(DownloadedObject): + + def __init__(self, squeak: CSqueak): + self.squeak = squeak + + def matches_requested_squeak_range(self, interest: CInterested) -> bool: + return squeak_matches_interest(self.squeak, interest) + + def matches_requested_squeak_hash(self, squeak_hash: bytes) -> bool: + return squeak_hash == get_hash(self.squeak) + + def matches_requested_offer_hash(self, squeak_hash: bytes) -> bool: + return False + + +class DownloadedOffer(DownloadedObject): + + def __init__(self, offer: Offer): + self.offer = offer + + def matches_requested_squeak_range(self, interest: CInterested) -> bool: + return False + + def matches_requested_squeak_hash(self, squeak_hash: bytes) -> bool: + return False + + def matches_requested_offer_hash(self, squeak_hash: bytes) -> bool: + return self.offer.squeak_hash == squeak_hash diff --git a/squeaknode/node/squeak_controller.py b/squeaknode/node/squeak_controller.py index 09b681a3..759c0378 100644 --- a/squeaknode/node/squeak_controller.py +++ b/squeaknode/node/squeak_controller.py @@ -64,6 +64,8 @@ from squeaknode.core.update_subscriptions_event import UpdateSubscriptionsEvent from squeaknode.core.update_twitter_stream_event import UpdateTwitterStreamEvent from squeaknode.core.user_config import UserConfig from squeaknode.node.active_download_manager import ActiveDownload +from squeaknode.node.downloaded_object import DownloadedOffer +from squeaknode.node.downloaded_object import DownloadedSqueak from squeaknode.node.listener_subscription_client import EventListener from squeaknode.node.price_policy import PricePolicy from squeaknode.node.received_payments_subscription_client import ReceivedPaymentsSubscriptionClient @@ -166,10 +168,12 @@ class SqueakController: ) < self.config.node.max_squeaks_per_address_in_block_range def get_download_squeak_counter(self, squeak: CSqueak) -> Optional[ActiveDownload]: - return self.active_download_manager.lookup_counter(squeak) + downloaded_squeak = DownloadedSqueak(squeak) + return self.active_download_manager.lookup_counter(downloaded_squeak) def get_download_offer_counter(self, offer: Offer) -> Optional[ActiveDownload]: - return self.active_download_manager.lookup_counter(offer) + downloaded_offer = DownloadedOffer(offer) + return self.active_download_manager.lookup_counter(downloaded_offer) def get_offer_or_secret_key(self, squeak_hash: bytes, peer_address: PeerAddress) -> Optional[Union[bytes, Offer]]: squeak = self.get_squeak(squeak_hash) diff --git a/tests/node/test_active_download.py b/tests/node/test_active_download.py index 22f672ab..77e4701c 100644 --- a/tests/node/test_active_download.py +++ b/tests/node/test_active_download.py @@ -30,6 +30,7 @@ from squeak.net import CSqueakLocator from squeaknode.core.download_result import DownloadResult from squeaknode.node.active_download_manager import HashDownload from squeaknode.node.active_download_manager import InterestDownload +from squeaknode.node.downloaded_object import DownloadedSqueak from tests.utils import gen_squeak @@ -53,25 +54,29 @@ def download_interest(interest): def test_download_hash_is_interested(download_hash, squeak): + downloaded_squeak = DownloadedSqueak(squeak) - assert download_hash.is_interested(squeak) + assert download_hash.is_interested(downloaded_squeak) def test_download_hash_is_not_interested(download_hash, signing_key, block_count): other_squeak = gen_squeak(signing_key, block_count) + downloaded_squeak = DownloadedSqueak(other_squeak) - assert not download_hash.is_interested(other_squeak) + assert not download_hash.is_interested(downloaded_squeak) def test_download_interest_is_interested(download_interest, squeak): + downloaded_squeak = DownloadedSqueak(squeak) - assert download_interest.is_interested(squeak) + assert download_interest.is_interested(downloaded_squeak) def test_download_interest_is_not_interested(download_interest, signing_key, block_count): other_squeak = gen_squeak(signing_key, block_count + 200) + downloaded_squeak = DownloadedSqueak(other_squeak) - assert not download_interest.is_interested(other_squeak) + assert not download_interest.is_interested(downloaded_squeak) def test_download_hash_initiate(download_hash, squeak_hash):