Use single method to save squeak in controller (#834)

* Use single save_squeak method in controller

* Remove new lines from save_squeak method

* Add separate methods for save uploaded, downloaded, created squeak

* Simplify insert squeak method in database

* Add TODO comment for insert squeak method

* Improve code style of parse squeak entry method

* Add itest for failed upload without decryption key
This commit is contained in:
Jonathan Zernik 2021-02-13 04:13:11 -08:00 committed by GitHub
parent 5604cae842
commit 791b76bd1b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 97 additions and 61 deletions

View file

@ -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):

View file

@ -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(

View file

@ -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"]

View file

@ -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

View file

@ -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
))