Add some extra checks for size and existence to notification backend (#14398)

* Add some extra checks for size and existence to notification backend

* mypy
This commit is contained in:
Matt Hauff 2023-01-25 22:45:23 -07:00 committed by GitHub
parent 56f8ff0dc8
commit 1d4700dceb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 21 additions and 5 deletions

View File

@ -43,15 +43,16 @@ class NotificationManager:
return self
async def potentially_add_new_notification(self, coin_state: CoinState, parent_spend: CoinSpend) -> bool:
coin_name: bytes32 = coin_state.coin.name()
if (
coin_state.spent_height is None
or not self.wallet_state_manager.wallet_node.config.get("accept_notifications", False)
or self.wallet_state_manager.wallet_node.config.get("required_notification_amount", 0)
> coin_state.coin.amount
or await self.notification_store.notification_exists(coin_name)
):
return False
else:
coin_name: bytes32 = coin_state.coin.name()
memos: Dict[bytes32, List[bytes]] = compute_memos_for_spend(parent_spend)
coin_memos: List[bytes] = memos.get(coin_name, [])
if (
@ -59,6 +60,8 @@ class NotificationManager:
and construct_notification(bytes32(coin_memos[0]), uint64(coin_state.coin.amount)).get_tree_hash()
== coin_state.coin.puzzle_hash
):
if len(coin_memos[1]) > 10000: # 10KB
return False
await self.notification_store.add_notification(
Notification(
coin_state.coin.name(),

View File

@ -152,3 +152,11 @@ class NotificationStore:
# Delete from storage
cursor = await conn.execute("DELETE FROM notifications")
await cursor.close()
async def notification_exists(self, id: bytes32) -> bool:
async with self.db_wrapper.reader_no_transaction() as conn:
async with conn.execute("SELECT EXISTS (SELECT 1 from notifications WHERE coin_id=?)", (id,)) as cursor:
row = await cursor.fetchone()
assert row is not None
exists: bool = row[0] > 0
return exists

View File

@ -83,12 +83,12 @@ async def test_notifications(self_hostname: str, two_wallet_nodes: Any, trusted:
await server_0.start_client(PeerInfo(self_hostname, uint16(full_node_server._port)), None)
await server_1.start_client(PeerInfo(self_hostname, uint16(full_node_server._port)), None)
for i in range(0, 1):
for i in range(0, 2):
await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph_1))
await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph_token))
await time_out_assert(30, wallets_are_synced, True, [wallet_node_1, wallet_node_2], full_node_api)
funds_1 = sum([calculate_pool_reward(uint32(i)) + calculate_base_farmer_reward(uint32(i)) for i in range(1, 2)])
funds_1 = sum([calculate_pool_reward(uint32(i)) + calculate_base_farmer_reward(uint32(i)) for i in range(1, 3)])
funds_2 = 0
await time_out_assert(30, wallet_1.get_unconfirmed_balance, funds_1)
@ -97,7 +97,8 @@ async def test_notifications(self_hostname: str, two_wallet_nodes: Any, trusted:
notification_manager_1 = wsm_1.notification_manager
notification_manager_2 = wsm_2.notification_manager
for case in ("block all", "block too low", "allow", "allow_larger"):
for case in ("block all", "block too low", "allow", "allow_larger", "block_too_large"):
msg: bytes = bytes(case, "utf8")
if case == "block all":
wallet_node_2.config["required_notification_amount"] = 100
AMOUNT = uint64(100)
@ -113,13 +114,17 @@ async def test_notifications(self_hostname: str, two_wallet_nodes: Any, trusted:
else:
AMOUNT = uint64(750000000000)
FEE = uint64(1)
elif case == "block_too_large":
msg = bytes([0] * 10001)
AMOUNT = uint64(750000000000)
FEE = uint64(0)
peak = full_node_api.full_node.blockchain.get_peak()
assert peak is not None
if case == "allow":
allow_height = peak.height + 1
if case == "allow_larger":
allow_larger_height = peak.height + 1
tx = await notification_manager_1.send_new_notification(ph_2, bytes(case, "utf8"), AMOUNT, fee=FEE)
tx = await notification_manager_1.send_new_notification(ph_2, msg, AMOUNT, fee=FEE)
await wsm_1.add_pending_transaction(tx)
await time_out_assert_not_none(
5,