Saved tests

config.txt

@@ -67,10 +67,9 @@ static_root = static
 
 ## Saved Tests
 
-# Where to keep files when users `save` them. Note that files will be automatically
-# deleted from this directory, so it needs to be only used for REDbot.
+# Directory to keep test results when users `save` them.
 # Comment out to disable saving.
-save_dir = /tmp/redbot/
+save_dir = /tmp/
 
 # How long to store things when users save them, in days.
 save_days = 30

pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
     "markdown >= 3.4.4",
     "MarkupSafe >= 2.1.3",
     "netaddr >= 0.9.0",
-    "thor >= 0.9.6",
+    "thor >= 0.9.9",
     "typing-extensions >= 4.8.0",
 ]
 

redbot/daemon.py

@@ -13,9 +13,11 @@
 import locale
 import os
 from pstats import Stats
+import signal
 import sys
 import traceback
-from typing import Dict, Optional
+from types import FrameType
+from typing import Dict, Optional, Any, Union
 from urllib.parse import urlsplit
 
 from importlib_resources import files as resource_files
@@ -27,7 +29,7 @@
 import redbot
 from redbot.type import RawHeaderListType
 from redbot.webui import RedWebUi
-from redbot.webui.saved_tests import clean_saved_tests
+from redbot.webui.saved_tests import SavedTests
 
 if os.environ.get("SYSTEMD_WATCHDOG"):
     try:
@@ -64,6 +66,14 @@ def __init__(self, config: SectionProxy) -> None:
         if self.config.get("extra_base_dir"):
             self.extra_files = self.walk_files(self.config["extra_base_dir"])
 
+        # Set up signal handlers
+        signal.signal(signal.SIGINT, self.shutdown_handler)
+        signal.signal(signal.SIGABRT, self.shutdown_handler)
+        signal.signal(signal.SIGTERM, self.shutdown_handler)
+
+        # open the save db
+        self.saved = SavedTests(config, self.console)
+
         # Start garbage collection
         if config.get("save_dir", ""):
             thor.schedule(10, self.gc_state)
@@ -74,20 +84,22 @@ def __init__(self, config: SectionProxy) -> None:
             self.config.getint("port", fallback=8000),
         )
         server.on("exchange", self.handler)
-        try:
-            thor.run()
-        except KeyboardInterrupt:
-            self.console("Stopping...")
-            thor.stop()
+        thor.run()
 
     def watchdog_ping(self) -> None:
         notify(Notification.WATCHDOG)
         thor.schedule(self.watchdog_freq, self.watchdog_ping)
 
     def gc_state(self) -> None:
-        clean_saved_tests(self.config)
+        self.saved.clean()
         thor.schedule(self.config.getint("gc_mins", fallback=2) * 60, self.gc_state)
 
+    def shutdown_handler(self, sig: int, frame: Union[FrameType, None]) -> Any:
+        self.console("Stopping...")
+        thor.stop()
+        self.saved.shutdown()
+        sys.exit(0)
+
     def walk_files(self, dir_name: str, uri_base: bytes = b"") -> Dict[bytes, bytes]:
         out: Dict[bytes, bytes] = {}
         for root, _, files in os.walk(dir_name):
