async def readuntil(self, delim):
try:
return await self.reader.readuntil(delim)
except (IncompleteReadError, LimitOverrunError):
raise MalformedDataError(f"Failed to find {delim} in read data")
except ConnectionError as e:
raise MalformedDataError(f"Connection error: {short_exc(e)}")
async def readexactly(self, amount):
try:
return await self.reader.readexactly(amount)
except IncompleteReadError:
raise MalformedDataError(
f"Stream ended while reading exactly {amount}")
except ConnectionError as e:
raise MalformedDataError(f"Connection error: {short_exc(e)}")
async def _read_game_data(cls, connection):
try:
line = await connection.readuntil(b'\0')
line = line[:-1].decode()
game_id, game_name = line.split("/", 1)
i, n = int(game_id), game_name
if i < 0:
raise MalformedDataError("Negative game ID!")
return i, n
except (ValueError, UnicodeDecodeError):
raise MalformedDataError("Malformed connection header")
async def from_connection(cls, connection):
generator = cls._generate(cls.MAXLEN)
generator.send(None)
while True:
data = await connection.read(4096) # TODO - configure?
if not data:
raise MalformedDataError("Replay header ended prematurely")
try:
generator.send(data)
except ValueError as e:
raise MalformedDataError("Invalid replay header") from e
except StopIteration as v:
return v.value
async def handle_connection(self, header, connection):
with self._track_connection(connection):
if header.type == ConnectionHeader.Type.WRITER:
await self.merger.handle_connection(connection)
elif header.type == ConnectionHeader.Type.READER:
await self.sender.handle_connection(connection)
else:
raise MalformedDataError("Invalid connection type")
async def _read_type(cls, connection):
prefix = await connection.readexactly(2)
if prefix == b"P/":
return cls.Type.WRITER
elif prefix == b"G/":
return cls.Type.READER
else:
raise MalformedDataError(
f"Expected reader or writer prefix, got '{prefix}'")
async def write(self, data):
if self._closed:
return False
try:
self.writer.write(data)
await self.writer.drain()
except ConnectionError as e:
raise MalformedDataError("Connection error") from e
return True
async def handle_connection(self, header, connection):
with self._track_connection(connection):
logger.debug(f"{self} - new connection, {header}")
if header.type == ConnectionHeader.Type.WRITER:
await self.merger.handle_connection(connection)
elif header.type == ConnectionHeader.Type.READER:
await self.sender.handle_connection(connection)
else:
raise MalformedDataError("Invalid connection type")
logger.debug(f"{self} - connection over, {header}")
async def write(self, data):
if self._closed:
return False
try:
self.writer.write(data)
await self.writer.drain()
except ConnectionResetError:
return False
except (TimeoutError, ConnectionError, OSError) as e:
raise MalformedDataError(f"Connection error: {short_exc(e)}")
return True
async def _read_type(cls, connection):
try:
prefix = await connection.readexactly(2)
except MalformedDataError:
# Vast majority of these will be connections that entered lobby,
# but quit without starting the game. This also ignores very early
# connection errors and reads of length exactly 1; consider that a
# FIXME.
raise EmptyConnectionError
if prefix == b"P/":
return cls.Type.WRITER
elif prefix == b"G/":
return cls.Type.READER
else:
raise MalformedDataError(
f"Expected reader or writer prefix, got '{prefix}'")
async def from_connection(cls, connection):
gen = GeneratorData(cls.MAXLEN)
generator = cls._generate(gen)
generator.send(None)
while True:
try:
data = await connection.read(4096) # TODO - configure?
if not data:
raise ValueError("Replay header ended prematurely")
generator.send(data)
except ValueError as e:
raise MalformedDataError(
f"Invalid replay header while {gen.state}"
f" at {gen.position} bytes: {e}")
except StopIteration as v:
return v.value
async def read(self, size):
try:
data = await self.reader.read(size)
return data
except ConnectionError as e:
raise MalformedDataError(f"Connection error: {short_exc(e)}")
async def read(cls, connection, timeout):
try:
return await asyncio.wait_for(cls._do_read(connection), timeout)
except asyncio.TimeoutError:
raise MalformedDataError("Timed out while reading header")
async def _write_header(self, connection):
header = await self._stream.wait_for_header()
if header is None:
raise MalformedDataError("Malformed replay header")
await connection.write(header.data)
async def read(self, size):
try:
data = await self.reader.read(size)
return data
except ConnectionError as e:
raise MalformedDataError("Connection error") from e