154 lines
6.1 KiB
Python
154 lines
6.1 KiB
Python
import asyncio
|
|
import dataclasses
|
|
import unittest
|
|
from typing import *
|
|
|
|
from hippolyzer.lib.base.datatypes import UUID
|
|
from hippolyzer.lib.base.message.message import Block, Message
|
|
from hippolyzer.lib.base.message.message_handler import MessageHandler
|
|
from hippolyzer.lib.base.message.circuit import ConnectionHolder
|
|
from hippolyzer.lib.base.templates import (
|
|
AssetType,
|
|
EstateAssetType,
|
|
TransferSourceType,
|
|
TransferRequestParamsSimEstate,
|
|
TransferChannelType,
|
|
TransferTargetType,
|
|
TransferStatus,
|
|
)
|
|
from hippolyzer.lib.proxy.circuit import ProxiedCircuit
|
|
from hippolyzer.lib.base.network.transport import Direction
|
|
from hippolyzer.lib.base.transfer_manager import TransferManager, Transfer
|
|
from hippolyzer.lib.base.xfer_manager import XferManager
|
|
|
|
|
|
class MockHandlingCircuit(ProxiedCircuit):
|
|
def __init__(self, handler: MessageHandler[Message, str]):
|
|
super().__init__(("127.0.0.1", 1), ("127.0.0.1", 2), None)
|
|
self.handler = handler
|
|
|
|
def _send_prepared_message(self, message: Message, transport=None):
|
|
loop = asyncio.get_event_loop_policy().get_event_loop()
|
|
loop.call_soon(self.handler.handle, message)
|
|
|
|
|
|
class MockConnectionHolder(ConnectionHolder):
|
|
def __init__(self, circuit, message_handler):
|
|
self.circuit = circuit
|
|
self.message_handler = message_handler
|
|
|
|
|
|
class BaseTransferTests(unittest.IsolatedAsyncioTestCase):
|
|
SMALL_PAYLOAD = b"foobar"
|
|
LARGE_PAYLOAD = b"foobar" * 500
|
|
|
|
def setUp(self) -> None:
|
|
self.server_message_handler: MessageHandler[Message, str] = MessageHandler()
|
|
self.client_message_handler: MessageHandler[Message, str] = MessageHandler()
|
|
# The client side should send messages to the server side's message handler
|
|
# and vice-versa
|
|
self.client_circuit = MockHandlingCircuit(self.server_message_handler)
|
|
self.server_circuit = MockHandlingCircuit(self.client_message_handler)
|
|
self.server_connection = MockConnectionHolder(self.server_circuit, self.server_message_handler)
|
|
self.client_connection = MockConnectionHolder(self.client_circuit, self.client_message_handler)
|
|
|
|
|
|
class XferManagerTests(BaseTransferTests):
|
|
def setUp(self) -> None:
|
|
super().setUp()
|
|
self.secure_session_id = UUID.random()
|
|
self.xfer_manager = XferManager(self.client_connection, self.secure_session_id)
|
|
self.received_bytes: Optional[bytes] = None
|
|
|
|
async def _handle_vfile_upload(self):
|
|
msg = await self.server_message_handler.wait_for(('AssetUploadRequest',), timeout=0.01)
|
|
asset_block = msg["AssetBlock"]
|
|
transaction_id = asset_block["TransactionID"]
|
|
asset_id = UUID.combine(transaction_id, self.secure_session_id)
|
|
if asset_block["AssetData"]:
|
|
self.received_bytes = asset_block["AssetData"]
|
|
else:
|
|
manager = XferManager(self.server_connection)
|
|
xfer = await manager.request(vfile_id=asset_id, vfile_type=AssetType.BODYPART)
|
|
self.received_bytes = xfer.reassemble_chunks()
|
|
self.server_circuit.send(Message(
|
|
"AssetUploadComplete",
|
|
Block("AssetBlock", UUID=asset_id, Type=asset_block["Type"], Success=True),
|
|
direction=Direction.IN,
|
|
))
|
|
|
|
async def test_small_xfer_upload(self):
|
|
asyncio.create_task(self._handle_vfile_upload())
|
|
await asyncio.wait_for(self.xfer_manager.upload_asset(
|
|
AssetType.BODYPART, self.SMALL_PAYLOAD
|
|
), timeout=0.1)
|
|
self.assertEqual(self.received_bytes, self.SMALL_PAYLOAD)
|
|
|
|
async def test_large_xfer_upload(self):
|
|
# Larger payloads take a different path
|
|
asyncio.create_task(self._handle_vfile_upload())
|
|
await asyncio.wait_for(self.xfer_manager.upload_asset(
|
|
AssetType.BODYPART, self.LARGE_PAYLOAD
|
|
), timeout=0.1)
|
|
self.assertEqual(self.received_bytes, self.LARGE_PAYLOAD)
|
|
|
|
|
|
class TestTransferManager(BaseTransferTests):
|
|
def setUp(self) -> None:
|
|
super().setUp()
|
|
self.transfer_manager = TransferManager(
|
|
self.client_connection,
|
|
UUID.random(),
|
|
UUID.random(),
|
|
)
|
|
|
|
async def _handle_covenant_download(self):
|
|
msg = await self.server_message_handler.wait_for(('TransferRequest',), timeout=0.01)
|
|
self.assertEqual(TransferSourceType.SIM_ESTATE, msg["TransferInfo"]["SourceType"])
|
|
tid = msg["TransferInfo"]["TransferID"]
|
|
params: TransferRequestParamsSimEstate = msg["TransferInfo"][0].deserialize_var("Params")
|
|
self.assertEqual(EstateAssetType.COVENANT, params.EstateAssetType)
|
|
data = self.LARGE_PAYLOAD
|
|
|
|
self.server_circuit.send(Message(
|
|
'TransferInfo',
|
|
Block(
|
|
'TransferInfo',
|
|
TransferID=tid,
|
|
ChannelType=TransferChannelType.MISC,
|
|
TargetType_=TransferTargetType.UNKNOWN,
|
|
Status=TransferStatus.OK,
|
|
Size=len(self.LARGE_PAYLOAD),
|
|
Params_=dataclasses.asdict(params),
|
|
),
|
|
))
|
|
packet_num = 0
|
|
while True:
|
|
chunk = data[:1000]
|
|
data = data[1000:]
|
|
self.server_circuit.send(Message(
|
|
'TransferPacket',
|
|
Block(
|
|
'TransferData',
|
|
TransferID=tid,
|
|
ChannelType=TransferChannelType.MISC,
|
|
Packet=packet_num,
|
|
Status=TransferStatus.OK if data else TransferStatus.DONE,
|
|
Data=chunk,
|
|
)
|
|
))
|
|
if not data:
|
|
break
|
|
packet_num += 1
|
|
|
|
async def test_simple_transfer(self):
|
|
asyncio.create_task(self._handle_covenant_download())
|
|
transfer: Transfer = await asyncio.wait_for(self.transfer_manager.request(
|
|
source_type=TransferSourceType.SIM_ESTATE,
|
|
params=TransferRequestParamsSimEstate(
|
|
EstateAssetType=EstateAssetType.COVENANT,
|
|
),
|
|
), 0.1)
|
|
self.assertEqual(len(self.LARGE_PAYLOAD), transfer.expected_size)
|
|
self.assertEqual(self.LARGE_PAYLOAD, transfer.reassemble_chunks())
|