perf(cert_manager): slightly faster self signed generation

stapler/cert_manager.py

@@ -106,7 +106,7 @@ class CertManager:
                     "req",
                     "-new",
                     "-newkey",
-                    "rsa:4096",
+                    "rsa:2048",
                     "-days",
                     str(self.SELF_SIGNED_DAYS),
                     "-nodes",
@@ -192,7 +192,7 @@ class CertManager:
         _: ssl.SSLSocket,
         /,
     ) -> None | int:
-        if host is None:
+        if host is None or is_valid_host(host):
             return None
         self.logger.debug("servername callback: %s", host)
         if not self.exists(host) and not self.create_or_update(host):

stapler/data_dir.py

@@ -99,7 +99,8 @@ class DataDir:
             self.logger.debug("Deleted %s", target_path)
 
     def empty(self, path: str) -> None:
-        self.remove(path)
+        if self.exists(path):
+            self.remove(path)
         target_path = self.root_path / path
         target_path.mkdir()
         self.logger.debug("Created empty %s", target_path)

stapler/handlers.py

@@ -252,7 +252,10 @@ class BaseHandler(abc.ABC, http.server.BaseHTTPRequestHandler):
         return self.__in_size
 
     def _get_length(self) -> int:
-        return int(self._get_header("Content-Length", "0"))
+        try:
+            return max(0, int(self._get_header("Content-Length", "0")))
+        except ValueError:
+            return 0
 
     def _get_header(self, key: str, default_value: str = "") -> str:
         if self._has_header(key):
@@ -369,7 +372,7 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
     @property
     def target_redirect(self) -> str:
         if self.__target_redirect is None:
-            self.__target_redirect = self._get_header(self.REDIRECT_HEADER).lower()
+            self.__target_redirect = self._get_header(self.REDIRECT_HEADER)
         return self.__target_redirect
 
     @property
@@ -379,7 +382,7 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
     @property
     def target_proxy(self) -> str:
         if self.__target_proxy is None:
-            self.__target_proxy = self._get_header(self.PROXY_HEADER).lower()
+            self.__target_proxy = self._get_header(self.PROXY_HEADER)
         return self.__target_proxy
 
     @property
@@ -389,7 +392,7 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
     @property
     def target_spa(self) -> str:
         if self.__target_spa is None:
-            self.__target_spa = self._get_header(self.SPA_HEADER).lower()
+            self.__target_spa = self._get_header(self.SPA_HEADER)
         return self.__target_spa
 
     @property
@@ -493,16 +496,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 +521,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 +533,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:
@@ -536,7 +545,7 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
         return True
 
     def _proxy_or_redirect(self) -> bool:
-        if self.has_token or self.path.startswith(self.CERTBOT_CHALLENGE_PATH):
+        if self.has_token or self._is_certbot_challenge(self.path):
             return False
         if (page := self.__get_page(self.path)) is None:
             return False
@@ -556,9 +565,14 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
         """Disable default directory listing."""
         self.send_error(http.HTTPStatus.NOT_FOUND, "File not found")
 
+    def _is_certbot_challenge(self, path: str) -> bool:
+        return path.startswith(self.CERTBOT_CHALLENGE_PATH) and pathlib.Path(
+            self.certbot_www + path
+        ).resolve().is_relative_to(self.certbot_www)
+
     @typing.override
     def translate_path(self, path: str) -> str:
-        if path.startswith(self.CERTBOT_CHALLENGE_PATH):
+        if self._is_certbot_challenge(path):
             return self.certbot_www + path
         page = self.__get_page(path)
         if page is None:
@@ -655,7 +669,7 @@ class UpgradeHandler(RequestHandler):
 
     def do_GET(self) -> None:
         with self.handle_errors():
-            if self.path.startswith(self.CERTBOT_CHALLENGE_PATH):
+            if self._is_certbot_challenge(self.path):
                 super().do_GET()
                 self.close_connection = True
             else:

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/params.py

@@ -40,7 +40,7 @@ def __get_env_str(var: str, default: str) -> str:
 
 def __get_env_int(var: str, default: int) -> int:
     value = __get_env_str(var, str(default))
-    if value.isnumeric():
+    if value.isdecimal():
         return int(value)
     return default
 

stapler/registry.py

@@ -10,6 +10,7 @@ if typing.TYPE_CHECKING:
 
 class Registry:
     __slots__ = [
+        "_host_pages",
         "data_dir",
         "logger",
         "pages",
@@ -26,14 +27,24 @@ class Registry:
         self.logger: logging.Logger = logging.getLogger(self.__class__.__name__)
         self.pages: dict[str, Page] = {}
         self.data_dir = DataDir(params.data_dir)
+        self._host_pages: dict[str, Page] | None = None
+
+    @property
+    def host_pages(self) -> dict[str, Page]:
+        if self._host_pages is None:
+            self._host_pages = {
+                p.host: p for p in self.pages.values() if p.host is not None
+            }
+        return self._host_pages
 
     def load_pages(self) -> None:
         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]
+        return list(self.host_pages.keys())
 
     def add(self, path: str) -> None:
         host = self.data_dir.get_file(path, self.HOST_FILE)