@@ -158,6 +170,7 @@ def request_done(self, trailers: RawHeaderListType) -> None:
             try:
                 RedWebUi(
                     self.server.config,
+                    self.server.saved,
                     self.method.decode(self.server.config["charset"]),
                     p_uri.query,
                     self.req_hdrs,
@@ -179,6 +192,7 @@ def request_done(self, trailers: RawHeaderListType) -> None:
                 dump = traceback.format_exc()
                 thor.stop()
                 self.server.console(dump)
+                self.server.saved.shutdown()
                 sys.exit(1)
         else:
             return self.serve_static(p_uri.path)

redbot/formatter/__init__.py

@@ -2,7 +2,6 @@
 Formatters for REDbot output.
 """
 
-
 from collections import defaultdict
 from configparser import SectionProxy
 import inspect

redbot/formatter/har.py

@@ -2,7 +2,6 @@
 HAR Formatter for REDbot.
 """
 
-
 import datetime
 import json
 from typing import Optional, Any, Dict, List

redbot/resource/__init__.py

@@ -48,6 +48,7 @@ def __init__(self, config: SectionProxy, descend: bool = False) -> None:
         self.ims_support: bool = False
         self.gzip_support: bool = False
         self.gzip_savings: int = 0
+        self.save_expires: float = 0.0
         self._task_map: Set[RedFetcher] = set([])
         self.subreqs = {ac.check_name: ac(config, self) for ac in active_checks}
         self.once("fetch_done", self.run_active_checks)

redbot/resource/fetch.py

@@ -81,7 +81,10 @@ def __getstate__(self) -> Dict[str, Any]:
             del state["exchange"]
         except KeyError:
             pass
-        del state["response_content_processors"]
+        try:
+            del state["response_content_processors"]
+        except KeyError:
+            pass
         return state
 
     def __repr__(self) -> str:

redbot/type.py

@@ -11,11 +11,8 @@
 class HttpResponseExchange(Protocol):
     def response_start(
         self, status_code: bytes, status_phrase: bytes, res_hdrs: RawHeaderListType
-    ) -> None:
-        ...
+    ) -> None: ...
 
-    def response_body(self, chunk: bytes) -> None:
-        ...
+    def response_body(self, chunk: bytes) -> None: ...
 
-    def response_done(self, trailers: RawHeaderListType) -> None:
-        ...
+    def response_done(self, trailers: RawHeaderListType) -> None: ...

redbot/webui/__init__.py

@@ -20,12 +20,7 @@
 from redbot import __version__
 from redbot.webui.captcha import CaptchaHandler
 from redbot.webui.ratelimit import ratelimiter
-from redbot.webui.saved_tests import (
-    init_save_file,
-    save_test,
-    extend_saved_test,
-    load_saved_test,
-)
+from redbot.webui.saved_tests import SavedTests
 from redbot.webui.slack import slack_run, slack_auth
 from redbot.resource import HttpResource
 from redbot.formatter import find_formatter, html, Formatter
