fix(registry): race condition handled with page 'ready'

stapler/handlers.py

@@ -493,16 +493,20 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
                 "Archive too large",
             )
             return False
+        existing = self.registry.remove(path)
         try:
             file_bytes = io.BytesIO(self.rfile.read(self.in_size))
             self.data_dir.extract_tar_bytes(path, file_bytes)
         except tarfile.TarError:
             self.send_error(http.HTTPStatus.BAD_REQUEST, "Invalid tar archive")
+            if existing:
+                self.registry.add(path)  # restore path on error
             return False
         self.registry.add(path)
         self.token_manager.set_token(path, self.token)
         if self.has_target_spa:
             self.registry.set_spa(path, self.target_spa)
+        self.registry.mark_ready(path)
         return True
 
     def _update_redirect(self, path: str) -> bool:
@@ -514,6 +518,7 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
             return False
         self.registry.set_redirect(path, self.target_redirect)
         self.token_manager.set_token(path, self.token)
+        self.registry.mark_ready(path)
         return True
 
     def _update_proxy(self, path: str) -> bool:
@@ -525,6 +530,7 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
             return False
         self.registry.set_proxy(path, self.target_proxy)
         self.token_manager.set_token(path, self.token)
+        self.registry.mark_ready(path)
         return True
 
     def _update_remove(self, path: str) -> bool:

stapler/page.py

@@ -11,6 +11,7 @@ class Page:
     redirect: str | None = None
     proxy: str | None = None
     spa: str | None = None
+    ready: bool = True
 
     def __repr__(self) -> str:
         out = f"/{self.path}/"

stapler/registry.py

@@ -31,6 +31,7 @@ class Registry:
         self.pages = {}
         for path in self.data_dir.list_paths():
             self.add(path)
+            self.mark_ready(path)
 
     def get_hosts(self) -> list[str]:
         return [p.host for p in self.pages.values() if p.host is not None]
@@ -47,6 +48,7 @@ class Registry:
             redirect=self.data_dir.get_file(path, self.REDIRECT_FILE),
             proxy=self.data_dir.get_file(path, self.PROXY_FILE),
             spa=self.data_dir.get_file(path, self.SPA_FILE),
+            ready=False,
         )
         self.logger.info("Updated %s", self.pages[path])
 
@@ -77,19 +79,23 @@ class Registry:
 
     def set_redirect(self, path: str, redirect: str) -> None:
         if path not in self.pages or self.pages[path].redirect != redirect:
+            if path in self.pages:
+                self.pages[path].ready = False
             self.data_dir.empty(path)
             self.data_dir.set_file(path, self.REDIRECT_FILE, redirect)
             if path not in self.pages:
-                self.pages[path] = Page(path)
+                self.pages[path] = Page(path, ready=False)
             self.pages[path].redirect = redirect
             self.logger.debug("Updated %s", self.pages[path])
 
     def set_proxy(self, path: str, proxy: str) -> None:
         if path not in self.pages or self.pages[path].proxy != proxy:
+            if path in self.pages:
+                self.pages[path].ready = False
             self.data_dir.empty(path)
             self.data_dir.set_file(path, self.PROXY_FILE, proxy)
             if path not in self.pages:
-                self.pages[path] = Page(path)
+                self.pages[path] = Page(path, ready=False)
             self.pages[path].proxy = proxy
             self.logger.debug("Updated %s", self.pages[path])
 
@@ -99,19 +105,25 @@ class Registry:
             self.pages[path].spa = spa
             self.logger.debug("Updated %s", self.pages[path])
 
-    def remove(self, path: str) -> None:
+    def mark_ready(self, path: str) -> None:
+        if path in self.pages:
+            self.pages[path].ready = True
+
+    def remove(self, path: str) -> bool:
         if path in self.pages:
             page = self.pages[path]
             del self.pages[path]
             self.logger.info("Removed %s", page)
+            return True
+        return False
 
     def get_from_path(self, path: str) -> Page | None:
-        if path in self.pages:
+        if path in self.pages and self.pages[path].ready:
             return self.pages[path]
         return None
 
     def get_from_host(self, host: str) -> Page | None:
         for p in self.pages.values():
-            if p.host == host:
+            if p.host == host and p.ready:
                 return p
         return None

tests/test_handlers.py

@@ -436,6 +436,8 @@ class TestRequestHandler(BaseHandlerTestCase):
                 ["secret", "path"],
                 True,  # noqa: FBT003
             ),
+            self.mock_call(self.registry.remove, ["path"], True),  # noqa: FBT003
+            self.mock_call(self.registry.add, ["path"]),
             self.expects_error(
                 handler, http.HTTPStatus.BAD_REQUEST, "Invalid tar archive"
             ),
@@ -457,6 +459,7 @@ class TestRequestHandler(BaseHandlerTestCase):
                 ["secret", "path"],
                 True,  # noqa: FBT003
             ),
+            self.mock_call(self.registry.remove, ["path"], False),  # noqa: FBT003
             self.expects_error(handler, http.HTTPStatus.INTERNAL_SERVER_ERROR, ""),
             self.seal_mocks(),
         ):
@@ -476,8 +479,10 @@ class TestRequestHandler(BaseHandlerTestCase):
                 True,  # noqa: FBT003
             ),
             self.mock_call_unchecked(self.data_dir.extract_tar_bytes),
+            self.mock_call(self.registry.remove, ["path"], False),  # noqa: FBT003
             self.mock_call(self.registry.add, ["path"]),
             self.mock_call(self.token_manager.set_token, ["path", "secret"]),
+            self.mock_call(self.registry.mark_ready, ["path"]),
             self.mock_call(self.registry.get_from_path, ["path"]),
             self.expects_status_only(
                 handler, http.HTTPStatus.CREATED, "Resource updated"
@@ -501,9 +506,11 @@ class TestRequestHandler(BaseHandlerTestCase):
             ),
             self.mock_call(self.registry.get_from_host, ["example.com"], Page("path")),
             self.mock_call_unchecked(self.data_dir.extract_tar_bytes),
+            self.mock_call(self.registry.remove, ["path"], False),  # noqa: FBT003
             self.mock_call(self.registry.add, ["path"]),
             self.mock_call(self.token_manager.set_token, ["path", "secret"]),
             self.mock_call(self.registry.set_host, ["path", "example.com"]),
+            self.mock_call(self.registry.mark_ready, ["path"]),
             self.mock_call(self.registry.get_from_path, ["path"]),
             self.expects_status_only(
                 handler, http.HTTPStatus.CREATED, "Resource updated"
@@ -525,9 +532,11 @@ class TestRequestHandler(BaseHandlerTestCase):
                 True,  # noqa: FBT003
             ),
             self.mock_call_unchecked(self.data_dir.extract_tar_bytes),
+            self.mock_call(self.registry.remove, ["path"], False),  # noqa: FBT003
             self.mock_call(self.registry.add, ["path"]),
             self.mock_call(self.token_manager.set_token, ["path", "secret"]),
             self.mock_call(self.registry.set_spa, ["path", "index.html"]),
+            self.mock_call(self.registry.mark_ready, ["path"]),
             self.mock_call(self.registry.get_from_path, ["path"]),
             self.expects_status_only(
                 handler, http.HTTPStatus.CREATED, "Resource updated"
@@ -578,6 +587,7 @@ class TestRequestHandler(BaseHandlerTestCase):
             ),
             self.mock_call(self.registry.set_redirect, ["path", "https://example.com"]),
             self.mock_call(self.token_manager.set_token, ["path", "secret"]),
