From 5af0ead5b9b1232526511ab81ad1f59bf370a290 Mon Sep 17 00:00:00 2001
From: Itamar Turner-Trauring <itamar@pythonspeed.com>
Date: Tue, 6 Jun 2023 10:58:16 -0400
Subject: [PATCH 1/8] Refactor HTTP client creation to be more centralized.

---
 src/allmydata/storage/http_client.py    | 170 +++++++++++++++---------
 src/allmydata/storage_client.py         |  50 ++-----
 src/allmydata/test/common_system.py     |   4 +-
 src/allmydata/test/test_storage_http.py |  12 +-
 4 files changed, 128 insertions(+), 108 deletions(-)

diff --git a/src/allmydata/storage/http_client.py b/src/allmydata/storage/http_client.py
index 670d84be3..fe2545c03 100644
--- a/src/allmydata/storage/http_client.py
+++ b/src/allmydata/storage/http_client.py
@@ -16,6 +16,7 @@ from typing import (
     Set,
     Dict,
     Callable,
+    ClassVar,
 )
 from base64 import b64encode
 from io import BytesIO
@@ -60,6 +61,15 @@ from .http_common import (
 from .common import si_b2a, si_to_human_readable
 from ..util.hashutil import timing_safe_compare
 from ..util.deferredutil import async_to_deferred
+from ..util.tor_provider import _Provider as TorProvider
+
+try:
+    from txtorcon import Tor  # type: ignore
+except ImportError:
+
+    class Tor:
+        pass
+
 
 _OPENSSL = Binding().lib
 
@@ -302,18 +312,30 @@ class _StorageClientHTTPSPolicy:
         )
 
 
-@define(hash=True)
-class StorageClient(object):
+@define
+class StorageClientFactory:
     """
-    Low-level HTTP client that talks to the HTTP storage server.
+    Create ``StorageClient`` instances, using appropriate
+    ``twisted.web.iweb.IAgent`` for different connection methods: normal TCP,
+    Tor, and eventually I2P.
+
+    There is some caching involved since there might be shared setup work, e.g.
+    connecting to the local Tor service only needs to happen once.
     """
 
-    # If set, we're doing unit testing and we should call this with
-    # HTTPConnectionPool we create.
-    TEST_MODE_REGISTER_HTTP_POOL = None
+    _default_connection_handlers: dict[str, str]
+    _tor_provider: Optional[TorProvider]
+    # Cache the Tor instance created by the provider, if relevant.
+    _tor_instance: Optional[Tor] = None
+
+    # If set, we're doing unit testing and we should call this with any
+    # HTTPConnectionPool that gets passed/created to ``create_agent()``.
+    TEST_MODE_REGISTER_HTTP_POOL = ClassVar[
+        Optional[Callable[[HTTPConnectionPool], None]]
+    ]
 
     @classmethod
