From c0f82c97e2759b63f3009d8998da8a73f96d8eac Mon Sep 17 00:00:00 2001
From: Martin Pitt <martin@piware.de>
Date: Mon, 19 Sep 2022 11:30:02 +0200
Subject: [PATCH] appservice: Add initial status API

Track session status. Right after creation, it starts as `wait_target`.
After the target machine has connected (first connection to the /ws
path), it changes to `running`. We'll probably add more states in the
future.

Add /sessions/ID/status API to query it, so that e.g. the target machine
and the UI can poll this to know when to connect.

Fixes #28
---
 appservice/multiplexer.py | 25 ++++++++++++++++++++-----
 test/test_basic.py        | 21 +++++++++++++++++++++
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/appservice/multiplexer.py b/appservice/multiplexer.py
index cd0e4c0..4b498f8 100644
--- a/appservice/multiplexer.py
+++ b/appservice/multiplexer.py
@@ -17,6 +17,7 @@
 SESSION_INSTANCE_DOMAIN = os.getenv('SESSION_INSTANCE_DOMAIN', '')
 
 PODMAN_SOCKET = '/run/podman/podman.sock'
+# states: wait_target or running
 SESSIONS = {}
 REDIS = redis.Redis(host=os.environ['REDIS_SERVICE_HOST'], port=int(os.environ.get('REDIS_SERVICE_PORT', '6379')))
 
@@ -84,7 +85,6 @@ async def new_session_podman(self, sessionid):
         return status, content
 
     async def handle_session_new(self, request):
-        global SESSIONS
         sessionid = str(uuid.uuid4())
         assert sessionid not in SESSIONS
 
@@ -96,15 +96,19 @@ async def handle_session_new(self, request):
 
         if pod_status >= 200 and pod_status < 300:
             response = web.json_response({'id': sessionid})
-            SESSIONS[sessionid] = True
-            dumped_sessions = json.dumps(SESSIONS)
-            await REDIS.set('sessions', dumped_sessions)
-            await REDIS.publish('sessions', dumped_sessions)
+            await update_session(sessionid, 'wait_target')
         else:
             response = web.Response(status=pod_status, text=f'creating session container failed: {content}')
 
         return response
 
+    async def handle_session_status(self, request):
+        sessionid = request.match_info['sessionid']
+        try:
+            return web.Response(text=SESSIONS[sessionid])
+        except KeyError:
+            return web.HTTPNotFound(text='unknown session ID')
+
     async def handle_session_id(self, upstream_req):
         sessionid = upstream_req.match_info['sessionid']
         if sessionid not in SESSIONS:
@@ -115,6 +119,8 @@ async def handle_session_id(self, upstream_req):
             target_url = f'http://session-{sessionid}{SESSION_INSTANCE_DOMAIN}:9090{upstream_req.path_qs}'
         elif path == 'ws':
             target_url = f'http://session-{sessionid}{SESSION_INSTANCE_DOMAIN}:8080{upstream_req.path_qs}'
+            if SESSIONS[sessionid] == 'wait_target':
+                await update_session(sessionid, 'running')
         else:
             return web.HTTPNotFound(text=f'invalid session path prefix: {path}')
 
@@ -188,6 +194,14 @@ async def init_sessions():
     logger.debug('initial sessions: %s', SESSIONS)
 
 
+async def update_session(session_id, status):
+    global SESSIONS
+    SESSIONS[session_id] = status
+    dumped_sessions = json.dumps(SESSIONS)
+    await REDIS.set('sessions', dumped_sessions)
+    await REDIS.publish('sessions', dumped_sessions)
+
+
 async def watch_redis(channel):
     global SESSIONS
     while True:
@@ -217,6 +231,7 @@ async def main():
         app = web.Application()
         app.router.add_route('GET', f'{config.ROUTE_API}/ping', handlers.handle_ping)
         app.router.add_route('GET', f'{config.ROUTE_API}/sessions/new', handlers.handle_session_new)
+        app.router.add_route('GET', f'{config.ROUTE_API}/sessions/{{sessionid}}/status', handlers.handle_session_status)
         app.router.add_route('*', f'{config.ROUTE_WSS}/sessions/{{sessionid}}/{{path:.*}}', handlers.handle_session_id)
 
         runner = web.AppRunner(app, auto_decompress=False)
diff --git a/test/test_basic.py b/test/test_basic.py
index 3448dc0..47e40f2 100755
--- a/test/test_basic.py
+++ b/test/test_basic.py
@@ -103,6 +103,11 @@ def newSession(self):
         sessionid = json.load(response)['id']
         self.assertIsInstance(sessionid, str)
 
+        # inital status
+        response = self.request(f'{self.api_url}{config.ROUTE_API}/sessions/{sessionid}/status')
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.read(), b'wait_target')
+
         # API URL is on the container host's localhost; translate for the container DNS
         websocket_url = self.api_url.replace('localhost', 'host.containers.internal').replace('https:', 'wss:')
         podman = ['podman', 'run', '-d', '--pod', 'webconsoleapp',
@@ -113,6 +118,17 @@ def newSession(self):
 
         subprocess.check_call(podman + cmd)
 
+        # successful bridge connection updates status
+        for retry in range(10):
+            response = self.request(f'{self.api_url}{config.ROUTE_API}/sessions/{sessionid}/status')
+            self.assertEqual(response.status, 200)
+            status = response.read()
+            if status == b'running':
+                break
+            time.sleep(0.5)
+        else:
+            self.fail(f'session status was not updated to running, still at {status}')
+
         return sessionid
 
     def checkSession(self, sessionid):
@@ -142,6 +158,11 @@ def testSessions(self):
         # first session still works
         self.checkSession(s1)
 
+        # unknown session ID
+        with self.assertRaises(urllib.error.HTTPError) as cm:
+            self.request(f'{self.api_url}{config.ROUTE_API}/sessions/123unknown/status')
+        self.assertEqual(cm.exception.code, 404)
+
         # crash container for s2; use --time 0 once we have podman 4.0 everywhere
         subprocess.check_call(['podman', 'rm', '--force', f'session-{s2}'])
         # first session still works