+            self.mock_call(self.registry.mark_ready, ["path"]),
             self.mock_call(self.registry.get_from_path, ["path"]),
             self.expects_status_only(
                 handler, http.HTTPStatus.CREATED, "Resource updated"
@@ -606,6 +616,7 @@ class TestRequestHandler(BaseHandlerTestCase):
             self.mock_call(self.registry.set_redirect, ["path", "https://example.com"]),
             self.mock_call(self.token_manager.set_token, ["path", "secret"]),
             self.mock_call(self.registry.set_host, ["path", "example.com"]),
+            self.mock_call(self.registry.mark_ready, ["path"]),
             self.mock_call(self.registry.get_from_path, ["path"]),
             self.expects_status_only(
                 handler, http.HTTPStatus.CREATED, "Resource updated"
@@ -634,6 +645,7 @@ class TestRequestHandler(BaseHandlerTestCase):
             self.mock_call(self.registry.set_redirect, ["path", "https://example.com"]),
             self.mock_call(self.token_manager.set_token, ["path", "secret"]),
             self.mock_call(self.registry.set_host_only, ["path", "example.com"]),
+            self.mock_call(self.registry.mark_ready, ["path"]),
             self.mock_call(self.registry.get_from_path, ["path"]),
             self.expects_status_only(
                 handler, http.HTTPStatus.CREATED, "Resource updated"
@@ -684,6 +696,7 @@ class TestRequestHandler(BaseHandlerTestCase):
             ),
             self.mock_call(self.registry.set_proxy, ["path", "https://example.com"]),
             self.mock_call(self.token_manager.set_token, ["path", "secret"]),
+            self.mock_call(self.registry.mark_ready, ["path"]),
             self.mock_call(self.registry.get_from_path, ["path"]),
             self.expects_status_only(
                 handler, http.HTTPStatus.CREATED, "Resource updated"
@@ -712,6 +725,7 @@ class TestRequestHandler(BaseHandlerTestCase):
             self.mock_call(self.registry.set_proxy, ["path", "https://example.com"]),
             self.mock_call(self.token_manager.set_token, ["path", "secret"]),
             self.mock_call(self.registry.set_host, ["path", "example.com"]),
+            self.mock_call(self.registry.mark_ready, ["path"]),
             self.mock_call(self.registry.get_from_path, ["path"]),
             self.expects_status_only(
                 handler, http.HTTPStatus.CREATED, "Resource updated"

tests/test_registry.py

@@ -185,6 +185,7 @@ class TestRegistry(BaseTestCase):
         self.assertEqual(
             self.registry.pages["test_1"].redirect, "https://new-example.com"
         )
+        assert not self.registry.pages["test_1"].ready
 
     def test_set_redirect_no_change(self) -> None:
         self.registry.pages["test_1"] = Page(
@@ -214,6 +215,7 @@ class TestRegistry(BaseTestCase):
         self.assertEqual(
             self.registry.pages["test_1"].redirect, "https://new-example.com"
         )
+        assert not self.registry.pages["test_1"].ready
 
     def test_set_proxy(self) -> None:
         self.registry.pages["test_1"] = Page(
@@ -233,6 +235,7 @@ class TestRegistry(BaseTestCase):
         ):
             self.registry.set_proxy("test_1", "https://new-example.com")
         self.assertEqual(self.registry.pages["test_1"].proxy, "https://new-example.com")
+        assert not self.registry.pages["test_1"].ready
 
     def test_set_proxy_no_change(self) -> None:
         self.registry.pages["test_1"] = Page(
@@ -260,6 +263,7 @@ class TestRegistry(BaseTestCase):
             self.registry.set_proxy("test_1", "https://new-example.com")
         self.assertIn("test_1", self.registry.pages)
         self.assertEqual(self.registry.pages["test_1"].proxy, "https://new-example.com")
+        assert not self.registry.pages["test_1"].ready
 
     def test_set_spa(self) -> None:
         self.registry.pages["test_1"] = Page(
@@ -298,9 +302,27 @@ class TestRegistry(BaseTestCase):
             "test_1",
         )
         self.seal_mocks()
-        self.registry.remove("test_1")
+        assert self.registry.remove("test_1")
         self.assertNotIn("test_1", self.registry.pages)
 
+    def test_remove_not_found(self) -> None:
+        self.seal_mocks()
+        assert not self.registry.remove("test_1")
+
+    def test_mark_ready(self) -> None:
+        self.registry.pages["test_1"] = Page("test_1", ready=False)
+        with (
+            self.seal_mocks(),
+        ):
+            self.registry.mark_ready("test_1")
+            assert self.registry.pages["test_1"].ready
+
+    def test_mark_ready_not_found(self) -> None:
+        with (
+            self.seal_mocks(),
+        ):
+            self.registry.mark_ready("test_1")
+
     def test_get_from_path(self) -> None:
         self.registry.pages["test_1"] = (
             target := Page(
@@ -313,6 +335,14 @@ class TestRegistry(BaseTestCase):
         self.seal_mocks()
         self.assertEqual(self.registry.get_from_path("test_1"), target)
 
+    def test_get_from_path_not_ready(self) -> None:
+        self.registry.pages["test_1"] = Page(
+            "test_1",
+            ready=False,
+        )
+        self.seal_mocks()
+        self.assertIsNone(self.registry.get_from_path("test_1"))
+
     def test_get_from_path_not_found(self) -> None:
         self.registry.pages["test_1"] = Page(
             "test_1",
@@ -329,6 +359,11 @@ class TestRegistry(BaseTestCase):
         self.seal_mocks()
         self.assertEqual(self.registry.get_from_host("host_1"), target)
 
+    def test_get_from_host_not_ready(self) -> None:
+        self.registry.pages["test_1"] = Page("test_1", host="host_1", ready=False)
+        self.seal_mocks()
+        self.assertIsNone(self.registry.get_from_host("host_1"))
+
     def test_get_from_host_not_found(self) -> None:
         self.registry.pages["test_1"] = Page("test_1", host="host_1")
         self.registry.pages["test_2"] = Page("test_2", host="host_2")