diff --git a/itests/tests/test_squeak_node.py b/itests/tests/test_squeak_node.py index 56c7e865..55b1a2fd 100644 --- a/itests/tests/test_squeak_node.py +++ b/itests/tests/test_squeak_node.py @@ -108,9 +108,29 @@ def test_post_squeak_not_following( ) squeak_msg = build_squeak_msg(squeak) - with pytest.raises(Exception): + with pytest.raises(Exception) as e: server_stub.UploadSqueak( squeak_server_pb2.UploadSqueakRequest(squeak=squeak_msg)) + assert "Squeak address not in followed list." in e.details() + + +def test_post_squeak_without_decryption_key( + server_stub, admin_stub, lightning_client, following_signing_key +): + # Post a squeak with a direct request to the server + block_height, block_hash = get_latest_block_info(lightning_client) + squeak = make_squeak( + following_signing_key, "hello from itest!", block_hash, block_height + ) + + # Clear the decryption key from the squeak + squeak.ClearDecryptionKey() + + squeak_msg = build_squeak_msg(squeak) + with pytest.raises(Exception) as e: + server_stub.UploadSqueak( + squeak_server_pb2.UploadSqueakRequest(squeak=squeak_msg)) + assert "Squeak must contain decryption key." in e.details() def test_lookup_squeaks(server_stub, admin_stub, signing_profile_id, saved_squeak_hash): @@ -375,8 +395,8 @@ def test_make_reply_squeak( def test_post_squeak_rate_limit(server_stub, admin_stub, lightning_client, following_signing_key): - # Make 10 squeak - for i in range(10): + # Make 11 squeak (rate limit is 5, so this should fail even with two blocks) + for i in range(11): try: block_height, block_hash = get_latest_block_info(lightning_client) squeak = make_squeak( @@ -393,7 +413,8 @@ def test_post_squeak_rate_limit(server_stub, admin_stub, lightning_client, follo except Exception as e: post_squeak_exception = e assert post_squeak_exception is not None - assert "Excedeed allowed number of squeaks per block" in post_squeak_exception.details() + assert "Exceeded allowed number of squeaks per address per block." in \ + post_squeak_exception.details() def test_make_signing_profile(server_stub, admin_stub): diff --git a/squeaknode/core/squeak_controller.py b/squeaknode/core/squeak_controller.py index 43e53a4a..a2957810 100644 --- a/squeaknode/core/squeak_controller.py +++ b/squeaknode/core/squeak_controller.py @@ -41,35 +41,46 @@ class SqueakController: self.create_offer_lock = threading.Lock() def save_uploaded_squeak(self, squeak: CSqueak) -> bytes: - if not self.squeak_rate_limiter.should_rate_limit_allow(squeak): - raise Exception( - "Excedeed allowed number of squeaks per block.") - # Only allow uploaded squeak if decryption key included. - if not squeak.HasDecryptionKey(): - raise Exception( - "Uploaded squeak must contain decryption key.") - decryption_key = squeak.GetDecryptionKey() - squeak_entry = self.squeak_core.validate_squeak(squeak) - logger.info("Saving uploaded squeak: {}".format( - get_hash(squeak).hex() - )) - inserted_squeak_hash = self.squeak_db.insert_squeak( - squeak, squeak_entry.block_header) - logger.info("Unlocking uploaded squeak: {}".format( - get_hash(squeak).hex() - )) - self.squeak_db.set_squeak_decryption_key( - inserted_squeak_hash, decryption_key - ) - return inserted_squeak_hash + return self.save_squeak(squeak, require_decryption_key=True) def save_downloaded_squeak(self, squeak: CSqueak) -> bytes: + return self.save_squeak(squeak, require_decryption_key=False) + + def save_created_squeak(self, squeak: CSqueak) -> bytes: + return self.save_squeak(squeak, require_decryption_key=True) + + def save_squeak( + self, + squeak: CSqueak, + require_decryption_key: bool, + ) -> bytes: + # Check if squeak is valid. + squeak_entry = self.squeak_core.validate_squeak(squeak) + # Check if squeak has decryption key. + if require_decryption_key and not squeak.HasDecryptionKey(): + raise Exception( + "Squeak must contain decryption key.") + # Check if rate limit is violated. if not self.squeak_rate_limiter.should_rate_limit_allow(squeak): raise Exception( - "Excedeed allowed number of squeaks per block.") - squeak_entry = self.squeak_core.validate_squeak(squeak) + "Exceeded allowed number of squeaks per address per block.") + # Save the squeak. + logger.info("Saving squeak: {}".format( + get_hash(squeak).hex(), + )) inserted_squeak_hash = self.squeak_db.insert_squeak( squeak, squeak_entry.block_header) + # Unlock the squeak if decryption key exists. + if squeak.HasDecryptionKey(): + decryption_key = squeak.GetDecryptionKey() + logger.info("Unlocking squeak: {}".format( + get_hash(squeak).hex(), + )) + self.unlock_squeak( + inserted_squeak_hash, + decryption_key, + ) + # Return the squeak hash. return inserted_squeak_hash def get_squeak(self, squeak_hash: bytes, clear_decryption_key: bool = False): @@ -81,17 +92,10 @@ class SqueakController: squeak.ClearDecryptionKey() return squeak - def get_public_squeak(self, squeak_hash: bytes): + def get_squeak_without_decryption_key(self, squeak_hash: bytes): return self.get_squeak(squeak_hash, clear_decryption_key=True) def lookup_allowed_addresses(self, addresses: List[str]): - # TODO: Implement db_whitelist class - # following_profiles_from_addresses = self.squeak_db.get_following_profiles_from_addreses( - # addresses) - # return [ - # profile.address - # for profile in following_profiles_from_addresses - # ] followed_addresses = self.get_followed_addresses() return set(followed_addresses) & set(addresses) @@ -229,10 +233,10 @@ class SqueakController: squeak_profile = self.squeak_db.get_profile(profile_id) squeak_entry = self.squeak_core.make_squeak( squeak_profile, content_str, replyto_hash) - # return self.save_created_squeak(squeak_entry.squeak) - inserted_squeak_hash = self.squeak_db.insert_squeak( - squeak_entry.squeak, squeak_entry.block_header) - return inserted_squeak_hash + return self.save_created_squeak(squeak_entry.squeak) + # inserted_squeak_hash = self.squeak_db.insert_squeak( + # squeak_entry.squeak, squeak_entry.block_header) + # return inserted_squeak_hash def delete_squeak(self, squeak_hash: bytes): num_deleted_offers = self.squeak_db.delete_offers_for_squeak( diff --git a/squeaknode/db/squeak_db.py b/squeaknode/db/squeak_db.py index 9a4cb931..ffb6d9ec 100644 --- a/squeaknode/db/squeak_db.py +++ b/squeaknode/db/squeak_db.py @@ -122,13 +122,11 @@ class SqueakDb: def insert_squeak(self, squeak: CSqueak, block_header: CBlockHeader) -> bytes: """ Insert a new squeak. + TODO: Clear the decryption key from the serialized bytes + without modifying the passed squeak object. Return the hash (bytes) of the inserted squeak. """ - secret_key_hex = ( - squeak.GetDecryptionKey().hex() if squeak.HasDecryptionKey() else None - ) - squeak.ClearDecryptionKey() ins = self.squeaks.insert().values( hash=get_hash(squeak).hex(), squeak=squeak.serialize(), @@ -137,7 +135,7 @@ class SqueakDb: n_block_height=squeak.nBlockHeight, n_time=squeak.nTime, author_address=str(squeak.GetAddress()), - secret_key=secret_key_hex, + secret_key=None, block_header=block_header.serialize(), ) with self.get_connection() as connection: @@ -964,19 +962,26 @@ class SqueakDb: def _parse_squeak_entry(self, row) -> SqueakEntry: secret_key_column = row["secret_key"] - secret_key = bytes.fromhex( - secret_key_column) if secret_key_column else b"" + secret_key = ( + bytes.fromhex(secret_key_column) if + secret_key_column else b"" + ) squeak = CSqueak.deserialize(row["squeak"]) if secret_key: squeak.SetDecryptionKey(secret_key) block_header_column = row["block_header"] - block_header_bytes = bytes( - block_header_column) if block_header_column else None - block_header = ( - parse_block_header( - block_header_bytes) if block_header_bytes else None + block_header_bytes = ( + bytes(block_header_column) if + block_header_column else None + ) + block_header = ( + parse_block_header(block_header_bytes) if + block_header_bytes else None + ) + return SqueakEntry( + squeak=squeak, + block_header=block_header, ) - return SqueakEntry(squeak=squeak, block_header=block_header) def _parse_squeak_profile(self, row) -> SqueakProfile: private_key_column = row["private_key"] diff --git a/squeaknode/server/squeak_server_handler.py b/squeaknode/server/squeak_server_handler.py index 4fe71a0d..c15aa423 100644 --- a/squeaknode/server/squeak_server_handler.py +++ b/squeaknode/server/squeak_server_handler.py @@ -23,16 +23,17 @@ class SqueakServerHandler(object): squeak.nBlockHeight > block_range.max_block: raise Exception("Invalid block range for upload.") followed_addresses = self.squeak_controller.get_followed_addresses() - squeak_address = squeak.GetAddress() - squeak_address_str = str(squeak_address) - if squeak_address_str not in followed_addresses: - raise Exception("Invalid squeak address for upload.") - # Save the squeak + squeak_address = str(squeak.GetAddress()) + if squeak_address not in followed_addresses: + raise Exception("Squeak address not in followed list.") + # Save the uploaded squeak self.squeak_controller.save_uploaded_squeak(squeak) def handle_get_squeak(self, squeak_hash: bytes): logger.info("Handle get squeak by hash: {}".format(squeak_hash.hex())) - return self.squeak_controller.get_public_squeak(squeak_hash) + return self.squeak_controller.get_squeak_without_decryption_key( + squeak_hash, + ) def handle_lookup_squeaks_to_download(self, request): network = request.network diff --git a/squeaknode/sync/peer_connection.py b/squeaknode/sync/peer_connection.py index 79c965fc..095a2461 100644 --- a/squeaknode/sync/peer_connection.py +++ b/squeaknode/sync/peer_connection.py @@ -5,6 +5,7 @@ from typing import Optional from squeaknode.core.block_range import BlockRange from squeaknode.core.received_offer_with_peer import ReceivedOfferWithPeer +from squeaknode.core.squeak_controller import SqueakController from squeaknode.network.peer_client import PeerClient logger = logging.getLogger(__name__) @@ -14,7 +15,7 @@ class PeerConnection: def __init__( self, - squeak_controller, + squeak_controller: SqueakController, peer, timeout_s, ): @@ -140,14 +141,18 @@ class PeerConnection: squeak_address_str = str(squeak_address) if squeak_address_str not in followed_addresses: raise Exception("Invalid squeak address for download.") - self.squeak_controller.save_downloaded_squeak(squeak) + self.squeak_controller.save_downloaded_squeak( + squeak, + ) logger.info("Downloaded squeak {} from peer {}".format( squeak_hash.hex(), self.peer )) def _force_download_squeak(self, squeak_hash: bytes): squeak = self.peer_client.download_squeak(squeak_hash) - self.squeak_controller.save_downloaded_squeak(squeak) + self.squeak_controller.save_downloaded_squeak( + squeak, + ) logger.info("Force downloaded squeak {} from peer {}".format( squeak_hash.hex(), self.peer ))