Qubes/core-admin
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

api: keep track of established connections

Browse Source 
This will be needed to gracefuly cleanup them between individual tests.
This commit is contained in:
Marek Marczykowski-Górecki 2017-07-26 02:56:19 +02:00
parent 4be325a383
commit 074c705f77
No known key found for this signature in database
GPG Key ID: 063938BA42CFA724
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
qubes/api/__init__.py
View File

@ -200,6 +200,9 @@ class AbstractQubesAPI(object):
class QubesDaemonProtocol(asyncio.Protocol): class QubesDaemonProtocol(asyncio.Protocol):
buffer_size = 65536 buffer_size = 65536
header = struct.Struct('Bx') header = struct.Struct('Bx')
# keep track of connections, to gracefully close them at server exit
# (including cleanup of integration test)
connections = set()
def __init__(self, handler, *args, app, debug=False, **kwargs): def __init__(self, handler, *args, app, debug=False, **kwargs):
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
@ -214,6 +217,7 @@ class QubesDaemonProtocol(asyncio.Protocol):
def connection_made(self, transport): def connection_made(self, transport):
self.transport = transport self.transport = transport
self.connections.add(self)
def connection_lost(self, exc): def connection_lost(self, exc):
self.untrusted_buffer.close() self.untrusted_buffer.close()
@ -221,6 +225,7 @@ class QubesDaemonProtocol(asyncio.Protocol):
if self.mgmt is not None: if self.mgmt is not None:
self.mgmt.cancel() self.mgmt.cancel()
self.transport = None self.transport = None
self.connections.remove(self)
def data_received(self, untrusted_data): # pylint: disable=arguments-differ def data_received(self, untrusted_data): # pylint: disable=arguments-differ
if self.len_untrusted_buffer + len(untrusted_data) > self.buffer_size: if self.len_untrusted_buffer + len(untrusted_data) > self.buffer_size: