From fa1c9cb42a3c78d7ed0e7a6b89ca033791b29482 Mon Sep 17 00:00:00 2001
From: "Dustin C. Hatch" <dustin@hatch.name>
Date: Sat, 30 Apr 2022 15:53:57 -0500
Subject: [PATCH] svc: hud: Protect window switches with a lock

Any time we need to switch Firefox windows, we need to use a lock to
prevent multiple simultaneous requests.  If we do not, interleaved
Marionette commands may result in performing operations on the wrong
window.  For example, making two simultaneous requests for screenshots
is liable to return the wrong window for one of them.
---
 svc/src/hudctrl/hud.py | 91 +++++++++++++++++++++++-------------------
 1 file changed, 49 insertions(+), 42 deletions(-)

diff --git a/svc/src/hudctrl/hud.py b/svc/src/hudctrl/hud.py
index d76891c..4e8618c 100644
--- a/svc/src/hudctrl/hud.py
+++ b/svc/src/hudctrl/hud.py
@@ -35,19 +35,21 @@ class HUDService:
         self.windows: Dict[str, str] = {}
 
         self.urls_file = Path('urls.json')
+        self.lock = asyncio.Lock()
 
     async def get_screens(self) -> Dict[str, HUDScreen]:
         assert self.marionette
         screens = {}
-        for name, handle in self.windows.items():
-            await self.marionette.switch_to_window(handle)
-            title = await self.marionette.get_title()
-            url = await self.marionette.get_url()
-            rect = await self.marionette.get_window_rect()
-            screens[name] = HUDScreen(
-                handle=handle, title=title, url=url, dimensions=rect
-            )
-        return screens
+        async with self.lock:
+            for name, handle in self.windows.items():
+                await self.marionette.switch_to_window(handle)
+                title = await self.marionette.get_title()
+                url = await self.marionette.get_url()
+                rect = await self.marionette.get_window_rect()
+                screens[name] = HUDScreen(
+                    handle=handle, title=title, url=url, dimensions=rect
+                )
+            return screens
 
     async def initialize(self) -> None:
         assert self.marionette
@@ -56,36 +58,39 @@ class HUDService:
             raise NoMonitorConfig(
                 'Cannot initialize display: No monitor config supplied'
             )
-        log.info(
-            'Initializing display for %d monitors',
-            len(self.monitor_config.monitors),
-        )
-        window: Optional[str] = await self.marionette.new_window('window')
-        for handle in await self.marionette.get_window_handles():
-            if handle == window:
-                continue
-            await self.marionette.switch_to_window(handle)
-            await self.marionette.close_window()
-        tasks = []
-        for monitor in self.monitor_config.monitors:
-            try:
-                url = self.urls[monitor.name]
-            except KeyError:
-                continue
-            if window is None:
-                window = await self.marionette.new_window('window')
-            self.windows[monitor.name] = window
-            await self.marionette.switch_to_window(window)
-            await self.marionette.set_window_rect(
-                x=monitor.pos_x,
-                y=1,
+        async with self.lock:
+            log.info(
+                'Initializing display for %d monitors',
+                len(self.monitor_config.monitors),
             )
-            await self.marionette.fullscreen()
-            log.info('Screen %s: Opening URL %s', monitor.name, url)
-            tasks.append(asyncio.create_task(self.marionette.navigate(url)))
-            window = None
-        if tasks:
-            await asyncio.wait(tasks)
+            window: Optional[str] = await self.marionette.new_window('window')
+            for handle in await self.marionette.get_window_handles():
+                if handle == window:
+                    continue
+                await self.marionette.switch_to_window(handle)
+                await self.marionette.close_window()
+            tasks = []
+            for monitor in self.monitor_config.monitors:
+                try:
+                    url = self.urls[monitor.name]
+                except KeyError:
+                    continue
+                if window is None:
+                    window = await self.marionette.new_window('window')
+                self.windows[monitor.name] = window
+                await self.marionette.switch_to_window(window)
+                await self.marionette.set_window_rect(
+                    x=monitor.pos_x,
+                    y=1,
+                )
+                await self.marionette.fullscreen()
+                log.info('Screen %s: Opening URL %s', monitor.name, url)
+                tasks.append(
+                    asyncio.create_task(self.marionette.navigate(url))
+                )
+                window = None
+            if tasks:
+                await asyncio.wait(tasks)
 
     def load_urls(self) -> None:
         try:
@@ -113,8 +118,9 @@ class HUDService:
 
     async def refresh_screen(self, name: str) -> None:
         assert self.marionette
-        await self.marionette.switch_to_window(self.windows[name])
-        await self.marionette.refresh()
+        async with self.lock:
+            await self.marionette.switch_to_window(self.windows[name])
+            await self.marionette.refresh()
 
     async def set_display(self, host: str, port: int) -> None:
         if self.marionette:
@@ -138,6 +144,7 @@ class HUDService:
     async def take_screenshot(self, screen: str) -> bytes:
         assert self.marionette
         handle = self.windows[screen]
-        await self.marionette.switch_to_window(handle)
-        data = await self.marionette.take_screenshot()
+        async with self.lock:
+            await self.marionette.switch_to_window(handle)
+            data = await self.marionette.take_screenshot()
         return base64.b64decode(data)