async def closeConnection(self, connection: Connection, reason: Optional[str]=None) -> None:
'''
End and close an existing connection:
Acknowledge or inform the peer about EOF, then close.
'''
if connection and not connection.writer.transport.is_closing():
# Tell transport the reason
if reason:
await self.send_message(reason, connection)
try:
connection.state['closed'] = time.time()
if connection.reader.at_eof():
self.logger.debug('Transport closed by peer')
connection.reader.feed_eof()
await self.handlePeerDisconnect(connection)
else:
self.logger.debug(f'Transport closed by {self.name}')
if connection.writer.can_write_eof(): connection.writer.write_eof()
await asyncio.sleep(.1)
connection.writer.close()
await connection.writer.wait_closed()
except:
pass
finally:
# Make sure we set this in any case
if not connection.state['closed']:
connection.state['closed'] = time.time()
self.logger.debug('Transport writer closed')
def _assembleConnection(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> Connection:
'''
Template method: Assemble a connection object based on the info we get from the reader/writer.
Might be overridden by a subclass method to set different parameters
'''
return Connection(
socket=writer.get_extra_info('socket'),
sslctx=None,
sslobj=None,
pid=None,
uid=None,
gid=None,
cert=None,
peer_address=None,
host_address=None,
reader=reader,
writer=writer,
read_handlers=set(),
write_handlers=set(),
state={
'mode': self.mode or ConnectionType.PERSISTENT,
'created': time.time(),
'updated': time.time(),
'closed': None
}
)
def _assembleConnection(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter) -> Connection:
sock = writer.get_extra_info('socket')
credentials = sock.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED,
struct.calcsize('3i'))
pid, uid, gid = struct.unpack('3i', credentials)
return Connection(socket=sock,
sslctx=writer.get_extra_info('sslcontext')
if self.sslctx else None,
sslobj=writer.get_extra_info('ssl_object')
if self.sslctx else None,
pid=pid,
uid=uid,
gid=gid,
peer_cert=writer.get_extra_info('peercert')
if self.sslctx else None,
peer_address=None,
host_address=None,
reader=reader,
writer=writer,
read_handlers=set(),
write_handlers=set(),
state={
'mode': self.mode or ConnectionType.PERSISTENT,
'created': time.time(),
'updated': time.time(),
'closed': None
})
def _assembleConnection(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter) -> Connection:
sock = writer.get_extra_info('socket')
return Connection(socket=sock,
sslctx=writer.get_extra_info('sslcontext')
if self.sslctx else None,
sslobj=writer.get_extra_info('ssl_object')
if self.sslctx else None,
pid=os.getpid(),
uid=os.getuid(),
gid=os.getgid(),
peer_cert=writer.get_extra_info('peercert')
if self.sslctx else None,
peer_address=sock.getpeername(),
host_address=sock.getsockname(),
reader=reader,
writer=writer,
read_handlers=set(),
write_handlers=set(),
state={
'mode': self.mode or ConnectionType.PERSISTENT,
'created': time.time(),
'updated': time.time(),
'closed': None
})
async def _handleConnection(self, connection: Connection) -> None:
'''
Handle the connection and listen for incoming messages until we receive an EOF.
'''
try:
# Authenticate
auth = await self.authenticateConnection(connection)
if auth:
self.logger.debug(
'Client connection was successfully authenticated')
else:
await self.rejectConnection(connection,
'Could not authenticate')
return
# Read data depending on the connection type and handle it
raw = None
chunks = 0
chunksize = self.chunksize
while not connection.reader.at_eof():
try:
# Update the connection
connection.state['updated'] = time.time()
# Read up to the limit or as many chunks until the limit
if self.chunksize:
# Make sure we read as many chunks till reaching the limit
if self.limit:
chunksize = min(self.limit, self.chunksize)
consumed = chunks * self.chunksize
rest = self.limit - consumed
if self.limit == chunksize and chunks > 0:
self.logger.error(
f'Set message size limit of {self.limit} bytes is smaller than chunksize {chunksize}. Closing connection.'
)
break
if rest <= 0:
self.logger.debug(
f'Total message size limit of {self.limit} bytes reached. Closing connection.'
)
break
chunksize = chunksize if rest > chunksize else rest
raw = await connection.reader.read(chunksize)
chunks += 1
# Read line by line (Respecting set limit)
elif self.lines:
try:
raw = await connection.reader.readuntil(
separator=self.line_separator.encode())
except asyncio.LimitOverrunError as e:
raw = await connection.reader.read(
e.consumed + len(self.line_separator.encode()))
await self.handleLimitExceedance(
connection, raw, True)
continue
except asyncio.IncompleteReadError as e:
if len(e.partial) > 0:
raw = e.partial
else:
continue
except asyncio.CancelledError as e:
raise e
except Exception as e:
continue
# Default read (Respecting set limit)
else:
raw = await connection.reader.read(self.limit or -1)
# Handle the received message or message fragment
if raw and not len(raw) == 0:
message = Message(data=raw, sender=connection)
# Never launch another message handler while we're being shutdown (this task getting already cancelled)
if not self._shutdown and not connection.state[
'closed']:
reader = asyncio.ensure_future(
self.handleMessage(message), loop=self.loop)
reader.add_done_callback(
lambda task: connection.read_handlers.remove(
task) if connection.read_handlers and task
in connection.read_handlers else None)
connection.read_handlers.add(reader)
except asyncio.CancelledError:
return # We return as the connection closing is handled by self.close()
except Exception as e:
self.logger.error(f'Transport error ({e})')
except asyncio.CancelledError:
return # We return as the connection closing is handled by self.close()
except ConnectionError as e:
await self.handleConnectionFailure(connection)
return
except Exception as e:
self.logger.error(f'Transport error ({e})')
# Close the connection (It will be automatically removed from the pool)
await self.closeConnection(connection)