From c5f3e53d880e88e1ed3554d86834211db26db801 Mon Sep 17 00:00:00 2001
From: klemek <klemek.dev@proton.me>
Date: Wed, 3 Jun 2026 18:56:07 +0200
Subject: [PATCH] feat: stapler challenge for self discovery

---
 stapler/handlers.py    | 19 +++++++++++++++++--
 tests/test_handlers.py | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/stapler/handlers.py b/stapler/handlers.py
index 5ef5046..5fd08bd 100644
--- a/stapler/handlers.py
+++ b/stapler/handlers.py
@@ -289,6 +289,7 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
     protocol_version = "HTTP/1.1"
     server_version = "StaplerServer/" + PKG_VERSION
     CERTBOT_CHALLENGE_PATH = "/.well-known/acme-challenge"
+    STAPLER_CHALLENGE_PATH = "/.well-known/stapler"
     UPDATE_PATH_REGEX = re.compile(r"^\/([\w-]+)\/?$")
     GET_PATH_REGEX = re.compile(r"^\/([\w-]+)($|\/)")
     HOST_PART_REGEX = re.compile(r"^([a-z0-9]|[a-z0-9][a-z0-9-]{,61}[a-z0-9])$")
@@ -405,6 +406,8 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
             self._pre_log_request()
             if self._proxy_or_redirect():
                 return None
+            if self._is_stapler_challenge(self.path):
+                return self.send_status_only(http.HTTPStatus.NO_CONTENT)
             if self.path == "/" and self.host == self.default_host:
                 return self.send_basic_body(self.server_signature())
             super().do_HEAD()
@@ -417,6 +420,8 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
             self._pre_log_request()
             if self._proxy_or_redirect():
                 return None
+            if self._is_stapler_challenge(self.path):
+                return self.send_status_only(http.HTTPStatus.NO_CONTENT)
             if self.path == "/" and self.host == self.default_host:
                 return self.send_basic_body(self.server_signature())
             super().do_GET()
@@ -545,7 +550,11 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
         return True
 
     def _proxy_or_redirect(self) -> bool:
-        if self.has_token or self._is_certbot_challenge(self.path):
+        if (
+            self.has_token
+            or self._is_certbot_challenge(self.path)
+            or self._is_stapler_challenge(self.path)
+        ):
             return False
         if (page := self.__get_page(self.path)) is None:
             return False
@@ -570,6 +579,9 @@ class RequestHandler(http.server.SimpleHTTPRequestHandler, BaseHandler):
             self.certbot_www + path
         ).resolve().is_relative_to(self.certbot_www)
 
+    def _is_stapler_challenge(self, path: str) -> bool:
+        return path.startswith(self.STAPLER_CHALLENGE_PATH)
+
     @typing.override
     def translate_path(self, path: str) -> str:
         if self._is_certbot_challenge(path):
@@ -664,7 +676,10 @@ class UpgradeHandler(RequestHandler):
     def do_HEAD(self) -> None:
         with self.handle_errors():
             self._pre_log_request()
-            self.send_redirect(f"https://{self.host}{self.path}")
+            if self._is_stapler_challenge(self.path):
+                self.send_status_only(http.HTTPStatus.NO_CONTENT)
+            else:
+                self.send_redirect(f"https://{self.host}{self.path}")
             self.close_connection = True
 
     def do_GET(self) -> None:
diff --git a/tests/test_handlers.py b/tests/test_handlers.py
index 1684c2d..2ee1047 100644
--- a/tests/test_handlers.py
+++ b/tests/test_handlers.py
@@ -204,6 +204,17 @@ class TestRequestHandler(BaseHandlerTestCase):
         ):
             handler.do_HEAD()
 
+    def test_do_head_stapler(self) -> None:
+        handler = self._get_handler("/.well-known/stapler/something")
+        with (
+            self.expects_status_only(
+                handler,
+                http.HTTPStatus.NO_CONTENT,
+            ),
+            self.seal_mocks(),
+        ):
+            handler.do_HEAD()
+
     def test_do_head_forward(self) -> None:
         handler = self._get_handler("/file")
         with (
@@ -222,6 +233,17 @@ class TestRequestHandler(BaseHandlerTestCase):
         ):
             handler.do_GET()
 
+    def test_do_get_stapler(self) -> None:
+        handler = self._get_handler("/.well-known/stapler/something")
+        with (
+            self.expects_status_only(
+                handler,
+                http.HTTPStatus.NO_CONTENT,
+            ),
+            self.seal_mocks(),
+        ):
+            handler.do_GET()
+
     def test_do_get_forward_on_other_path(self) -> None:
         handler = self._get_handler("/file")
         with (
@@ -1373,3 +1395,14 @@ class TestUpgradeHandler(BaseHandlerTestCase):
             self.seal_mocks(),
         ):
             handler.do_HEAD()
+
+    def test_do_head_stapler(self) -> None:
+        handler = self._get_handler("/.well-known/stapler/something")
+        with (
+            self.expects_status_only(
+                handler,
+                http.HTTPStatus.NO_CONTENT,
+            ),
+            self.seal_mocks(),
+        ):
+            handler.do_HEAD()