-    def start_test_mode(cls, callback):
+    def start_test_mode(cls, callback: Callable[[HTTPConnectionPool], None]) -> None:
         """Switch to testing mode.
 
         In testing mode we register the pool with test system using the given
@@ -328,66 +350,84 @@ class StorageClient(object):
         """Stop testing mode."""
         cls.TEST_MODE_REGISTER_HTTP_POOL = None
 
-    # The URL is a HTTPS URL ("https://...").  To construct from a NURL, use
-    # ``StorageClient.from_nurl()``.
-    _base_url: DecodedURL
-    _swissnum: bytes
-    _treq: Union[treq, StubTreq, HTTPClient]
-    _pool: Optional[HTTPConnectionPool]
-    _clock: IReactorTime
-
-    @classmethod
-    def from_nurl(
-        cls,
+    async def _create_agent(
+        self,
         nurl: DecodedURL,
-        reactor,
-        # TODO default_connection_handlers should really be a class, not a dict
-        # of strings...
-        default_connection_handlers: dict[str, str],
-        pool: Optional[HTTPConnectionPool] = None,
-        agent_factory: Optional[
-            Callable[[object, IPolicyForHTTPS, HTTPConnectionPool], IAgent]
-        ] = None,
-    ) -> StorageClient:
-        """
-        Create a ``StorageClient`` for the given NURL.
-        """
-        # Safety check: if we're using normal TCP connections, we better not be
-        # configured for Tor or I2P.
-        if agent_factory is None:
-            assert default_connection_handlers["tcp"] == "tcp"
+        reactor: object,
+        tls_context_factory: IPolicyForHTTPS,
+        pool: HTTPConnectionPool,
+    ) -> IAgent:
+        """Create a new ``IAgent``, possibly using Tor."""
+        if self.TEST_MODE_REGISTER_HTTP_POOL is not None:
+            self.TEST_MODE_REGISTER_HTTP_POOL(pool)
 
+        # TODO default_connection_handlers should really be an object, not a
+        # dict, so we can ask "is this using Tor" without poking at a
+        # dictionary with arbitrary strings... See
+        # https://tahoe-lafs.org/trac/tahoe-lafs/ticket/4032
+        handler = self._default_connection_handlers["tcp"]
+
+        if handler == "tcp":
+            return Agent(reactor, tls_context_factory, pool=pool)
+        if handler == "tor":  # TODO or nurl.scheme == "pb+tor":
+            assert self._tor_provider is not None
+            if self._tor_instance is None:
+                self._tor_instance = await self._tor_provider.get_tor_instance(reactor)
+            return self._tor_instance.web_agent(
+                pool=pool, tls_context_factory=tls_context_factory
+            )
+        else:
+            raise RuntimeError(f"Unsupported tcp connection handler: {handler}")
+
+    async def create_storage_client(
+        self,
+        nurl: DecodedURL,
+        reactor: IReactorTime,
+        pool: Optional[HTTPConnectionPool] = None,
+    ) -> StorageClient:
+        """Create a new ``StorageClient`` for the given NURL."""
         assert nurl.fragment == "v=1"
-        assert nurl.scheme == "pb"
-        swissnum = nurl.path[0].encode("ascii")
-        certificate_hash = nurl.user.encode("ascii")
+        assert nurl.scheme in ("pb", "pb+tor")
         if pool is None:
             pool = HTTPConnectionPool(reactor)
             pool.maxPersistentPerHost = 10
 
-        if cls.TEST_MODE_REGISTER_HTTP_POOL is not None:
-            cls.TEST_MODE_REGISTER_HTTP_POOL(pool)
-
-        def default_agent_factory(
-            reactor: object,
-            tls_context_factory: IPolicyForHTTPS,
-            pool: HTTPConnectionPool,
-        ) -> IAgent:
-            return Agent(reactor, tls_context_factory, pool=pool)
-
-        if agent_factory is None:
-            agent_factory = default_agent_factory
-
-        treq_client = HTTPClient(
-            agent_factory(
-                reactor,
-                _StorageClientHTTPSPolicy(expected_spki_hash=certificate_hash),
-                pool,
-            )
+        certificate_hash = nurl.user.encode("ascii")
+        agent = await self._create_agent(
+            nurl,
+            reactor,
+            _StorageClientHTTPSPolicy(expected_spki_hash=certificate_hash),
+            pool,
+        )
+        treq_client = HTTPClient(agent)
+        https_url = DecodedURL().replace(scheme="https", host=nurl.host, port=nurl.port)
+        swissnum = nurl.path[0].encode("ascii")
+        return StorageClient(
+            https_url,
+            swissnum,
+            treq_client,
+            pool,
+            reactor,
+            self.TEST_MODE_REGISTER_HTTP_POOL is not None,
         )
 
-        https_url = DecodedURL().replace(scheme="https", host=nurl.host, port=nurl.port)
-        return cls(https_url, swissnum, treq_client, pool, reactor)
+
+@define(hash=True)
+class StorageClient(object):
+    """
+    Low-level HTTP client that talks to the HTTP storage server.
+
+    Create using a ``StorageClientFactory`` instance.
+    """
+
+    # The URL should be a HTTPS URL ("https://...")
+    _base_url: DecodedURL
+    _swissnum: bytes
+    _treq: Union[treq, StubTreq, HTTPClient]
+    _pool: HTTPConnectionPool
+    _clock: IReactorTime
+    # Are we running unit tests?
+    _test_mode: bool
 
     def relative_url(self, path: str) -> DecodedURL:
         """Get a URL relative to the base URL."""
@@ -495,12 +535,11 @@ class StorageClient(object):
             method, url, headers=headers, timeout=timeout, **kwargs
         )
 
-        if self.TEST_MODE_REGISTER_HTTP_POOL is not None:
-            if response.code != 404:
-                # We're doing API queries, HTML is never correct except in 404, but
-                # it's the default for Twisted's web server so make sure nothing
-                # unexpected happened.
-                assert get_content_type(response.headers) != "text/html"
+        if self._test_mode and response.code != 404:
+            # We're doing API queries, HTML is never correct except in 404, but
+            # it's the default for Twisted's web server so make sure nothing
+            # unexpected happened.
+            assert get_content_type(response.headers) != "text/html"
 
         return response
 
@@ -529,8 +568,7 @@ class StorageClient(object):
 
     def shutdown(self) -> Deferred:
         """Shutdown any connections."""
-        if self._pool is not None:
-            return self._pool.closeCachedConnections()
+        return self._pool.closeCachedConnections()
 
 
 @define(hash=True)
diff --git a/src/allmydata/storage_client.py b/src/allmydata/storage_client.py
index 2ea154263..6a965aaac 100644
--- a/src/allmydata/storage_client.py
+++ b/src/allmydata/storage_client.py
@@ -89,7 +89,8 @@ from allmydata.util.deferredutil import async_to_deferred, race
 from allmydata.storage.http_client import (
     StorageClient, StorageClientImmutables, StorageClientGeneral,
     ClientException as HTTPClientException, StorageClientMutables,
-    ReadVector, TestWriteVectors, WriteVector, TestVector, ClientException
+    ReadVector, TestWriteVectors, WriteVector, TestVector, ClientException,
+    StorageClientFactory
 )
 from .node import _Config
 
@@ -1068,8 +1069,9 @@ class HTTPNativeStorageServer(service.MultiService):
         self._on_status_changed = ObserverList()
         self._reactor = reactor
         self._grid_manager_verifier = grid_manager_verifier
-        self._tor_provider = tor_provider
-        self._default_connection_handlers = default_connection_handlers
+        self._storage_client_factory = StorageClientFactory(
+            default_connection_handlers, tor_provider
+        )
 
         furl = announcement["anonymous-storage-FURL"].encode("utf-8")
         (
@@ -1232,26 +1234,6 @@ class HTTPNativeStorageServer(service.MultiService):
         self._connecting_deferred = connecting
         return connecting
 
-    async def _agent_factory(self) -> Optional[Callable[[object, IPolicyForHTTPS, HTTPConnectionPool],IAgent]]:
-        """Return a factory for ``twisted.web.iweb.IAgent``."""
-        # TODO default_connection_handlers should really be an object, not a
-        # dict, so we can ask "is this using Tor" without poking at a
-        # dictionary with arbitrary strings... See
-        # https://tahoe-lafs.org/trac/tahoe-lafs/ticket/4032
-        handler = self._default_connection_handlers["tcp"]
-        if handler == "tcp":
-            return None
-        if handler == "tor":
-            assert self._tor_provider is not None
-            tor_instance = await self._tor_provider.get_tor_instance(self._reactor)
-
-            def agent_factory(reactor: object, tls_context_factory: IPolicyForHTTPS, pool: HTTPConnectionPool) -> IAgent:
-                assert reactor == self._reactor
-                return tor_instance.web_agent(pool=pool, tls_context_factory=tls_context_factory)
-            return agent_factory
-        else:
-            raise RuntimeError(f"Unsupported tcp connection handler: {handler}")
-
     @async_to_deferred
     async def _pick_server_and_get_version(self):
         """