@@ -50,6 +45,7 @@ class RedWebUi:
     def __init__(
         self,
         config: SectionProxy,
+        saved: SavedTests,
         method: str,
         query_string: bytes,
         req_headers: RawHeaderListType,
@@ -58,13 +54,15 @@ def __init__(
         client_ip: str,
         console: Callable[[str], Optional[int]] = sys.stderr.write,
     ) -> None:
-        self.config: SectionProxy = config
+        self.config = config
+        self.saved = saved
+        self.method = method
         self.charset = self.config["charset"]
         self.charset_bytes = self.charset.encode("ascii")
         self.query_string = parse_qs(query_string.decode(self.charset, "replace"))
         self.req_headers = req_headers
         self.req_body = req_body
-        self.body_args = {}
+        self.body_args: Dict[str, list[str]] = {}
         self.exchange = exchange
         self.client_ip = client_ip
         self.console = console  # function to log errors to
@@ -88,25 +86,38 @@ def __init__(
         if not self.descend:
             self.check_name = self.query_string.get("check_name", [None])[0]
 
-        self.save_path: str
         self.timeout: Optional[thor.loop.ScheduledEvent] = None
 
         self.nonce: str = standard_b64encode(
             getrandbits(128).to_bytes(16, "big")
         ).decode("ascii")
         self.start = time.time()
+        self.handle_request()
 
-        if method == "POST":
+    def handle_request(self) -> None:
+        if self.method == "POST":
             req_ct = get_header(self.req_headers, b"content-type")
             if req_ct and req_ct[-1].lower() == b"application/x-www-form-urlencoded":
-                self.body_args = parse_qs(req_body.decode(self.charset, "replace"))
+                self.body_args = parse_qs(self.req_body.decode(self.charset, "replace"))
 
             if (
                 "save" in self.query_string
                 and self.config.get("save_dir", "")
                 and self.test_id
             ):
-                extend_saved_test(self)
+                try:
+                    self.saved.extend(self.test_id)
+                except KeyError:
+                    return self.error_response(
+                        None,
+                        b"404",
+                        b"Not Found",
+                        f"Can't find the test ID {self.test_id}",
+                    )
+                location = b"?id=%s" % self.test_id.encode("ascii")
+                if self.descend:
+                    location = b"%s&descend=True" % location
+                self.redirect_response(location)
             elif "slack" in self.query_string:
                 slack_run(self)
             elif "client_error" in self.query_string:
@@ -115,31 +126,78 @@ def __init__(
                 self.run_test()
             else:
                 self.show_default()
-        elif method in ["GET", "HEAD"]:
+        elif self.method in ["GET", "HEAD"]:
             if self.test_id:
-                load_saved_test(self)
+                self.show_saved_test()
             elif "code" in self.query_string:
                 slack_auth(self)
             else:
                 self.show_default()
         else:
             self.error_response(
-                find_formatter("html")(
-                    self.config,
-                    HttpResource(self.config),
-                    self.output,
-                    {
-                        "nonce": self.nonce,
-                    },
-                ),
+                None,
                 b"405",
                 b"Method Not Allowed",
                 "Method Not Allowed",
             )
+        return None
+
+    def show_saved_test(self) -> None:
+        """Show a saved test."""
+        try:
+            top_resource = self.saved.load(self)
+        except ValueError:
+            return self.error_response(
+                None,
+                b"400",
+                b"Bad Request",
+                "Saved tests are not available.",
+            )
+        except KeyError:
+            return self.error_response(
+                None,
+                b"404",
+                b"Not Found",
+                f"Can't find the test ID {self.test_id}",
+            )
+
+        if self.check_name:
+            display_resource = cast(
+                HttpResource, top_resource.subreqs.get(self.check_name, top_resource)
+            )
+        else:
+            display_resource = top_resource
+
+        formatter = find_formatter(self.format, "html", top_resource.descend)(
+            self.config,
+            display_resource,
+            self.output,
+            {
+                "allow_save": True,
+                "is_saved": True,
+                "test_id": self.test_id,
+                "nonce": self.nonce,
+            },
+        )
+        self.exchange.response_start(
+            b"200",
+            b"OK",
+            [
+                (b"Content-Type", formatter.content_type()),
+                (b"Cache-Control", b"max-age=3600, must-revalidate"),
+            ],
+        )
+
+        @thor.events.on(formatter)
+        def formatter_done() -> None:
+            self.exchange.response_done([])
+
+        formatter.bind_resource(display_resource)
+        return None
 
     def run_test(self) -> None:
         """Test a URI."""
-        self.test_id = init_save_file(self)
+        self.test_id = self.saved.get_test_id()
         top_resource = HttpResource(self.config, descend=self.descend)
         top_resource.set_request(self.test_uri, headers=self.req_hdrs)
         formatter = find_formatter(self.format, "html", self.descend)(
@@ -224,7 +282,7 @@ def formatter_done() -> None:
                 self.timeout.delete()
                 self.timeout = None
             self.exchange.response_done([])
-            save_test(self, top_resource)
+            self.saved.save(self, top_resource)
 
             # log excessive traffic
             ti = sum(
@@ -313,7 +371,7 @@ def show_default(self) -> None:
 
     def error_response(
         self,
-        formatter: Formatter,
+        formatter: Optional[Formatter],
         status_code: bytes,
         status_phrase: bytes,
         message: str,
@@ -323,6 +381,15 @@ def error_response(
         if self.timeout:
             self.timeout.delete()
             self.timeout = None
+        if formatter is None:
+            formatter = find_formatter("html")(
+                self.config,
+                HttpResource(self.config),
+                self.output,
+                {
+                    "nonce": self.nonce,
+                },
+            )
         self.exchange.response_start(
             status_code,
             status_phrase,
@@ -341,6 +408,11 @@ def error_response(
         if log_message:
             self.error_log(log_message)
 
+    def redirect_response(self, location: bytes) -> None:
+        self.exchange.response_start(b"303", b"See Other", [(b"Location", location)])
+        self.output("Redirecting...")
+        self.exchange.response_done([])
+
     def output(self, chunk: str) -> None:
         self.exchange.response_body(chunk.encode(self.charset, "replace"))