@@ -47,7 +58,9 @@ 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._host_pages = None
         self.logger.info("Updated %s", self.pages[path])
 
     def set_host(self, path: str, host: str) -> None:
@@ -57,6 +70,7 @@ class Registry:
             self.data_dir.set_file(path, self.HOST_FILE, host)
             self.data_dir.remove_file(path, self.HOST_ONLY_FILE)
             self.pages[path].host = host
+            self._host_pages = None
             self.logger.debug("Updated %s", self.pages[path])
 
     def set_host_only(self, path: str, host: str) -> None:
@@ -67,6 +81,7 @@ class Registry:
             self.data_dir.remove_file(path, self.HOST_FILE)
             self.pages[path].host = host
             self.pages[path].host_only = True
+            self._host_pages = None
             self.logger.debug("Updated %s", self.pages[path])
 
     def set_token_hash(self, path: str, token_hash: str) -> None:
@@ -77,19 +92,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 +118,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._host_pages = None
             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:
-                return p
+        if host in self.host_pages and self.host_pages[host].ready:
+            return self.host_pages[host]
         return None

stapler/token_manager.py

@@ -15,6 +15,7 @@ class TokenManager:
     __slots__ = [
         "last_file_change",
         "logger",
+        "pbkdf2_iterations",
         "registry",
         "token_hashes",
         "token_salt",
@@ -23,11 +24,14 @@ class TokenManager:
 
     FILE = ".tokens"
 
-    def __init__(self, params: Parameters, registry: Registry) -> None:
+    def __init__(
+        self, params: Parameters, registry: Registry, pbkdf2_iterations: int = 500_000
+    ) -> None:
         self.logger: logging.Logger = logging.getLogger(self.__class__.__name__)
-        self.token_salt: str = params.token_salt
+        self.token_salt: bytes = params.token_salt.encode()
         self.tokens_file: pathlib.Path = pathlib.Path(params.data_dir) / self.FILE
         self.registry: Registry = registry
+        self.pbkdf2_iterations: int = pbkdf2_iterations
         self.token_hashes: list[str] = []
         self.last_file_change: int | float = 0
 
@@ -63,17 +67,18 @@ class TokenManager:
     def detect_file_change(self) -> bool:
         if (
             self.tokens_file.exists()
-            and self.tokens_file.stat().st_mtime != self.last_file_change
+            and (file_change := self.tokens_file.stat().st_mtime)
+            != self.last_file_change
         ):
             self.logger.debug("Detected change: %s", self.tokens_file)
-            self.last_file_change = self.tokens_file.stat().st_mtime
+            self.last_file_change = file_change
             return True
         return False
 
     def __hash_token(self, token: str) -> str:
-        return hashlib.sha512(
-            (self.token_salt + token).encode(), usedforsecurity=True
-        ).hexdigest()
+        return hashlib.pbkdf2_hmac(
+            "sha256", token.encode(), self.token_salt, self.pbkdf2_iterations
+        ).hex()
 
     def __load_hashes(self) -> list[str]:
         if self.tokens_file.is_file():

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")

tests/test_token_manager.py

@@ -11,9 +11,9 @@ from . import BaseTestCase
 
 
 class TestTokenManager(BaseTestCase):
-    EMPTY_SALT_HASH = "a04ca803c9fd73c21b721ece14b8b30cd3d9ca1bff752904a46982b881e152d0cdaa463a32e6bce71408de611953bc304ca8000d40d4b06b3f2a70769f69fecc"
-    SALT_HASH = "a5f2d8785eb4f064eae60f94e6025f93be32c2c93d2bbd73a982ee5c7ebcc484536487a4f60cfdfcb9ba72da7cebe0ce11afa91f191272e51d8c14be6874824b"
-    SECRET_HASH = "9901847ff8c76bd5fb473b7bd2e4f4ddd110332a52a888fd69deb276613885ddf382e5cf1210ed0decdb8010ae3994331a9e0639c3ca7e9e8b110dd50978ce76"  # noqa: S105
+    EMPTY_SALT_HASH = "5f88941ac5e26c430d97411ac1103af7a35c753f14aec088fbf34801c099135a"
+    SALT_HASH = "d71b1f52657c77d00b2a8c59b8d12d13c1c1bb2bcfbb85d2a9b804c36ad57a70"
+    SECRET_HASH = "38df428b309308e48c3687e7f90bda0e9cf253568c21ec754a0e076ab4ab6423"  # noqa: S105
 
     @typing.override
     def setUp(self) -> None:
@@ -21,6 +21,7 @@ class TestTokenManager(BaseTestCase):
         self.token_manager = TokenManager(
             Parameters(data_dir=self.get_tmp_dir(), token_salt="salt"),  # noqa: S106
             self.registry,
+            pbkdf2_iterations=1,
         )
         self.token_manager.logger = unittest.mock.Mock(logging.Logger)
         self.tmp_tokens_file = self.tmp_path / TokenManager.FILE
@@ -34,7 +35,7 @@ class TestTokenManager(BaseTestCase):
         self.assertListEqual(self.token_manager.token_hashes, [])
 
     def test_init_weak_salt(self) -> None:
-        self.token_manager.token_salt = ""
+        self.token_manager.token_salt = b""
         self.seal_mocks()
         self.token_manager.init()
         self.assert_file_content(