@@ -1270,28 +1252,24 @@ class HTTPNativeStorageServer(service.MultiService):
             # version() calls before we are live talking to a server, it could only
             # be one. See https://tahoe-lafs.org/trac/tahoe-lafs/ticket/3992
 
-            agent_factory = await self._agent_factory()
-
-            def request(reactor, nurl: DecodedURL):
+            @async_to_deferred
+            async def request(reactor, nurl: DecodedURL):
                 # Since we're just using this one off to check if the NURL
                 # works, no need for persistent pool or other fanciness.
                 pool = HTTPConnectionPool(reactor, persistent=False)
                 pool.retryAutomatically = False
-                return StorageClientGeneral(
-                    StorageClient.from_nurl(
-                        nurl, reactor, self._default_connection_handlers,
-                        pool=pool, agent_factory=agent_factory)
-                ).get_version()
+                storage_client = await self._storage_client_factory.create_storage_client(
+                    nurl, reactor, pool
+                )
+                return await StorageClientGeneral(storage_client).get_version()
 
             nurl = await _pick_a_http_server(reactor, self._nurls, request)
 
             # If we've gotten this far, we've found a working NURL.
-            self._istorage_server = _HTTPStorageServer.from_http_client(
-                StorageClient.from_nurl(
-                    nurl, reactor, self._default_connection_handlers,
-                    agent_factory=agent_factory
-                )
+            storage_client = await self._storage_client_factory.create_storage_client(
+                    nurl, reactor, None
             )
+            self._istorage_server = _HTTPStorageServer.from_http_client(storage_client)
             return self._istorage_server
 
         try:
diff --git a/src/allmydata/test/common_system.py b/src/allmydata/test/common_system.py
index fa8d943e5..cfb6c9f04 100644
--- a/src/allmydata/test/common_system.py
+++ b/src/allmydata/test/common_system.py
@@ -686,8 +686,8 @@ class SystemTestMixin(pollmixin.PollMixin, testutil.StallMixin):
 
     def setUp(self):
         self._http_client_pools = []
-        http_client.StorageClient.start_test_mode(self._got_new_http_connection_pool)
-        self.addCleanup(http_client.StorageClient.stop_test_mode)
+        http_client.StorageClientFactory.start_test_mode(self._got_new_http_connection_pool)
+        self.addCleanup(http_client.StorageClientFactory.stop_test_mode)
         self.port_assigner = SameProcessStreamEndpointAssigner()
         self.port_assigner.setUp()
         self.addCleanup(self.port_assigner.tearDown)
diff --git a/src/allmydata/test/test_storage_http.py b/src/allmydata/test/test_storage_http.py
index 1380ab7e7..233d82989 100644
--- a/src/allmydata/test/test_storage_http.py
+++ b/src/allmydata/test/test_storage_http.py
@@ -58,6 +58,7 @@ from ..storage.http_server import (
 )
 from ..storage.http_client import (
     StorageClient,
+    StorageClientFactory,
     ClientException,
     StorageClientImmutables,
     ImmutableCreateResult,
@@ -323,10 +324,10 @@ class CustomHTTPServerTests(SyncTestCase):
 
     def setUp(self):
         super(CustomHTTPServerTests, self).setUp()
-        StorageClient.start_test_mode(
+        StorageClientFactory.start_test_mode(
             lambda pool: self.addCleanup(pool.closeCachedConnections)
         )
-        self.addCleanup(StorageClient.stop_test_mode)
+        self.addCleanup(StorageClientFactory.stop_test_mode)
         # Could be a fixture, but will only be used in this test class so not
         # going to bother:
         self._http_server = TestApp()
@@ -341,6 +342,7 @@ class CustomHTTPServerTests(SyncTestCase):
             # fixed if https://github.com/twisted/treq/issues/226 were ever
             # fixed.
             clock=treq._agent._memoryReactor,
+            test_mode=True,
         )
         self._http_server.clock = self.client._clock
 
@@ -529,10 +531,10 @@ class HttpTestFixture(Fixture):
     """
 
     def _setUp(self):
-        StorageClient.start_test_mode(
+        StorageClientFactory.start_test_mode(
             lambda pool: self.addCleanup(pool.closeCachedConnections)
         )
-        self.addCleanup(StorageClient.stop_test_mode)
+        self.addCleanup(StorageClientFactory.stop_test_mode)
         self.clock = Reactor()
         self.tempdir = self.useFixture(TempDir())
         # The global Cooperator used by Twisted (a) used by pull producers in
@@ -558,6 +560,7 @@ class HttpTestFixture(Fixture):
             treq=self.treq,
             pool=None,
             clock=self.clock,
+            test_mode=True,
         )
 
     def result_of_with_flush(self, d):
@@ -671,6 +674,7 @@ class GenericHTTPAPITests(SyncTestCase):
                 treq=StubTreq(self.http.http_server.get_resource()),
                 pool=None,
                 clock=self.http.clock,
+                test_mode=True,
             )
         )
         with assert_fails_with_http_code(self, http.UNAUTHORIZED):

From 74a121da74af8928cb212e550869fc4c7d511cde Mon Sep 17 00:00:00 2001
From: Itamar Turner-Trauring <itamar@pythonspeed.com>
Date: Tue, 6 Jun 2023 11:47:36 -0400
Subject: [PATCH 2/8] Fix bug which meant object could not be created.

---
 src/allmydata/storage/http_client.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/allmydata/storage/http_client.py b/src/allmydata/storage/http_client.py
index fe2545c03..8cf79843f 100644
--- a/src/allmydata/storage/http_client.py
+++ b/src/allmydata/storage/http_client.py
@@ -330,9 +330,9 @@ class StorageClientFactory:
 
     # If set, we're doing unit testing and we should call this with any
     # HTTPConnectionPool that gets passed/created to ``create_agent()``.
-    TEST_MODE_REGISTER_HTTP_POOL = ClassVar[
+    TEST_MODE_REGISTER_HTTP_POOL: ClassVar[
         Optional[Callable[[HTTPConnectionPool], None]]
-    ]
+    ] = None
 
     @classmethod
     def start_test_mode(cls, callback: Callable[[HTTPConnectionPool], None]) -> None:

From e8744f91e5176d6adf2f1fc4c2065f7041162573 Mon Sep 17 00:00:00 2001
From: Itamar Turner-Trauring <itamar@pythonspeed.com>
Date: Tue, 6 Jun 2023 12:06:51 -0400
Subject: [PATCH 3/8] Hook up HTTP storage for servers listening on .onion
 addresses

---
 src/allmydata/protocol_switch.py     | 14 +++++++++++---
 src/allmydata/storage/http_client.py |  2 +-
 src/allmydata/storage/http_server.py | 11 +++++++++--
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/src/allmydata/protocol_switch.py b/src/allmydata/protocol_switch.py
index 208efec6c..941b104be 100644
--- a/src/allmydata/protocol_switch.py
+++ b/src/allmydata/protocol_switch.py
@@ -102,8 +102,15 @@ class _FoolscapOrHttps(Protocol, metaclass=_PretendToBeNegotiation):
         for location_hint in chain.from_iterable(
             hints.split(",") for hints in cls.tub.locationHints
         ):
-            if location_hint.startswith("tcp:"):
-                _, hostname, port = location_hint.split(":")
+            if location_hint.startswith("tcp:") or location_hint.startswith("tor:"):
+                scheme, hostname, port = location_hint.split(":")
+                if scheme == "tcp":
+                    subscheme = None
+                else:
+                    subscheme = "tor"
+                    # If we're listening on Tor, the hostname needs to have an
+                    # .onion TLD.
+                    assert hostname.endswith(".onion")
                 port = int(port)
                 storage_nurls.add(
                     build_nurl(
@@ -111,9 +118,10 @@ class _FoolscapOrHttps(Protocol, metaclass=_PretendToBeNegotiation):
                         port,
                         str(swissnum, "ascii"),
                         cls.tub.myCertificate.original.to_cryptography(),
+                        subscheme
                     )
                 )
-            # TODO this is probably where we'll have to support Tor and I2P?
+            # TODO this is where we'll have to support Tor and I2P as well.
             # See https://tahoe-lafs.org/trac/tahoe-lafs/ticket/3888#comment:9
             # for discussion (there will be separate tickets added for those at
             # some point.)
diff --git a/src/allmydata/storage/http_client.py b/src/allmydata/storage/http_client.py
index 8cf79843f..cd3143924 100644
--- a/src/allmydata/storage/http_client.py
+++ b/src/allmydata/storage/http_client.py
@@ -369,7 +369,7 @@ class StorageClientFactory:
 
         if handler == "tcp":
             return Agent(reactor, tls_context_factory, pool=pool)
-        if handler == "tor":  # TODO or nurl.scheme == "pb+tor":
+        if handler == "tor" or nurl.scheme == "pb+tor":
             assert self._tor_provider is not None
             if self._tor_instance is None:
                 self._tor_instance = await self._tor_provider.get_tor_instance(reactor)
diff --git a/src/allmydata/storage/http_server.py b/src/allmydata/storage/http_server.py
index 924ae5a43..028ebf1c7 100644
--- a/src/allmydata/storage/http_server.py
+++ b/src/allmydata/storage/http_server.py
@@ -995,13 +995,20 @@ class _TLSEndpointWrapper(object):
 
 
 def build_nurl(
-    hostname: str, port: int, swissnum: str, certificate: CryptoCertificate
+    hostname: str,
+    port: int,
+    swissnum: str,
+    certificate: CryptoCertificate,
+    subscheme: Optional[str] = None,
 ) -> DecodedURL:
     """
     Construct a HTTPS NURL, given the hostname, port, server swissnum, and x509
     certificate for the server.  Clients can then connect to the server using
     this NURL.
     """
+    scheme = "pb"
+    if subscheme is not None:
+        scheme = f"{scheme}+{subscheme}"
     return DecodedURL().replace(
         fragment="v=1",  # how we know this NURL is HTTP-based (i.e. not Foolscap)
         host=hostname,
@@ -1013,7 +1020,7 @@ def build_nurl(
                 "ascii",
             ),
         ),
-        scheme="pb",
+        scheme=scheme,
     )
 
 

From 57a6721670da156bbf94b72c0b904f04633f6757 Mon Sep 17 00:00:00 2001
From: Itamar Turner-Trauring <itamar@pythonspeed.com>
Date: Tue, 6 Jun 2023 12:07:13 -0400
Subject: [PATCH 4/8] News file.

---
 newsfragments/3910.minor | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 newsfragments/3910.minor

diff --git a/newsfragments/3910.minor b/newsfragments/3910.minor
new file mode 100644
index 000000000..e69de29bb

From a977180baf0b138c8dd3824d95ae3f66bec1540d Mon Sep 17 00:00:00 2001
From: Itamar Turner-Trauring <itamar@pythonspeed.com>
Date: Tue, 6 Jun 2023 12:15:31 -0400
Subject: [PATCH 5/8] Fix lint

---
 src/allmydata/storage_client.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/allmydata/storage_client.py b/src/allmydata/storage_client.py
index 6a965aaac..a614c17db 100644
--- a/src/allmydata/storage_client.py
+++ b/src/allmydata/storage_client.py
@@ -51,7 +51,6 @@ from zope.interface import (
 )
 from twisted.python.failure import Failure
 from twisted.web import http
-from twisted.web.iweb import IAgent, IPolicyForHTTPS
 from twisted.internet.task import LoopingCall
 from twisted.internet import defer, reactor
 from twisted.application import service

From 20d4175abcbe2948f6e219fd4dd564e7fea47ae3 Mon Sep 17 00:00:00 2001
From: Itamar Turner-Trauring <itamar@pythonspeed.com>
Date: Tue, 6 Jun 2023 12:18:02 -0400
Subject: [PATCH 6/8] Fix typecheck complaint

---
 src/allmydata/storage/http_client.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/allmydata/storage/http_client.py b/src/allmydata/storage/http_client.py
index cd3143924..e7df3709d 100644
--- a/src/allmydata/storage/http_client.py
+++ b/src/allmydata/storage/http_client.py
@@ -67,7 +67,7 @@ try:
     from txtorcon import Tor  # type: ignore
 except ImportError:
 
-    class Tor:
+    class Tor:  # type: ignore[no-redef]
         pass
 
 

From 122e0a73a979f379e6bc7a3f795847be8dc6db0b Mon Sep 17 00:00:00 2001
From: meejah <meejah@meejah.ca>
Date: Thu, 22 Jun 2023 01:29:55 -0600
Subject: [PATCH 7/8] more-generic testing hook

---
 src/allmydata/storage/http_client.py    |  9 ++-------
 src/allmydata/test/test_storage_http.py | 17 ++++++++++++++---
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/allmydata/storage/http_client.py b/src/allmydata/storage/http_client.py
index e7df3709d..8c0100656 100644
--- a/src/allmydata/storage/http_client.py
+++ b/src/allmydata/storage/http_client.py
@@ -427,7 +427,7 @@ class StorageClient(object):
     _pool: HTTPConnectionPool
     _clock: IReactorTime
     # Are we running unit tests?
-    _test_mode: bool
+    _analyze_response: Callable[[IResponse], None] = lambda _: None
 
     def relative_url(self, path: str) -> DecodedURL:
         """Get a URL relative to the base URL."""
@@ -534,12 +534,7 @@ class StorageClient(object):
         response = await self._treq.request(
             method, url, headers=headers, timeout=timeout, **kwargs
         )
-
-        if self._test_mode and response.code != 404:
-            # We're doing API queries, HTML is never correct except in 404, but
-            # it's the default for Twisted's web server so make sure nothing
-            # unexpected happened.
-            assert get_content_type(response.headers) != "text/html"
+        self._analyze_response(response)
 
         return response
 
diff --git a/src/allmydata/test/test_storage_http.py b/src/allmydata/test/test_storage_http.py
index 233d82989..aaa858db4 100644
--- a/src/allmydata/test/test_storage_http.py
+++ b/src/allmydata/test/test_storage_http.py
@@ -316,6 +316,17 @@ def result_of(d):
         + "This is probably a test design issue."
     )
 
+def response_is_not_html(response):
+    """
+    During tests, this is registered so we can ensure the web server
+    doesn't give us text/html.
+
+    HTML is never correct except in 404, but it's the default for
+    Twisted's web server so we assert nothing unexpected happened.
+    """
+    if response.code != 404:
+        assert get_content_type(response.headers) != "text/html"
+
 
 class CustomHTTPServerTests(SyncTestCase):
     """
@@ -342,7 +353,7 @@ class CustomHTTPServerTests(SyncTestCase):
             # fixed if https://github.com/twisted/treq/issues/226 were ever
             # fixed.
             clock=treq._agent._memoryReactor,
-            test_mode=True,
+            analyze_response=response_is_not_html,
         )
         self._http_server.clock = self.client._clock
 
@@ -560,7 +571,7 @@ class HttpTestFixture(Fixture):
             treq=self.treq,
             pool=None,
             clock=self.clock,
-            test_mode=True,
+            analyze_response=response_is_not_html,
         )
 
     def result_of_with_flush(self, d):
@@ -674,7 +685,7 @@ class GenericHTTPAPITests(SyncTestCase):
                 treq=StubTreq(self.http.http_server.get_resource()),
                 pool=None,
                 clock=self.http.clock,
-                test_mode=True,
+                analyze_response=response_is_not_html,
             )
         )
         with assert_fails_with_http_code(self, http.UNAUTHORIZED):

From 75b9c59846bffadbeca2c5941931d335790a23bd Mon Sep 17 00:00:00 2001
From: meejah <meejah@meejah.ca>
Date: Thu, 22 Jun 2023 01:54:53 -0600
Subject: [PATCH 8/8] refactor

---
 src/allmydata/storage/http_client.py    |  7 ++++++-
 src/allmydata/storage/http_common.py    | 13 +++++++++++++
 src/allmydata/test/test_storage_http.py | 18 +++++-------------
 3 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/src/allmydata/storage/http_client.py b/src/allmydata/storage/http_client.py
index 8c0100656..f2165ffda 100644
--- a/src/allmydata/storage/http_client.py
+++ b/src/allmydata/storage/http_client.py
@@ -57,6 +57,7 @@ from .http_common import (
     get_content_type,
     CBOR_MIME_TYPE,
     get_spki_hash,
+    response_is_not_html,
 )
 from .common import si_b2a, si_to_human_readable
 from ..util.hashutil import timing_safe_compare
@@ -402,13 +403,17 @@ class StorageClientFactory:
         treq_client = HTTPClient(agent)
         https_url = DecodedURL().replace(scheme="https", host=nurl.host, port=nurl.port)
         swissnum = nurl.path[0].encode("ascii")
+        response_check = lambda _: None
+        if self.TEST_MODE_REGISTER_HTTP_POOL is not None:
+            response_check = response_is_not_html
+
         return StorageClient(
             https_url,
             swissnum,
             treq_client,
             pool,
             reactor,
-            self.TEST_MODE_REGISTER_HTTP_POOL is not None,
+            response_check,
         )
 
 
diff --git a/src/allmydata/storage/http_common.py b/src/allmydata/storage/http_common.py
index e5f07898e..7ee137e1d 100644
--- a/src/allmydata/storage/http_common.py
+++ b/src/allmydata/storage/http_common.py
@@ -12,6 +12,7 @@ from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat
 
 from werkzeug.http import parse_options_header
 from twisted.web.http_headers import Headers
+from twisted.web.iweb import IResponse
 
 CBOR_MIME_TYPE = "application/cbor"
 
@@ -27,6 +28,18 @@ def get_content_type(headers: Headers) -> Optional[str]:
     return content_type
 
 
+def response_is_not_html(response: IResponse) -> None:
+    """
+    During tests, this is registered so we can ensure the web server
+    doesn't give us text/html.
+
+    HTML is never correct except in 404, but it's the default for
+    Twisted's web server so we assert nothing unexpected happened.
+    """
+    if response.code != 404:
+        assert get_content_type(response.headers) != "text/html"
+
+
 def swissnum_auth_header(swissnum: bytes) -> bytes:
     """Return value for ``Authorization`` header."""
     return b"Tahoe-LAFS " + b64encode(swissnum).strip()
