Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

mypy annotation for test_dos.py #19036

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

mypy annotation for test_dos.py #19036

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
135a5ed
remove unused function, get_block_path()
arvidn Dec 12, 2024
e92225e
add type annotation to test_dos.py
arvidn Dec 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 0 additions & 10 deletions chia/_tests/core/full_node/test_full_node.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,16 +129,6 @@ async def new_transaction_requested(incoming, new_spend):
return False


async def get_block_path(full_node: FullNodeAPI):
blocks_list = [await full_node.full_node.blockchain.get_full_peak()]
assert blocks_list[0] is not None
while blocks_list[0].height != 0:
b = await full_node.full_node.block_store.get_full_block(blocks_list[0].prev_header_hash)
assert b is not None
blocks_list.insert(0, b)
return blocks_list


@pytest.mark.anyio
async def test_sync_no_farmer(
setup_two_nodes_and_wallet,
Expand Down
11 changes: 0 additions & 11 deletions chia/_tests/core/full_node/test_performance.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@
from chia._tests.util.time_out_assert import time_out_assert
from chia.consensus.block_record import BlockRecord
from chia.consensus.pot_iterations import is_overflow_block
from chia.full_node.full_node_api import FullNodeAPI
from chia.protocols import full_node_protocol as fnp
from chia.types.condition_opcodes import ConditionOpcode
from chia.types.condition_with_args import ConditionWithArgs
Expand All @@ -23,16 +22,6 @@
log = logging.getLogger(__name__)


async def get_block_path(full_node: FullNodeAPI):
blocks_list = [await full_node.full_node.blockchain.get_full_peak()]
assert blocks_list[0] is not None
while blocks_list[0].height != 0:
b = await full_node.full_node.block_store.get_full_block(blocks_list[0].prev_header_hash)
assert b is not None
blocks_list.insert(0, b)
return blocks_list


class TestPerformance:
@pytest.mark.anyio
async def test_full_block_performance(
Expand Down
64 changes: 36 additions & 28 deletions chia/_tests/core/server/test_dos.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,13 @@

import pytest
from aiohttp import ClientSession, ClientTimeout, WSCloseCode, WSMessage, WSMsgType, WSServerHandshakeError
from chia_rs.sized_bytes import bytes32

import chia.server.server
from chia._tests.util.time_out_assert import time_out_assert
from chia.full_node.full_node_api import FullNodeAPI
from chia.protocols import full_node_protocol
from chia.protocols.protocol_message_types import ProtocolMessageTypes
from chia.protocols.shared_protocol import Handshake
from chia.protocols.shared_protocol import Capability, Handshake
from chia.server.outbound_message import Message, make_msg
from chia.server.rate_limits import RateLimiter
from chia.server.server import ChiaServer
Expand All @@ -22,7 +22,7 @@
from chia.simulator.full_node_simulator import FullNodeSimulator
from chia.types.peer_info import PeerInfo
from chia.util.errors import Err
from chia.util.ints import uint64
from chia.util.ints import uint8, uint16, uint64
from chia.util.timing import adjusted_timeout
from chia.wallet.wallet_node import WalletNode

Expand All @@ -33,18 +33,10 @@ def not_localhost(host: str) -> bool:
return False


async def get_block_path(full_node: FullNodeAPI):
blocks_list = [await full_node.full_node.blockchain.get_full_peak()]
assert blocks_list[0] is not None
while blocks_list[0].height != 0:
b = await full_node.full_node.block_store.get_full_block(blocks_list[0].prev_header_hash)
assert b is not None
blocks_list.insert(0, b)
return blocks_list


class FakeRateLimiter:
def process_msg_and_check(self, msg, capa, capb) -> Optional[str]:
def process_msg_and_check(
self, message: Message, our_capabilities: list[Capability], peer_capabilities: list[Capability]
) -> Optional[str]:
return None


Expand Down Expand Up @@ -142,7 +134,11 @@ async def test_bad_handshake_and_ban(
await ws.close()

@pytest.mark.anyio
async def test_invalid_protocol_handshake(self, setup_two_nodes_fixture, self_hostname):
async def test_invalid_protocol_handshake(
self,
setup_two_nodes_fixture: tuple[list[FullNodeSimulator], list[tuple[WalletNode, ChiaServer]], BlockTools],
self_hostname: str,
) -> None:
nodes, _, _ = setup_two_nodes_fixture
server_1 = nodes[0].full_node.server
server_2 = nodes[1].full_node.server
Expand All @@ -159,8 +155,8 @@ async def test_invalid_protocol_handshake(self, setup_two_nodes_fixture, self_ho
)

# Construct an otherwise valid handshake message
handshake: Handshake = Handshake("test", "0.0.32", "1.0.0.0", 3456, 1, [(1, "1")])
outbound_handshake: Message = Message(2, None, bytes(handshake)) # 2 is an invalid ProtocolType
handshake: Handshake = Handshake("test", "0.0.32", "1.0.0.0", uint16(3456), uint8(1), [(uint16(1), "1")])
outbound_handshake: Message = Message(uint8(2), None, bytes(handshake)) # 2 is an invalid ProtocolType
await ws.send_bytes(bytes(outbound_handshake))

response: WSMessage = await ws.receive()
Expand All @@ -173,7 +169,11 @@ async def test_invalid_protocol_handshake(self, setup_two_nodes_fixture, self_ho
await asyncio.sleep(1) # give some time for cleanup to work

@pytest.mark.anyio
async def test_spam_tx(self, setup_two_nodes_fixture, self_hostname):
async def test_spam_tx(
self,
setup_two_nodes_fixture: tuple[list[FullNodeSimulator], list[tuple[WalletNode, ChiaServer]], BlockTools],
self_hostname: str,
) -> None:
nodes, _, _ = setup_two_nodes_fixture
_full_node_1, full_node_2 = nodes
server_1 = nodes[0].full_node.server
Expand All @@ -191,7 +191,7 @@ async def test_spam_tx(self, setup_two_nodes_fixture, self_hostname):

new_tx_message = make_msg(
ProtocolMessageTypes.new_transaction,
full_node_protocol.NewTransaction(bytes([9] * 32), uint64(0), uint64(0)),
full_node_protocol.NewTransaction(bytes32([9] * 32), uint64(0), uint64(0)),
)
for i in range(4000):
await ws_con._send_message(new_tx_message)
Expand All @@ -215,20 +215,24 @@ async def test_spam_tx(self, setup_two_nodes_fixture, self_hostname):
await asyncio.sleep(0)
await asyncio.sleep(1)

def is_closed():
def is_closed() -> bool:
return ws_con.closed

await time_out_assert(15, is_closed)

assert ws_con.closed

def is_banned():
def is_banned() -> bool:
return "1.2.3.4" in server_2.banned_peers

await time_out_assert(15, is_banned)

@pytest.mark.anyio
async def test_spam_message_non_tx(self, setup_two_nodes_fixture, self_hostname):
async def test_spam_message_non_tx(
self,
setup_two_nodes_fixture: tuple[list[FullNodeSimulator], list[tuple[WalletNode, ChiaServer]], BlockTools],
self_hostname: str,
) -> None:
nodes, _, _ = setup_two_nodes_fixture
_full_node_1, full_node_2 = nodes
server_1 = nodes[0].full_node.server
Expand All @@ -244,7 +248,7 @@ async def test_spam_message_non_tx(self, setup_two_nodes_fixture, self_hostname)
ws_con.peer_info = PeerInfo("1.2.3.4", ws_con.peer_info.port)
ws_con_2.peer_info = PeerInfo("1.2.3.4", ws_con_2.peer_info.port)

def is_closed():
def is_closed() -> bool:
return ws_con.closed

new_message = make_msg(
Expand All @@ -271,13 +275,17 @@ def is_closed():
await time_out_assert(15, is_closed)

# Banned
def is_banned():
def is_banned() -> bool:
return "1.2.3.4" in server_2.banned_peers

await time_out_assert(15, is_banned)

@pytest.mark.anyio
async def test_spam_message_too_large(self, setup_two_nodes_fixture, self_hostname):
async def test_spam_message_too_large(
self,
setup_two_nodes_fixture: tuple[list[FullNodeSimulator], list[tuple[WalletNode, ChiaServer]], BlockTools],
self_hostname: str,
) -> None:
nodes, _, _ = setup_two_nodes_fixture
_full_node_1, full_node_2 = nodes
server_1 = nodes[0].full_node.server
Expand All @@ -293,7 +301,7 @@ async def test_spam_message_too_large(self, setup_two_nodes_fixture, self_hostna
ws_con.peer_info = PeerInfo("1.2.3.4", ws_con.peer_info.port)
ws_con_2.peer_info = PeerInfo("1.2.3.4", ws_con_2.peer_info.port)

def is_closed():
def is_closed() -> bool:
return ws_con.closed

new_message = make_msg(
Expand All @@ -307,13 +315,13 @@ def is_closed():
assert not ws_con.closed

# Remove outbound rate limiter to test inbound limits
ws_con.outbound_rate_limiter = FakeRateLimiter()
ws_con.outbound_rate_limiter = FakeRateLimiter() # type: ignore[assignment]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we do without this type ignore?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not without coming up with a different solution to make things unlimited

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sent #19038 to address this.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure it's an improvement. it just hides the type-unsafety


await ws_con._send_message(new_message)
await time_out_assert(15, is_closed)

# Banned
def is_banned():
def is_banned() -> bool:
return "1.2.3.4" in server_2.banned_peers

await time_out_assert(15, is_banned)
1 change: 0 additions & 1 deletion mypy-exclusions.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,6 @@ chia._tests.core.full_node.test_full_node
chia._tests.core.full_node.test_node_load
chia._tests.core.full_node.test_performance
chia._tests.core.full_node.test_transactions
chia._tests.core.server.test_dos
chia._tests.core.server.test_rate_limits
chia._tests.core.ssl.test_ssl
chia._tests.core.test_crawler_rpc
Expand Down
Loading