mirror of
https://github.com/yzernik/squeaknode.git
synced 2026-08-20 13:28:20 +02:00
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:
parent
5604cae842
commit
791b76bd1b
5 changed files with 97 additions and 61 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
))
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue