Fix "socket already in use", "frozen router" bugs in tests

mitre · Feb 8, 2024 · 7d6a4c8 · 7d6a4c8
1 parent e293ba0
commit 7d6a4c8
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 2 deletions.
diff --git a/app/api/packs/advanced.py b/app/api/packs/advanced.py
@@ -14,6 +14,7 @@ def __init__(self, services):
         self.rest_svc = services.get('rest_svc')
 
     async def enable(self):
+        self.app_svc.application.router._frozen = False
         self.app_svc.application.router.add_route('GET', '/advanced/sources', self._section_sources)
         self.app_svc.application.router.add_route('GET', '/advanced/objectives', self._section_objectives)
         self.app_svc.application.router.add_route('GET', '/advanced/planners', self._section_planners)

diff --git a/app/api/packs/campaign.py b/app/api/packs/campaign.py
@@ -16,6 +16,7 @@ def __init__(self, services):
         self.rest_svc = services.get('rest_svc')
 
     async def enable(self):
+        self.app_svc.application.router._frozen = False
         self.app_svc.application.router.add_route('GET', '/campaign/agents', self._section_agent)
         self.app_svc.application.router.add_route('GET', '/campaign/abilities', self._section_abilities)
         self.app_svc.application.router.add_route('GET', '/campaign/adversaries', self._section_profiles)

diff --git a/app/contacts/contact_dns.py b/app/contacts/contact_dns.py
@@ -18,12 +18,17 @@ def __init__(self, services):
         self.contact_svc = services.get('contact_svc')
         self.domain = self.get_config('app.contact.dns.domain')
         self.handler = Handler(self.domain, services, self.name)
+        self.transport = None
 
     async def start(self):
         loop = asyncio.get_event_loop()
         dns = self.get_config('app.contact.dns.socket')
         addr, port = dns.split(':')
-        await loop.create_datagram_endpoint(lambda: self.handler, local_addr=(addr, port))
+        self.transport, _ = await loop.create_datagram_endpoint(lambda: self.handler, local_addr=(addr, port))
+
+    async def stop(self):
+        if self.transport:
+            self.transport.close()
 
 
 class DnsPacket:

diff --git a/app/contacts/contact_udp.py b/app/contacts/contact_udp.py
@@ -13,12 +13,17 @@ def __init__(self, services):
         self.log = self.create_logger('contact_udp')
         self.contact_svc = services.get('contact_svc')
         self.handler = Handler(services)
+        self.transport = None
 
     async def start(self):
         loop = asyncio.get_event_loop()
         udp = self.get_config('app.contact.udp')
         addr, port = udp.split(':')
-        loop.create_task(loop.create_datagram_endpoint(lambda: self.handler, local_addr=(addr, port)))
+        self.transport, _ = await loop.create_task(loop.create_datagram_endpoint(lambda: self.handler, local_addr=(addr, port)))
+
+    async def stop(self):
+        if self.transport:
+            self.transport.close()
 
 
 class Handler(asyncio.DatagramProtocol):