diff --git a/src/allmydata/test/test_storage_http.py b/src/allmydata/test/test_storage_http.py
index aaa858db4..f660342ae 100644
--- a/src/allmydata/test/test_storage_http.py
+++ b/src/allmydata/test/test_storage_http.py
@@ -43,7 +43,11 @@ from testtools.matchers import Equals
 from zope.interface import implementer
 
 from .common import SyncTestCase
-from ..storage.http_common import get_content_type, CBOR_MIME_TYPE
+from ..storage.http_common import (
+    get_content_type,
+    CBOR_MIME_TYPE,
+    response_is_not_html,
+)
 from ..storage.common import si_b2a
 from ..storage.lease import LeaseInfo
 from ..storage.server import StorageServer
@@ -316,18 +320,6 @@ def result_of(d):
         + "This is probably a test design issue."
     )
 
-def response_is_not_html(response):
-    """
-    During tests, this is registered so we can ensure the web server
-    doesn't give us text/html.
-
-    HTML is never correct except in 404, but it's the default for
-    Twisted's web server so we assert nothing unexpected happened.
-    """
-    if response.code != 404:
-        assert get_content_type(response.headers) != "text/html"
-
-
 class CustomHTTPServerTests(SyncTestCase):
     """
     Tests that use a custom HTTP server.